Update modeling_minicpmo.py

modeling_minicpmo.py

@@ -762,17 +762,39 @@ class MiniCPMO(MiniCPMOPreTrainedModel):
                     for i in range(bs):
                         audio_embs = audio_embeddings[i]
                         bounds = audio_bounds[i]
+
+                    one_to_one_match = len(audio_embs) == len(bounds) and all(
+                        embs.shape[0] == int(bound[1] - bound[0])
+                        for embs, bound in zip(audio_embs, bounds)
+                    )
+
+                    if one_to_one_match:
                         for embs, bound in zip(audio_embs, bounds):
-                            audio_indices = torch.arange(bound[0], bound[1], dtype=torch.long).to(
-                                input_embeddings.device
+                            input_embeddings[i, bound[0] : bound[1]] = embs.to(
+                                device=input_embeddings.device,
+                                dtype=input_embeddings.dtype,
                             )
-
-                            if embs.shape[0] != len(audio_indices):
-                                raise ValueError(
-                                    f"Shape mismatch: Trying to assign embeddings of shape {embs.shape} "
-                                    f"to input indices of length {len(audio_indices)}"
-                                )
-                            input_embeddings[i, audio_indices] = embs.to(input_embeddings.dtype)
+                    else:
+                        flat_audio_embs = torch.cat(audio_embs, dim=0).to(
+                            device=input_embeddings.device,
+                            dtype=input_embeddings.dtype,
+                        )
+            
+                        total_bound_len = sum(int(bound[1] - bound[0]) for bound in bounds)
+            
+                        if flat_audio_embs.shape[0] != total_bound_len:
+                            raise ValueError(
+                                f"Audio total length mismatch: {flat_audio_embs.shape[0]} != {total_bound_len}"
+                            )
+            
+                        offset = 0
+                        for bound in bounds:
+                            audio_len = int(bound[1] - bound[0])
+                            input_embeddings[i, bound[0] : bound[1]] = flat_audio_embs[
+                                offset : offset + audio_len
+                            ]
+                            offset += audio_len
+                            
         elif self.training:
             for i in range(bs):
                 # dummy audio_embedings