import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load model
tokenizer = AutoTokenizer.from_pretrained("ai4bharat/Airavata")
model = AutoModelForCausalLM.from_pretrained("ai4bharat/Airavata")

def chatbot(input_text):
    # Encode the input text
    input_text = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors='pt')

    # Generate a response
    response = model.generate(input_text, max_length=1000, pad_token_id=tokenizer.eos_token_id)

    # Decode the response
    response_text = tokenizer.decode(response[:, input_text.shape[-1]:][0], skip_special_tokens=True)
    
    return response_text

iface = gr.Interface(fn=chatbot, inputs="text", outputs="text")
iface.launch()