Update main.py

main.py

@@ -1,6 +1,13 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel, get_peft_config
+import json
+import torch
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
 app = FastAPI()
 
 # 定义一个数据模型，用于POST请求的参数