import gradio as gr
from transformers import pipeline

# Load the model
pipe = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1", trust_remote_code=True)

# Define the function that will be called by Gradio
def generate_response(message):
    messages = [{"role": "user", "content": message}]
    response = pipe(messages)
    return response[0]['generated_text']

# Create a Gradio interface
iface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=2, placeholder="Enter your message here..."),
    outputs="text",
    title="DeepSeek-R1 Chatbot",
    description="Ask anything to the DeepSeek-R1 model."
)

# Launch the interface
iface.launch()