import gradio as gr
from transformers import pipeline

# 加载 Hugging Face 上的预训练模型（以 DeepSeek 为例）
from huggingface_hub import login
from transformers import pipeline

# 登录 Hugging Face
login(token="your_huggingface_token")

# 加载模型
model_name = "deepseek-ai/deepseek-7b"
pipe = pipeline("text-generation", model=model_name, tokenizer=model_name)


# 定义与模型对话的函数
def chat_with_ai(prompt):
    # 使用模型生成文本
    response = pipe(prompt, max_length=100, do_sample=True)
    return response[0]["generated_text"]

# 创建 Gradio 界面
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 AI Chatbot powered by DeepSeek")
    
    # 聊天框组件
    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="Type your message:")
    clear = gr.Button("Clear")

    # 定义响应函数，处理用户输入并更新聊天历史
    def respond(message, chat_history):
        response = chat_with_ai(message)
        chat_history.append((message, response))  # 将对话记录添加到聊天历史
        return "", chat_history

    # 提交消息并更新聊天记录
    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    # 清空聊天记录
    clear.click(lambda: [], None, chatbot)

# 启动 Gradio 界面
demo.launch()