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

model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"

print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16
)
print("Model loaded!")

def chat(message, history):
    if message == "":
        return ""
    
    try:
        # Build conversation history safely
        conversation = ""
        
        # Handle history - each item is a list [user_msg, bot_msg]
        for item in history:
            if len(item) >= 2:
                conversation += f"User: {item[0]}\nAssistant: {item[1]}\n"
        
        # Add current message
        conversation += f"User: {message}\nAssistant:"
        
        # Tokenize
        inputs = tokenizer(
            conversation,
            return_tensors="pt",
            truncation=True,
            max_length=1024
        ).to(model.device)
        
        # Generate response
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
        
        # Decode response
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extract only the assistant's reply
        if "Assistant:" in response:
            response = response.split("Assistant:")[-1].strip()
        
        # Remove any thinking tags if present
        if "</think>" in response:
            response = response.split("</think>")[-1].strip()
        
        return response
    
    except Exception as e:
        return f"Error: {str(e)}"

# Create the chat interface (removed theme parameter)
demo = gr.ChatInterface(
    fn=chat,
    title="DeepSeek Chat AI 🤖",
    description="Chat with DeepSeek-R1-Distill-Qwen-1.5B"
)

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