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


model_name = "Zhihu-ai/Zhi-writing-dsr1-14"
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)
    
    # 使用流式生成
    for response in 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,
        streamer=gr.TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    ):
        yield response.strip()

# 创建Gradio界面
demo = gr.ChatInterface(
    predict,
    title="Zhi-writing-dsr1-14",
    description="这是一个基于Zhi-writing-dsr1-14的文章生成器。",
    examples=["以鲁迅口吻写一篇500字关于桔了个仔的散文", "用知乎常见的表达方式讲讲什么是AI？", "告诉我一个我大概率不知道的人生哲理"],
    theme=gr.themes.Soft(),
    streaming=True
)

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