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

# Load Anki-2.5 model
model_name = "anktechsol/anki-2.5"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float32,
    device_map="auto"
)

def respond(message, history, system_message, max_tokens, temperature, top_p):
    """
    Indian Language Q&A powered by Anki-2.5
    Supports Hindi, Tamil, Bengali, Telugu, and other Indian languages
    """
    # Build prompt manually
    prompt = system_message + "\n\n"
    
    # Add history
    for msg in history:
        if isinstance(msg, dict):
            role = msg.get("role", "")
            content = msg.get("content", "")
            if role == "user":
                prompt += f"User: {content}\n"
            elif role == "assistant":
                prompt += f"Assistant: {content}\n"
        elif len(msg) == 2:
            prompt += f"User: {msg[0]}\n"
            prompt += f"Assistant: {msg[1]}\n"
    
    # Add current message
    prompt += f"User: {message}\nAssistant:"
    
    # Generate response
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            inputs["input_ids"],
            max_new_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    
    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
    return response

# Create ChatInterface for multilingual Q&A
chatbot = gr.ChatInterface(
    respond,
    title="🚀 Indian Language Q&A Assistant - Powered by Anki-2.5",
    description="Ask questions in Hindi, Tamil, Bengali, Telugu, and other Indian languages!",
    examples=[
        ["भारत की राजधानी क्या है?"],
        ["What is the capital of India?"],
        ["ভারতের স্বাধীনতা দিবস কবে?"],
        ["தமிழ் மொழியில் எண்ணிக்கையில் பேசுவது எப்படி?"],
        ["Explain the importance of education"],
    ],
    additional_inputs=[
        gr.Textbox(
            value="You are a helpful multilingual assistant that can answer questions in multiple Indian languages including Hindi, English, Bengali, Tamil, Telugu, and others. Provide accurate and helpful responses.",
            label="System Message"
        ),
        gr.Slider(minimum=50, maximum=1024, value=512, step=1, label="Max Tokens"),
        gr.Slider(minimum=0.1, maximum=1.5, 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"),
    ],
    theme=gr.themes.Soft(),
)

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