Aya / app.py
Tonic's picture
Update app.py
ff47bef verified
raw
history blame
1.47 kB
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.bfloat16, low_cpu_mem_usage=True, device_map=device)
@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()