Update app.py

app.py

@@ -1,17 +1,14 @@
-from accelerate import Accelerator
 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 os
 import torch
-
 def main():
-    # 初始化 Accelerator
-    accelerator = Accelerator()
-
     # 基础模型位置
     model_name = "dushuai112233/Qwen2-1.5B-Instruct"
     # 设备
-    device = accelerator.device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
 
     # 加载分词器和模型
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
@@ -39,7 +36,6 @@ def main():
     def tokenize_function(examples):
         return tokenizer(examples['question'], padding='max_length', truncation=True, max_length=128)
 
-    # 对训练集和验证集进行分词处理
     train_dataset = train_dataset.map(tokenize_function, batched=True)
     val_dataset = val_dataset.map(tokenize_function, batched=True)
 
@@ -56,9 +52,6 @@ def main():
         save_total_limit=2,                # 最大保存模型数
     )
 
-    # 将模型移到设备
-    model.to(device)
-
     # Define the Trainer
     trainer = Trainer(
         model=model,                       # 训练的模型
@@ -75,4 +68,4 @@ def main():
     model.save_pretrained('./output')
 
 if __name__ == '__main__':
-    main()
+    main()