from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import gradio as gr
from gtts import gTTS
from langdetect import detect

model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32)

def respond(user_input):
    if not user_input:
        return "Please ask something.", None

    detected_lang = detect(user_input)

    prompt = f"[INST] {user_input} [/INST]"
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=256, pad_token_id=tokenizer.eos_token_id)

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Voice output
    tts = gTTS(text=response, lang='hi' if detected_lang == 'hi' else 'en')
    tts.save("voice_response.mp3")
    return response, "voice_response.mp3"

iface = gr.Interface(
    fn=respond,
    inputs=gr.Textbox(lines=2, placeholder="Ask me anything..."),
    outputs=[gr.Textbox(label="TeachMe Says"), gr.Audio(label="Voice", autoplay=True)],
    title="TeachMe - Your Smart Tutor",
    description="Light AI bot with Hindi + English voice support."
)

iface.launch()