import spaces
import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import torch

title = """# 👋🏻Welcome To 🌟Tonic's🌐Aya-101"""
description = """ try this space to build downstream applications with [CohereForAI/aya-101](https://huggingface.co/CohereForAI/aya-101) use via API ;-) """

device = "cuda"
checkpoint = "CohereForAI/aya-101"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint, torch_dtype=torch.float16, low_cpu_mem_usage=True, device_map="auto")

@spaces.GPU
def aya(text):
    """
    Translates the input text to English using the Aya model.
    Assumes the model can automatically detect the input language.
    """
    model.to(device)

    inputs = tokenizer.encode(text, return_tensors="pt").to(device)
    
    outputs = model.generate(inputs, max_new_tokens=128)
    
    
    translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return translation

def main():
    with gr.Blocks() as demo:
        gr.Markdown(title)
        gr.Markdown(description)
        output_text = gr.Textbox(label="🌐Aya", interactive=False)
        with gr.Row():
            input_text = gr.Textbox(label="🗣️Input Text")
            submit_button = gr.Button("Translate")  # Add a button to submit the input
        submit_button.click(fn=aya, inputs=input_text, outputs=output_text)
    
    demo.launch()

if __name__ == "__main__":
    main()