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

# 1. Load the DeepSeek model and tokenizer
model_id = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_id)

print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

# 2. Define the inference function
def chat_function(message, history):
    messages = []
    
    # Smart history parser: handles both old and new Gradio formats smoothly
    for item in history:
        if isinstance(item, dict):  # New Gradio format
            messages.append({"role": item["role"], "content": item["content"]})
        elif isinstance(item, (list, tuple)):  # Older Gradio format [[user, bot], ...]
            messages.append({"role": "user", "content": item[0]})
            messages.append({"role": "assistant", "content": item[1]})
        
    # Add the user's newest message
    messages.append({"role": "user", "content": message})
    
    # Use the tokenizer's built-in template
    formatted_prompt = tokenizer.apply_chat_template(
        messages, 
        tokenize=False, 
        add_generation_prompt=True
    )

    # Tokenize the formatted prompt
    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
    
    # Generate the response
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            do_sample=True,
            temperature=0.6,
            top_p=0.95
        )
    
    # Decode only the newly generated tokens
    input_length = inputs.input_ids.shape[1]
    generated_tokens = outputs[0][input_length:]
    response = tokenizer.decode(generated_tokens, skip_special_tokens=True)
    
    return response

# 3. Launch the Gradio Interface
# Removed the `type` argument to ensure compatibility across all versions
demo = gr.ChatInterface(
    fn=chat_function,
    title="DeepSeek-R1 1.5B Chatbot",
    description="Running the open DeepSeek-R1-Distill-Qwen-1.5B model locally inside a Space."
)

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