Update modeling_qwen2_nomic_vision.py

modeling_qwen2_nomic_vision.py

@@ -770,6 +770,13 @@ QWEN2_INPUTS_DOCSTRING = r"""
 """
 
 
+@dataclass
+class Qwen2NomicVisionOutput(BaseModelOutputWithPast):
+    last_hidden_state: FloatTensor = None
+    past_key_values: Optional = None
+    hidden_states: Optional = None
+    attentions: Optional = None
+    processed_image: Optional = None
 # @add_start_docstrings(
 #     "The bare Qwen2 Model outputting raw hidden-states without any specific head on top.",
 #     QWEN2_START_DOCSTRING,
@@ -824,7 +831,7 @@ class Qwen2NomicVisionModel(Qwen2NomicVisionPreTrainedModel):
         return_dict: Optional[bool] = None,
         cache_position: Optional[torch.LongTensor] = None,
         image = None,
-    ) -> Union[Tuple, BaseModelOutputWithPast]:
+    ) -> Union[Tuple, Qwen2NomicVisionOutput]:
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
@@ -936,12 +943,15 @@ class Qwen2NomicVisionModel(Qwen2NomicVisionPreTrainedModel):
             next_cache = next_cache.to_legacy_cache()
 
         if not return_dict:
+            if image != None:
+                return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, mix] if v is not None)
             return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
-        return BaseModelOutputWithPast(
+        return Qwen2NomicVisionOutput(
             last_hidden_state=hidden_states,
             past_key_values=next_cache,
             hidden_states=all_hidden_states,
             attentions=all_self_attns,
+            processed_image=mix,
         )
 
     # Copied from transformers.models.phi3.modeling_phi3.Phi3Model._update_causal_mask
@@ -1196,7 +1206,7 @@ class Qwen2NomicVisionForCausalLM(Qwen2NomicVisionPreTrainedModel, GenerationMix
 
         loss = None
         if labels is not None:
-            loss = self.loss_function(logits[:, 1:], labels, self.vocab_size, **loss_kwargs)
+            loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs)
 
         if not return_dict:
             output = (logits,) + outputs[1:]