[fix] existing 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,8 @@ 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

@@ -347,7 +347,8 @@ 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 rp != 1.0:
+                seen = list(set(input_ids[0].tolist())); 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
@@ -370,13 +371,13 @@ class MiniMindOmni(MiniMindForCausalLM):
                     audio_codes[i].append(self.audio_pad_token)
                 else:
                     logits_i = al[0, -1, :].clone() / 0.2
-                    for prev_code in audio_codes[i][-3:]: logits_i[prev_code] /= 1.05
+                    for prev_code in audio_codes[i][-3:]: score = logits_i[prev_code]; logits_i[prev_code] = torch.where(score > 0, score / 1.05, score * 1.05)
                     top_val, top_idx = logits_i.topk(50)
                     code = top_idx[torch.multinomial(F.softmax(top_val, dim=-1), 1)].item()
                     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)