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):
    """
    Educational Tutor for Indian Students powered by Anki-2.5
    Helps with homework, exam prep, and learning in 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 focused on educational tutoring
chatbot = gr.ChatInterface(
    respond,
    title="🎓 Indian Education Tutor - Powered by Anki-2.5",
    description="Get help with homework, exam preparation, and learning in Indian languages. Ask questions about any subject!",
    examples=[
        ["धातु क्या है? धातुओं के उदाहरण दीजिए"],
        ["Explain photosynthesis in simple terms"],
        ["Help me solve this math problem: 2x + 5 = 15"],
        ["প্রথম বিশ্বযুদ্ধের কারণ কী?"],
        ["What are the main parts of a plant?"],
    ],
    additional_inputs=[
        gr.Textbox(
            value="You are an educational tutor for Indian students. You help with homework, exam preparation, and learning across all subjects including math, science, social studies, and languages. You can explain concepts in Hindi, English, Bengali, and other Indian languages. Be patient, encouraging, and provide step-by-step explanations.",
            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()