chinese_wyw / mymodule.py
test2023h5's picture
Create mymodule.py
ee98572 verified
#调用大模型
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')
# 加载预训练模型
model_name = "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)
print("model loading done")
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):
if method == 0:
prompt = ["翻译成现代文", text]
base_model.set_adapter("adapter1")
else:
prompt = ["翻译成古文", text]
base_model.set_adapter("adapter2")
print("debug", text)
response = generate_response(prompt[0], prompt[1])
print("debug2", response)
return response