[EXAMPLE]

example.py

@@ -1,30 +1,36 @@
 import torch
-from gemini_torch import Gemini
+from gemini_torch.model import Gemini
 
-# Initialize the model
+# Initialize model
 model = Gemini(
-    num_tokens=12608,
-    max_seq_len=2048,
-    dim=640,
-    depth=8,
-    dim_head=32,
-    heads=6,
+    num_tokens=50432,
+    max_seq_len=8192,
+    dim=2560,
+    depth=32,
+    dim_head=128,
+    heads=24,
     use_abs_pos_emb=False,
     alibi_pos_bias=True,
-    alibi_num_heads=3,
+    alibi_num_heads=12,
     rotary_xpos=True,
     attn_flash=True,
-    attn_kv_heads=1,
+    attn_kv_heads=2,
     qk_norm=True,
     attn_qk_norm=True,
     attn_qk_norm_dim_scale=True,
 )
 
-# Initialize the text random tokens
-x = torch.randint(0, 12608, (1, 2048))
+# Text shape: [batch, seq_len, dim]
+text = torch.randint(0, 50432, (1, 8192))
 
-# Apply model to x
-y = model(x)
+# Img shape: [batch, channels, height, width]
+img = torch.randn(1, 3, 256, 256)
 
-# Print logits
-print(y)
+# Audio shape: [batch, audio_seq_len, dim]
+audio = torch.randn(1, 128)
+
+# Apply model to text and img
+y = model(text, img, audio)
+
+# Output shape: [batch, seq_len, dim]
+print(y.shape)