Spaces:

fdbw
/

fzd12

Configuration error

fzd12 / app.py

Update app.py

94f8878 verified 11 months ago

2.51 kB

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()

	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()