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

# 本地模型路径（绝对路径）
LOCAL_MODEL_PATH = os.path.join(os.getcwd(), "model")
print("Loading model from:", LOCAL_MODEL_PATH)

# 加载 tokenizer 和模型
tokenizer = AutoTokenizer.from_pretrained(
    LOCAL_MODEL_PATH,
    local_files_only=True,
    trust_remote_code=False
)

model = AutoModelForCausalLM.from_pretrained(
    LOCAL_MODEL_PATH,
    local_files_only=True,
    trust_remote_code=False,
    torch_dtype=torch.float16
).to("cuda" if torch.cuda.is_available() else "cpu")

# 聊天函数
def chat(user_input):
    inputs = tokenizer(user_input, return_tensors="pt").to(model.device)
    outputs = model.generate(
        input_ids=inputs["input_ids"],
        attention_mask=inputs.get("attention_mask"),
        max_new_tokens=200,
        temperature=0.7,
        top_p=0.9,
        do_sample=True
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)

# Gradio 网页界面
interface = gr.Interface(
    fn=chat,
    inputs=gr.Textbox(lines=3, label="输入你的问题"),
    outputs=gr.Textbox(lines=10, label="模型回答"),
    title="我的AI模型",
    description="在 HuggingFace Space 上运行的 LLaMA/GPT 模型网页"
)

interface.launch()