74 lines
3.0 KiB
Python
74 lines
3.0 KiB
Python
import torch
|
||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||
from peft import PeftModel
|
||
from modelscope import snapshot_download
|
||
|
||
def generate_reply(model, tokenizer, prompt_text, system_prompt="你是由通信领域与知识产权专家调优的专业助手，擅长分析通信技术协议与专利权利要求书结构。"):
|
||
messages = [
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": prompt_text}
|
||
]
|
||
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
||
inputs = tokenizer([text], return_tensors="pt").to(model.device)
|
||
|
||
with torch.no_grad():
|
||
outputs = model.generate(
|
||
**inputs,
|
||
max_new_tokens=256,
|
||
temperature=0.7,
|
||
top_p=0.9,
|
||
do_sample=True,
|
||
repetition_penalty=1.1
|
||
)
|
||
# 截取生成的新 token
|
||
generated_ids = outputs[0][inputs.input_ids.shape[1]:]
|
||
return tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
|
||
|
||
def main():
|
||
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
||
print(f"[Device] 当前推理设备: {device}")
|
||
|
||
model_id = "Qwen/Qwen2.5-0.5B-Instruct"
|
||
lora_dir = "output/qwen_telecom_lora"
|
||
model_dir = snapshot_download(model_id)
|
||
|
||
print("[ModelScope] 正在加载基座模型 Tokenizer 与 Model ...")
|
||
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
|
||
base_model = AutoModelForCausalLM.from_pretrained(
|
||
model_dir,
|
||
torch_dtype=torch.bfloat16,
|
||
device_map=device,
|
||
trust_remote_code=True
|
||
)
|
||
|
||
test_prompt = (
|
||
"请针对以下公开专利进行技术方案与权利要求拆解：\n"
|
||
"【专利名称】：5G毫米波大规模天线阵列混合波束赋形方法及收发机\n"
|
||
"【技术方案】：针对高频通信中射频链路(RF Chain)硬件受限问题，提出一种模拟移相网络与数字基带迫零(ZF)预编码分步解耦优化架构，在保证系统信噪比的同时大幅降低硬件功耗。"
|
||
)
|
||
|
||
print("正在测试原版 Base 模型...")
|
||
base_reply = generate_reply(base_model, tokenizer, test_prompt)
|
||
|
||
print("正在挂载 LoRA 适配层权重...")
|
||
lora_model = PeftModel.from_pretrained(base_model, lora_dir)
|
||
print("正在测试 LoRA 微调后模型...")
|
||
lora_reply = generate_reply(lora_model, tokenizer, test_prompt)
|
||
|
||
result_text = []
|
||
result_text.append("="*40 + " 1. 原版 Base 模型输出 " + "="*40)
|
||
result_text.append(base_reply)
|
||
result_text.append("\n" + "="*40 + " 2. 加载 LoRA 微调权重后输出 " + "="*40)
|
||
result_text.append(lora_reply)
|
||
result_text.append("\n" + "="*90)
|
||
|
||
full_output = "\n".join(result_text)
|
||
print(full_output)
|
||
|
||
with open("output/comparison_result.txt", "w", encoding="utf-8") as f:
|
||
f.write(full_output)
|
||
print("\n[Done] 完整中文对比结果已保存到 output/comparison_result.txt")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|