Update modeling_i3.py

modeling_i3.py

@@ -1,10 +1,11 @@
+# modeling_i3.py
 import torch
 from transformers import PreTrainedModel
 from configuration_i3 import I3Config
-from i3_architecture import i3Model 
+from i3_architecture import i3Model  # your actual i3 implementation
 
 class I3ForCausalLM(PreTrainedModel):
-    config_class = I3Config  # tells Transformers which config class to use
+    config_class = I3Config
 
     def __init__(self, config):
         super().__init__(config)
@@ -17,22 +18,18 @@ class I3ForCausalLM(PreTrainedModel):
             rank=config.rank,
             d_state=config.d_state,
         )
-        self.post_init()  # ensures everything Transformers expects is initialized
+        self.post_init()
 
     def forward(self, input_ids, labels=None):
-        # Forward pass
         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):
-        # Generation method
         return self.model.generate(input_ids, max_new_tokens, temperature, top_k)
 
+# AutoClass registration (optional but recommended)
 from transformers import AutoConfig, AutoModelForCausalLM
-from configuration_i3 import I3Config
-from modeling_i3 import I3ForCausalLM
 
-# Register custom model
 AutoConfig.register("i3", I3Config)
 AutoModelForCausalLM.register(I3Config, I3ForCausalLM)