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"
)
print("Model loaded!")

def chat(message, history):
    try:
        # Build conversation from history (each item is [user_msg, bot_msg])
        conversation = ""
        for user_msg, bot_msg in history:
            conversation += f"User: {user_msg}\nAssistant: {bot_msg}\n"
        conversation += f"User: {message}\nAssistant:"

        inputs = tokenizer(
            conversation,
            return_tensors="pt",
            truncation=True,
            max_length=1024
        ).to(model.device)

        outputs = model.generate(
            **inputs,
            max_new_tokens=150,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.1,
            do_sample=True
        )

        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        response = response.split("Assistant:")[-1].strip()
        return response

    except Exception as e:
        return f"Error: {str(e)}"

iface = gr.ChatInterface(
    fn=chat,
    title="DeepSeek Chat AI",
    description="Chat with DeepSeek 1.5B model"
)

iface.launch()