feat(training): add evaluation with label masking and padding collator

- Add eval_strategy and eval_steps to training config
- Split dataset into train/eval subsets
- Implement label masking: only assistant responses contribute to loss
- Replace default data collator with MedicalCollator for proper padding
- Set pad_token to eos_token if missing

qwen_domain_expansion/device_config.py

@@ -184,6 +184,8 @@ def select_training_config(
         "learning_rate": 2e-4,
         "num_train_epochs": 1,
         "logging_steps": 5,
+        "eval_strategy": "steps",
+        "eval_steps": 50,
         "save_strategy": "no",
         "report_to": "none",
         **half_kwargs,

qwen_domain_expansion/train_medical_vocab.py

@@ -1,6 +1,7 @@
 import os
 
-from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForLanguageModeling
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
 from peft import LoraConfig, get_peft_model
 from datasets import load_dataset
 
@@ -27,6 +28,8 @@ print_config(cfg)
 # 1. 设置模型和本地数据集
 model_id = cfg["model_id"]
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
 
 # 2. 扩增医学核心专属词汇（防止它们被拆成破碎的单字）
 medical_tokens = ["[并发症_心肌炎]", "阿司匹林肠溶片", "脉弦滑", "靶向免疫治疗"]
@@ -41,17 +44,39 @@ dataset = load_dataset(
     split="train[:1000]"
 )
 
-# 4. 格式化为 Qwen 的 ChatML 对话格式
+split = dataset.train_test_split(test_size=100, seed=42)
+train_dataset_raw = split["train"]
+eval_dataset_raw = split["test"]
+print(f"训练集: {len(train_dataset_raw)} 条，验证集: {len(eval_dataset_raw)} 条")
+
+# 4. 格式化为 Qwen 的 ChatML 对话格式，并对 user 部分做 label mask
 def format_medical_prompts(batch):
-    texts = []
+    max_len = cfg["seq_len"]
+    input_ids_list = []
+    labels_list = []
     for i, o in zip(batch["input"], batch["output"]):
-        # 模拟拼接，让模型在上下文里深刻理解医学问答
-        text = f"<|im_start|>user\n{i}<|im_end|>\n<|im_start|>assistant\n{o}<|im_end|>"
-        texts.append(text)
-    
-    return tokenizer(texts, truncation=True, max_length=cfg["seq_len"], padding=False)
+        prompt = f"<|im_start|>user\n{i}<|im_end|>\n<|im_start|>assistant\n"
+        completion = f"{o}<|im_end|>"
+
+        prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
+        completion_ids = tokenizer(completion, add_special_tokens=False)["input_ids"]
+
+        # 截断：优先保留完整 prompt，剩余空间给 completion
+        if len(prompt_ids) > max_len:
+            prompt_ids = prompt_ids[:max_len]
+        completion_ids = completion_ids[: max_len - len(prompt_ids)]
+
+        ids = prompt_ids + completion_ids
+        # prompt 部分不计入 loss，只有 assistant 回答参与
+        labels = [-100] * len(prompt_ids) + completion_ids
+
+        input_ids_list.append(ids)
+        labels_list.append(labels)
+
+    return {"input_ids": input_ids_list, "labels": labels_list}
 
-tokenized_dataset = dataset.map(format_medical_prompts, batched=True, remove_columns=dataset.column_names)
+tokenized_train = train_dataset_raw.map(format_medical_prompts, batched=True, remove_columns=train_dataset_raw.column_names)
+tokenized_eval = eval_dataset_raw.map(format_medical_prompts, batched=True, remove_columns=eval_dataset_raw.column_names)
 
 # 5. 加载 Qwen 基础模型并调整嵌入层结构
 print(f"正在加载 {model_id} 权重 ({cfg['dtype']})...")
@@ -75,11 +100,40 @@ model.print_trainable_parameters()
 # 7. 使用根据设备内存自动选择的训练参数
 training_args = TrainingArguments(**cfg["training_args"])
 
+
+class MedicalCollator:
+    """Pad input_ids / labels to the longest sample in the batch.
+
+    Padding positions use pad_token_id for input_ids, 0 for attention_mask,
+    and -100 for labels so they are ignored by the loss.
+    """
+
+    def __init__(self, pad_token_id: int):
+        self.pad_token_id = pad_token_id
+
+    def __call__(self, features):
+        max_len = max(len(f["input_ids"]) for f in features)
+        input_ids, attention_mask, labels = [], [], []
+        for f in features:
+            ids = f["input_ids"]
+            lbl = f["labels"]
+            pad_len = max_len - len(ids)
+            input_ids.append(ids + [self.pad_token_id] * pad_len)
+            attention_mask.append([1] * len(ids) + [0] * pad_len)
+            labels.append(lbl + [-100] * pad_len)
+        return {
+            "input_ids": torch.tensor(input_ids, dtype=torch.long),
+            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
+            "labels": torch.tensor(labels, dtype=torch.long),
+        }
+
+
 trainer = Trainer(
     model=model,
-    train_dataset=tokenized_dataset,
+    train_dataset=tokenized_train,
+    eval_dataset=tokenized_eval,
     args=training_args,
-    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
+    data_collator=MedicalCollator(tokenizer.pad_token_id),
 )
 
 print("\n--- 启动医学专有词表扩展 LoRA 训练 ---")