122 lines
3.5 KiB
Python
122 lines
3.5 KiB
Python
# %%
|
||
import os
|
||
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
|
||
import torch
|
||
import pandas as pd
|
||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||
from tqdm import tqdm
|
||
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
|
||
# %%
|
||
# 加载微调后的完整模型（main.py 保存的 final_model 目录）
|
||
model_path = "./output/final_model"
|
||
|
||
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||
if tokenizer.pad_token is None:
|
||
tokenizer.pad_token = tokenizer.eos_token
|
||
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
model_path,
|
||
torch_dtype=torch.bfloat16,
|
||
trust_remote_code=True,
|
||
)
|
||
model.to(device)
|
||
model.eval()
|
||
print("模型加载完成:", model_path)
|
||
|
||
# %%
|
||
PROMPT_PREFIX = (
|
||
"你是一名材料科学助手，请根据问题给出准确、专业的回答。\n"
|
||
"问题：{q}\n"
|
||
"回答："
|
||
)
|
||
|
||
EOS_TOKEN = tokenizer.eos_token
|
||
|
||
|
||
def build_prompt(question):
|
||
return PROMPT_PREFIX.format(q=question)
|
||
|
||
|
||
@torch.no_grad()
|
||
def chat(
|
||
question,
|
||
max_new_tokens=256,
|
||
temperature=0.7,
|
||
top_p=0.9,
|
||
repetition_penalty=1.1,
|
||
system_prompt=None,
|
||
):
|
||
prefix = system_prompt if system_prompt else PROMPT_PREFIX
|
||
input_text = prefix.format(q=question)
|
||
inputs = tokenizer(input_text, return_tensors="pt").to(device)
|
||
output = model.generate(
|
||
**inputs,
|
||
max_new_tokens=max_new_tokens,
|
||
do_sample=True,
|
||
temperature=temperature,
|
||
top_p=top_p,
|
||
repetition_penalty=repetition_penalty,
|
||
pad_token_id=tokenizer.pad_token_id,
|
||
eos_token_id=tokenizer.eos_token_id,
|
||
)
|
||
response = tokenizer.decode(
|
||
output[0][inputs["input_ids"].shape[1]:],
|
||
skip_special_tokens=True,
|
||
)
|
||
return response.strip()
|
||
|
||
|
||
def interactive():
|
||
"""交互式问答（输入 quit/exit/q 退出）"""
|
||
print("=" * 60)
|
||
print("材料科学问答助手（输入 quit / exit / q 退出）")
|
||
print("=" * 60)
|
||
while True:
|
||
try:
|
||
q = input("\n问题：").strip()
|
||
except (EOFError, KeyboardInterrupt):
|
||
break
|
||
if not q:
|
||
continue
|
||
if q.lower() in ("quit", "exit", "q"):
|
||
break
|
||
print("回答：", chat(q))
|
||
print("已退出。")
|
||
|
||
|
||
# %%
|
||
# 从原始 csv 取问题做批量回测，对比模型回答与标准答案
|
||
def batch_eval(sample_n=10):
|
||
df = pd.read_csv("样本收集.csv", encoding="utf-8").dropna(subset=["问题", "回答"]).reset_index(drop=True)
|
||
df["问题"] = df["问题"].astype(str).str.replace("\u3000", " ").str.replace(r"\s+", " ", regex=True).str.strip()
|
||
df["回答"] = df["回答"].astype(str).str.replace("\u3000", " ").str.replace(r"\s+", " ", regex=True).str.strip()
|
||
|
||
sample_n = min(sample_n, len(df))
|
||
for i in tqdm(range(sample_n), desc="回测中"):
|
||
q = str(df.loc[i, "问题"])
|
||
a = str(df.loc[i, "回答"])
|
||
pred = chat(q)
|
||
print(f"\n{'='*60}")
|
||
print(f"[{i+1}] Q: {q}")
|
||
print(f" 标准答案: {a}")
|
||
print(f" 模型回答: {pred}")
|
||
|
||
|
||
# %%
|
||
if __name__ == "__main__":
|
||
import sys
|
||
|
||
# 用法：
|
||
# python test.py -> 默认进入交互式问答
|
||
# python test.py --eval -> 跑 csv 批量回测
|
||
# python test.py "你的问题" -> 单次问答
|
||
if len(sys.argv) > 1:
|
||
if sys.argv[1] == "--eval":
|
||
batch_eval(sample_n=10)
|
||
else:
|
||
q = " ".join(sys.argv[1:])
|
||
print("回答：", chat(q))
|
||
else:
|
||
interactive()
|