# app.py
import os
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import gradio as gr

# 设置缓存目录（节省空间）
os.environ["HF_HOME"] = "/tmp/hf_cache"

# 模型名称（来自 Hugging Face）
model_name = "UnfilteredAI/DAN-L3-R1-8B"

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 创建文本生成 pipeline
pipe = pipeline(
    "text-generation",
    model=model_name,
    tokenizer=tokenizer,
    model_kwargs={"torch_dtype": "auto"},  # 自动选择精度（如 float16）
    device_map="auto",  # 自动使用 GPU（如果有）
    max_new_tokens=512,
    temperature=0.7,
    top_p=0.9,
    repetition_penalty=1.1,
    return_full_text=False,  # 只返回生成的回复，不包含输入
)

# 定义生成函数
def generate_response(user_input, history):
    # 将对话历史转换为 messages 格式
    messages = []
    for idx, (user_msg, assistant_msg) in enumerate(history):
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": assistant_msg})
    
    # 添加当前用户输入
    messages.append({"role": "user", "content": user_input})

    # 调用模型生成回复
    try:
        outputs = pipe(messages)
        response = outputs[0]["generated_text"].strip()
    except Exception as e:
        response = f"❌ 模型生成出错：{str(e)}"

    # 返回更新后的对话历史
    history.append((user_input, response))
    return "", history

# 创建 Gradio 界面
with gr.Blocks(title="💬 DAN-L3-R1-8B AI 助手", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 💬 DAN-L3-R1-8B AI 助手")
    gr.Markdown("基于 [UnfilteredAI/DAN-L3-R1-8B](https://huggingface.co/UnfilteredAI/DAN-L3-R1-8B) 的对话系统")

    chatbot = gr.Chatbot(height=600, show_copy_button=True)
    with gr.Row():
        textbox = gr.Textbox(
            placeholder="请输入你的问题...",
            show_label=False,
            scale=7,
            container=False,
        )
        submit_btn = gr.Button("发送", scale=1)

    # 示例问题
    gr.Examples(
        examples=[
            "你是谁？",
            "讲个笑话",
            "写一首关于秋天的诗",
            "解释量子力学"
        ],
        inputs=textbox
    )

    # 清空按钮
    clear_btn = gr.Button("🗑️ 清空聊天")

    # 事件绑定
    textbox.submit(fn=generate_response, inputs=[textbox, chatbot], outputs=[textbox, chatbot])
    submit_btn.click(fn=generate_response, inputs=[textbox, chatbot], outputs=[textbox, chatbot])
    clear_btn.click(fn=lambda: ("", None), outputs=[textbox, chatbot])

# 启动应用
if __name__ == "__main__":
    demo.launch()