Update app.py

app.py

@@ -1,74 +1,76 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
 from peft import LoraConfig, get_peft_model, TaskType
 from datasets import load_dataset
-from torch.utils.tensorboard import SummaryWriter
 import torch
 import os
 
 def main():
     # 基础模型位置
-    model_name = "dushuai112233/Qwen2-1.5B-Instruct"  # 使用你提供的模型
-    # 设备
+    model_name = "dushuai112233/Qwen2-1.5B-Instruct"
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
     # 加载分词器和模型
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
     model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
 
-    # Setup PEFT (Low-Rank Adaption)
+    # Setup PEFT
     peft_config = LoraConfig(
         task_type=TaskType.CAUSAL_LM,
         target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
         inference_mode=False,
-        r=8,  # 低秩矩阵的秩
-        lora_alpha=32,  # LoRA的alpha超参数
+        r=8,
+        lora_alpha=32,
         lora_dropout=0.1
     )
     model = get_peft_model(model, peft_config)
 
-    # 加载 Hugging Face 数据集
-    ds = load_dataset("dushuai112233/medical")  # 自动加载 train 和 val 分区
-
-    # 提取训练集和验证集
+    # 加载数据集
+    ds = load_dataset("dushuai112233/medical")
     train_dataset = ds["train"]
     val_dataset = ds["validation"]
 
-    # Tokenize the datasets
+    # 数据集预处理
     def tokenize_function(examples):
-        # 注意: 对于 Causal LM，通常会使用输入文本作为标签（shifted label）
         encodings = tokenizer(examples['question'], padding='max_length', truncation=True, max_length=128)
-        encodings['labels'] = encodings['input_ids'].copy()  # Causal LM labels should be same as input_ids
+        encodings['labels'] = encodings['input_ids'].copy()
         return encodings
 
     train_dataset = train_dataset.map(tokenize_function, batched=True)
     val_dataset = val_dataset.map(tokenize_function, batched=True)
 
-    # Define Training Arguments
+    # 设置训练参数
     training_args = TrainingArguments(
-        output_dir="./output",             # 保存模型和日志的路径
-        evaluation_strategy="epoch",       # 每个epoch后进行验证
-        per_device_train_batch_size=1,     # 每个设备的batch size
-        per_device_eval_batch_size=1,      # 验证时的batch size
-        logging_dir="./logs",              # 日志目录
-        logging_steps=10,                  # 每10步记录一次日志
-        save_steps=100,                    # 每100步保存一次模型
-        num_train_epochs=10,               # 训练的epoch数
-        save_total_limit=2,                # 最大保存模型数
+        output_dir="./output",
+        evaluation_strategy="epoch",
+        per_device_train_batch_size=1,
+        per_device_eval_batch_size=1,
+        logging_dir="./logs",
+        logging_steps=10,
+        save_steps=100,                 # 每 100 步保存一次检查点
+        save_total_limit=2,             # 限制最多保存 2 个检查点
+        num_train_epochs=10,
+        load_best_model_at_end=False,   # 是否在训练结束时加载最优模型
     )
 
-    # Define the Trainer
+    # 定义 Trainer
     trainer = Trainer(
-        model=model,                       # 训练的模型
-        args=training_args,                # 训练的参数
-        train_dataset=train_dataset,       # 训练数据集
-        eval_dataset=val_dataset,          # 验证数据集
-        tokenizer=tokenizer,               # 用于预处理的分词器
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=val_dataset,
+        tokenizer=tokenizer,
     )
 
-    # Start Training
-    trainer.train()
+    # 检查是否有中断点
+    checkpoint = None
+    if os.path.exists("./output") and len(os.listdir("./output")) > 0:
+        checkpoint = max([os.path.join("./output", ckpt) for ckpt in os.listdir("./output")], key=os.path.getmtime)
+        print(f"Resuming training from checkpoint: {checkpoint}")
+
+    # 开始训练
+    trainer.train(resume_from_checkpoint=checkpoint)
 
-    # Save the model
+    # 保存最终模型
     model.save_pretrained('./output')
 
 if __name__ == '__main__':