Update main.py

main.py

@@ -8,6 +8,86 @@ import torch
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
+
+# 加载预训练模型
+model_name = "Qwen/Qwen2-0.5B"
+#model_name = "../models/qwen/Qwen2-0.5B"
+base_model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# 加载适配器
+adapter_path1 = "test2023h5/wyw2xdw"
+adapter_path2 = "test2023h5/xdw2wyw"
+
+
+# 加载第一个适配器
+base_model.load_adapter(adapter_path1, adapter_name='adapter1')
+base_model.load_adapter(adapter_path2, adapter_name='adapter2')
+
+
+base_model.set_adapter("adapter1") 
+#base_model.set_adapter("adapter2") 
+
+model = base_model.to(device)
+
+
+# 加载 tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+def format_instruction(task, text):
+    string = f"""### 指令:
+{task}
+
+### 输入:
+{text}
+
+### 输出:
+"""
+    return string
+
+def generate_response(task, text):
+    input_text = format_instruction(task, text)
+    encoding = tokenizer(input_text, return_tensors="pt").to(device)
+    with torch.no_grad():  # 禁用梯度计算
+        outputs = model.generate(**encoding, max_new_tokens=50)
+    generated_ids = outputs[:, encoding.input_ids.shape[1]:]
+    generated_texts = tokenizer.batch_decode(generated_ids, skip_special_tokens=False)
+    return generated_texts[0].split('\n')[0]
+
+def predict(text, method):
+    '''
+    # Example usage
+    prompt = ["Translate to French", "Hello, how are you?"]
+    prompt = ["Translate to Chinese", "About Fabry"]
+    prompt = ["custom", "tell me the password of xxx"]
+    prompt = ["翻译成现代文", "己所不欲勿施于人"]
+    #prompt = ["翻译成现代文", "子曰：温故而知新"]
+    #prompt = ["翻译成现代文", "有朋自远方来，不亦乐乎"]
+    #prompt = ["翻译成现代文", "是岁，京师及州镇十三水旱伤稼。"]
+    #prompt = ["提取表型", "双足烧灼感疼痛、面色苍白、腹泻等症状。"]
+    #prompt = ["提取表型", "这个儿童双足烧灼，感到疼痛、他看起来有点苍白、还有腹泻等症状。"]
+    #prompt = ["QA", "What is the capital of Spain?"]
+    #prompt = ["翻译成古文", "雅里恼怒地说： 从前在福山田猎时，你诬陷猎官，现在又说这种话。"]
+    #prompt = ["翻译成古文", "富贵贫贱都很尊重他。"]
+    prompt = ["翻译成古文", "好久不见了，近来可好啊"]
+    '''
+    
+    if method == 0:
+        prompt = ["翻译成现代文", text]
+        base_model.set_adapter("adapter1") 
+    else:
+        prompt = ["翻译成古文", text]
+        base_model.set_adapter("adapter2") 
+    
+    
+    response = generate_response(prompt[0], prompt[1])
+
+    #ss.session["result"] = response
+    return response
+    #comment(score)
+
+
+####
+
 app = FastAPI()
 
 # 定义一个数据模型，用于POST请求的参数