import gradio as gr
from transformers import pipeline

# Initialize the pipeline
pipe = pipeline("Summarization", model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B")

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    # Format the conversation history
    messages = [{"role": "system", "content": system_message}]
    
    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})
    
    # Add the current message
    messages.append({"role": "user", "content": message})
    
    # Convert messages to a single string format that the model can understand
    prompt = "\n".join([f"{msg['role']}: {msg['content']}" for msg in messages])
    
    # Generate response using the pipeline
    response = ""
    for output in pipe(
        prompt,
        max_new_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        stream=True,
    ):
        # Extract the generated text
        new_text = output[0]['generated_text'][len(response):]
        response = output[0]['generated_text']
        yield new_text

# Create the Gradio interface
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(
            value="You are a friendly and helpful assistant.",
            label="System message"
        ),
        gr.Slider(
            minimum=1,
            maximum=2048,
            value=512,
            step=1,
            label="Max new tokens"
        ),
        gr.Slider(
            minimum=0.1,
            maximum=4.0,
            value=0.7,
            step=0.1,
            label="Temperature"
        ),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)"
        ),
    ],
    title="DeepSeek Chat Interface",
    description="Chat with the DeepSeek-R1-Distill-Qwen-1.5B model",
)

# Launch the interface
if __name__ == "__main__":
    demo.launch()