app.py

import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load the model and tokenizer
model_name = "deepseek-ai/DeepSeek-R1"
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

def respond(message, history: list[tuple[str, str]]):
    # Prepare the conversation history
    messages = []
    for user_msg, assistant_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})

    # Tokenize the input
    inputs = tokenizer.apply_chat_template(messages, return_tensors="pt")

    # Generate the response
    outputs = model.generate(inputs, max_length=2048, temperature=0.7, top_p=0.95, do_sample=True)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return response

# Custom ChatInterface with undo and retry buttons
def chat_interface(message, history):
    return respond(message, history)

# Create the Gradio interface
demo = gr.ChatInterface(
    fn=chat_interface,
    retry_btn="Retry",
    undo_btn="Undo",
    clear_btn="Clear",
)

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