Update train.py

train.py

@@ -49,12 +49,15 @@ with open('data.json', 'r', encoding='utf-8') as f:
 # 假设 data 是一个列表，则使用 Dataset.from_list
 dataset = Dataset.from_list(data)
 
-# 定义预处理函数，对 'instruction' 文本进行 tokenize
+# 预处理数据
 def preprocess_data(example):
-    return tokenizer(example['instruction'], truncation=True, padding="max_length", max_length=128)
+    return {
+        "instruction": example["instruction"],
+        "output": example["output"]
+    }
 
 # 对数据集进行预处理，并移除原始文本列（此处同时移除了 'instruction' 与 'output'，仅保留 tokenize 后的输入）
-dataset = dataset.map(preprocess_data, batched=True, remove_columns=["instruction", "output"])
+dataset = dataset.map(preprocess_data, batched=True)
 
 # 自定义知识蒸馏 Trainer，结合交叉熵损失（hard target）和 KL 散度损失（soft target）
 class DistillationTrainer(Trainer):