import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 你自己的模型 repo
model_id = "Seunggg/lora-plant"

# 加载模型和 tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, device_map="auto")

# 定义接口函数
def plant_chat(user_input):
    prompt = f"用户提问：{user_input}\n请用人性化语言回答，并推荐相关的植物资料或文献：\n回答："
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=256)
    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return answer

# 启动 Gradio 接口
gr.Interface(fn=plant_chat,
             inputs="text",
             outputs="text",
             title="🌿 植物问答助手",
             description="根据你的问题，提供植物养护建议和文献线索。").launch()