import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "ai4bharat/Airavata"

# Load the model in 8-bit precision to reduce memory usage
model = AutoModelForCausalLM.from_pretrained(
    model_name, 
    device_map="auto", 
    load_in_8bit=True
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

def generate_text(prompt, max_length):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=max_length)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

interface = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.inputs.Textbox(label="Enter your prompt"),
        gr.inputs.Slider(10, 100, step=10, label="Max length")
    ],
    outputs="text",
    title="Airavata Text Generation Model",
    description="Generate text in Indic languages using the Airavata model."
)

interface.launch()