123 lines
3.5 KiB
Python
123 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()
|