Create modeling_i3.py

modeling_i3.py

@@ -0,0 +1,34 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from .configuration_i3 import I3Config
+
+# import or paste your existing i3Model and i3Block classes here
+# (or import them if you split them out into another module)
+from .i3_architecture import i3Model  # optional if you modularize
+
+class I3ForCausalLM(PreTrainedModel):
+    config_class = I3Config
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = i3Model(
+            vocab_size=config.vocab_size,
+            d_model=config.d_model,
+            n_layers=config.n_layers,
+            n_heads=config.n_heads,
+            max_seq_len=config.max_seq_len,
+            rank=config.rank,
+            d_state=config.d_state,
+        )
+
+        self.post_init()
+
+    def forward(self, input_ids, labels=None):
+        logits, loss = self.model(input_ids, labels)
+        return {"loss": loss, "logits": logits}
+
+    @torch.no_grad()
+    def generate(self, input_ids, max_new_tokens=50, temperature=1.0, top_k=None):
+        return self.model.generate(input_ids, max_new_tokens, temperature, top_k)