Update modeling_minicpmo.py

modeling_minicpmo.py

@@ -763,37 +763,36 @@ class MiniCPMO(MiniCPMOPreTrainedModel):
                         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)
-                    )
+                        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):
-                            input_embeddings[i, bound[0] : bound[1]] = embs.to(
+                        if one_to_one_match:
+                            for embs, bound in zip(audio_embs, bounds):
+                                input_embeddings[i, bound[0] : bound[1]] = embs.to(
+                                    device=input_embeddings.device,
+                                    dtype=input_embeddings.dtype,
+                                )
+                        else:
+                            flat_audio_embs = torch.cat(audio_embs, dim=0).to(
                                 device=input_embeddings.device,
                                 dtype=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
+
+                            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):
@@ -853,10 +852,12 @@ class MiniCPMO(MiniCPMOPreTrainedModel):
         result_text = []
         for result in result_ids:
             result = result[result != 0]
-            if result[0] == tokenizer.bos_id:
+            if len(result) > 0 and result[0] == tokenizer.bos_id:
                 result = result[1:]
-            if result[-1] in terminators:
-                result = result[:-1]
+            for idx, token_id in enumerate(result):
+                if token_id in terminators:
+                    result = result[:idx]
+                    break
             result_text.append(tokenizer.decode(result))
         return result_text
 
@@ -1129,8 +1130,10 @@ class MiniCPMO(MiniCPMOPreTrainedModel):
         input_images_list = []
         input_audios_list = []
         audio_parts_list = []
+        use_tts_template_list = []
 
         for image, msgs in zip(images_list, msgs_list):
+            sample_use_tts_template = use_tts_template
             if isinstance(msgs, str):
                 msgs = json.loads(msgs)
             copy_msgs = deepcopy(msgs)
@@ -1161,7 +1164,7 @@ class MiniCPMO(MiniCPMOPreTrainedModel):
                         audios.append(c)
                         audio_parts.append(i)
                         cur_msgs.append("<audio>./</audio>")
-                        use_tts_template = True
+                        sample_use_tts_template = True
                     elif isinstance(c, str):
                         cur_msgs.append(c)
 
@@ -1175,13 +1178,18 @@ class MiniCPMO(MiniCPMOPreTrainedModel):
                     copy_msgs,
                     tokenize=False,
                     add_generation_prompt=False if teacher_forcing else True,
-                    use_tts_template=use_tts_template,
+                    use_tts_template=sample_use_tts_template,
                     enable_thinking=enable_thinking,
                 )
             )
             input_images_list.append(images)
             input_audios_list.append(audios)
             audio_parts_list.append(audio_parts)
+            use_tts_template_list.append(sample_use_tts_template)
+
+        # Preserve the original downstream behavior for optional audio generation,
+        # while keeping prompt template selection local to each batch sample.
+        use_tts_template = any(use_tts_template_list)
 
         if not merge_audio_from_same_content:
             audio_parts_list = None
@@ -1253,9 +1261,8 @@ class MiniCPMO(MiniCPMOPreTrainedModel):
 
         tts_bound = (tts_bos_idx, tts_eos_idx)
 
-        answer = res[0]
-        if answer is not None:
-            answer = answer.split("<|tts_eos|>")[0]
+        answers = [answer.split("<|tts_eos|>")[0] if answer is not None else None for answer in res]
+        answer = answers if batched else answers[0]
 
         if use_tts_template and generate_audio and output_audio_path:
             import soundfile as sf
@@ -1284,7 +1291,7 @@ class MiniCPMO(MiniCPMOPreTrainedModel):
                 traceback.print_exc()
 
         if return_prompt:
-            return answer, prompts_lists[0]
+            return answer, prompts_lists if batched else prompts_lists[0]
         else:
             return answer
 

processing_minicpmo.py

@@ -1471,7 +1471,7 @@ class MiniCPMOProcessor(ProcessorMixin):
                     audio_feature_lens = torch.hstack(audio_feature_lens)
                     audio_feature_lens_list.append(audio_feature_lens)
             else:
-                audio_feature_lens_list.append([])
+                audio_feature_lens_list.append(torch.empty(0, dtype=torch.long))
 
         if audio_features_all:
             audio_features = [i.permute(1, 0) for i in audio_features_all]