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

# Use a smaller instruct-tuned model that runs on Hugging Face Spaces
model_name = "tiiuae/falcon-7b-instruct"  # Falcon-7B is lighter than Mistral

# Load model and tokenizer with optimizations
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name, 
    torch_dtype=torch.float16,  
    device_map="auto"  # Uses available GPU/CPU
)

# AI Response Function
def nithin_ai(question):
    inputs = tokenizer(question, return_tensors="pt").input_ids.to(model.device)
    outputs = model.generate(inputs, max_length=200)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Gradio Chat Interface
iface = gr.Interface(
    fn=nithin_ai,
    inputs="text",
    outputs="text",
    title="Nithin AI - Student Doubt Solver",
    description="Ask any question related to robotics, science, or math!"
)

iface.launch()