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

# 加载模型和分词器
model_name = "Zhihu-ai/Zhi-writing-dsr1-14b"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True
)
@spaces.GPU()
def predict(message, history):
    # 构建输入
    history_text = ""
    for human, assistant in history:
        history_text += f"Human: {human}\nAssistant: {assistant}\n"
    prompt = f"{history_text}Human: {message}\nAssistant:"
    
    # 生成回复
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=10000,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.1,
        pad_token_id=tokenizer.eos_token_id
    )
    response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
    
    return response.strip()

# 创建Gradio界面
demo = gr.ChatInterface(
    predict,
    title="测试Zhi-writing-dsr1-14b",
    description="Zhihu-ai/Zhi-writing-dsr1-14b",
    examples=["鲁迅口吻写五百字，描述桔猫的可爱！", "桔了个仔是谁", "介绍自己"],
    theme=gr.themes.Soft()
)

# with gr.Blocks(theme=gr.themes.Soft()) as demo:
#     gr.Markdown("# Zhi-writing-dsr1-14")
#     gr.Markdown("这是一个基于Zhi-writing-dsr1-14的文章生成器")
    
#     chatbot = gr.Chatbot()
#     msg = gr.Textbox(label="输入消息")
#     clear = gr.Button("清除对话")
    
#     def respond(message, chat_history):
#         bot_message = ""
#         for response in predict(message, chat_history):
#             bot_message = response
#             chat_history.append((message, bot_message))
#             yield chat_history
#         return "", chat_history
    
#     msg.submit(respond, [msg, chatbot], [msg, chatbot])
#     clear.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch(share=True)