File size: 2,512 Bytes
94f8878 f83685b 94f8878 f83685b 94f8878 f83685b 94f8878 f83685b 94f8878 61f6bc9 94f8878 61f6bc9 94f8878 f83685b 94f8878 61f6bc9 94f8878 6563cd2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | from transformers import (
AutoTokenizer,
AutoModelForCausalLM,
BitsAndBytesConfig
)
from peft import PeftModel
import torch
import gradio as gr
import os
# === 配置 ===
BASE_MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
ADAPTER_PATH = "./" # 因为模型文件就在 Space 根目录
print("🚀 正在加载基础模型...")
# 加载基础模型（自动使用 CPU）
model = AutoModelForCausalLM.from_pretrained(
BASE_MODEL_NAME,
torch_dtype=torch.float16, # 即使 CPU 也建议用 float16 减少内存
device_map="cpu", # 明确指定 CPU
trust_remote_code=False,
)
print("🔧 正在加载 LoRA 适配器...")
# 加载你微调的 LoRA 权重
model = PeftModel.from_pretrained(model, ADAPTER_PATH)
model.eval() # 切换到推理模式
print("🔤 正在加载分词器...")
tokenizer = AutoTokenizer.from_pretrained(ADAPTER_PATH)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
# === 推理函数 ===
def generate_response(prompt: str, history=None):
try:
# 构造输入（可根据你训练时的格式调整）
input_text = prompt.strip()
inputs = tokenizer(
input_text,
return_tensors="pt",
truncation=True,
max_length=256,
padding=True
).to("cpu")
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=128,
do_sample=True,
temperature=0.7,
top_p=0.9,
pad_token_id=tokenizer.pad_token_id,
eos_token_id=tokenizer.eos_token_id,
)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
# 去掉输入部分，只返回生成内容
if response.startswith(input_text):
response = response[len(input_text):].strip()
return response or "抱歉，我无法回答这个问题。"
except Exception as e:
return f"❌ 推理出错: {str(e)}"
# === Gradio 界面 ===
with gr.Blocks(title="冯氏家谱助手") as demo:
gr.Markdown("# 🧬 冯氏家族知识问答\n基于 TinyLlama 微调的家谱 AI 助手")
chatbot = gr.ChatInterface(
fn=generate_response,
examples=["冯国璋的字辈是什么？", "冯玉祥生于哪一年？", "冯家第几代是‘国’字辈？"],
title="冯氏家谱助手"
)
# 启动
demo.launch() |