Merge pull request #3 from zhuxiaoxuhit/fix/multiple-bugs

Fix bugs: label leak, repetition penalty, stop condition, O(n²) generation

dataset/omni_dataset.py

@@ -276,7 +276,8 @@ class OmniDataset(Dataset):
         # 生成labels（只训练最后一个assistant）
         text_labels, assistant_ranges = self.generate_text_labels(input_ids)
         for start, end in assistant_ranges[:-1]:
-            text_labels[start:end] = [-100] * (end - start)
+            mask_end = min(end + len(self.eos_id), self.max_length)
+            text_labels[start:mask_end] = [-100] * (mask_end - start)
         
         # 生成7层audio targets（只填充最后一个assistant）
         Y_audio_layers = [[self.audio_pad_token] * self.max_length for _ in range(8)]

model/model_minimind.py

@@ -266,7 +266,10 @@ class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
             attention_mask = torch.cat([attention_mask, attention_mask.new_ones(attention_mask.shape[0], 1)], -1) if attention_mask is not None else None
             logits = outputs.logits[:, -1, :] / temperature
             if repetition_penalty != 1.0:
-                for i in range(input_ids.shape[0]): logits[i, torch.unique(input_ids[i])] /= repetition_penalty
+                for i in range(input_ids.shape[0]):
+                    seen = torch.unique(input_ids[i])
+                    score = logits[i, seen]
+                    logits[i, seen] = torch.where(score > 0, score / repetition_penalty, score * repetition_penalty)
             if top_k > 0: 
                 logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')
             if top_p < 1.0:

model/model_omni.py

@@ -339,6 +339,7 @@ class MiniMindOmni(MiniMindForCausalLM):
         if spk_emb is not None and fill_start > 0:
             audio_buffer[:, :, fill_start - 1] = self.audio_spk_token
         think_end_step, generated_tokens = None, ([] if args.get('open_thinking', False) else None)
+        seen_tokens = set(input_ids[0].tolist()) if rp != 1.0 else None
         while input_ids.shape[1] < start_pos + max_new_tokens:
             if past_kvs is None or not use_cache:
                 out = self.forward(torch.cat((audio_buffer, input_ids.unsqueeze(1)), dim=1), past_key_values=past_kvs, use_cache=use_cache, **args)
@@ -347,7 +348,10 @@ class MiniMindOmni(MiniMindForCausalLM):
             past_kvs = out.past_key_values
 
             logits = out.logits[0, -1, :].clone() / (temperature + 1e-9)
-            logits[list(set(input_ids[0].tolist()))] /= rp
+            if seen_tokens is not None:
+                seen = list(seen_tokens)
+                score = logits[seen]
+                logits[seen] = torch.where(score > 0, score / rp, score * rp)
             if top_p and top_p < 1.0:
                 sorted_l, sorted_i = torch.sort(logits, descending=True)
                 mask = torch.cumsum(F.softmax(sorted_l, dim=-1), dim=-1) > top_p
@@ -358,6 +362,8 @@ class MiniMindOmni(MiniMindForCausalLM):
             if text_finished:
                 text_token = args.get('enter_token_id', 201) if first_finished else args.get('pad_token_id', 0)
                 first_finished = False
+            if seen_tokens is not None:
+                seen_tokens.add(text_token)
 
             step = input_ids.shape[1] - start_pos  # 已生成token数（0=首次，此时模型处理prompt末尾token）
             audio_step = step - 1  # 延迟1步：输出第1个text时无audio，输出第2个text时layer0开始
@@ -376,7 +382,7 @@ class MiniMindOmni(MiniMindForCausalLM):
                     audio_codes[i].append(code)
                     if audio_stop_pos[i] is None and code >= 2048: audio_stop_pos[i] = len(audio_codes[i]) - 1
 
-            if text_finished and audio_codes[7][-1] == self.audio_stop_token: break
+            if text_finished and all(audio_stop_pos[i] is not None for i in range(8)): break
 
             input_ids = torch.cat((input_ids, torch.tensor([[text_token]], device=input_ids.device)), dim=1)
             audio_buffer = torch.cat((audio_buffer, torch.full((1, 8, 1), self.audio_pad_token, dtype=torch.long, device=input_ids.device)), dim=2)