Update app.py

app.py

@@ -150,7 +150,9 @@ class SmolLM2ForCausalLM(PreTrainedModel):
         if config.tie_word_embeddings:
             self.lm_head.weight = self.embed_tokens.weight
             
-    def forward(self, input_ids, attention_mask=None, labels=None):
+    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=None, **kwargs):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
         hidden_states = self.embed_tokens(input_ids)
         
         # Create causal attention mask if none provided
@@ -171,8 +173,18 @@ class SmolLM2ForCausalLM(PreTrainedModel):
         loss = None
         if labels is not None:
             loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1))
-            
-        return logits if loss is None else (loss, logits)
+        
+        if return_dict:
+            from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+            return CausalLMOutputWithCrossAttentions(
+                loss=loss,
+                logits=logits,
+                past_key_values=None,
+                hidden_states=None,
+                attentions=None,
+                cross_attentions=None,
+            )
+        return (loss, logits) if loss is not None else logits
 
     def prepare_inputs_for_generation(self, input_ids, **kwargs):
         return {