Browse Source

fix: update past_key_values handling to support custom sequence lengths

myhloli 4 months ago
parent
commit
e7ed7660e3

+ 12 - 2
mineru/model/mfr/unimernet/unimernet_hf/unimer_mbart/modeling_unimer_mbart.py

@@ -1416,7 +1416,13 @@ class UnimerMBartDecoder(UnimerMBartPreTrainedModel):
             raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
 
         # past_key_values_length
-        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
+        # past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
+        past_key_values_length = 0
+        if past_key_values is not None:
+            if hasattr(past_key_values, 'get_seq_length'):
+                past_key_values_length = past_key_values.get_seq_length()
+            elif isinstance(past_key_values, (list, tuple)) and past_key_values:
+                past_key_values_length = past_key_values[0][0].shape[2]
 
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
@@ -1501,7 +1507,11 @@ class UnimerMBartDecoder(UnimerMBartPreTrainedModel):
                 if dropout_probability < self.layerdrop:
                     continue
 
-            past_key_value = past_key_values[idx] if past_key_values is not None else None
+            # past_key_value = past_key_values[idx] if past_key_values is not None else None
+            past_key_value = None
+            if past_key_values is not None and len(past_key_values) > idx:
+                if hasattr(past_key_values, 'get_usable_length') or isinstance(past_key_values, (list, tuple)):
+                    past_key_value = past_key_values[idx]
 
             if self.gradient_checkpointing and self.training:
                 layer_outputs = self._gradient_checkpointing_func(