Aya / app.py
Tonic's picture
Create app.py
d204ef0 verified
raw
history blame
1.4 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 ;-) """
checkpoint = "CohereForAI/aya-101"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
if device == "cuda":
model = model.half()
@spaces.GPU
def translate(text):
"""
Translates the input text to English using the Aya model.
Assumes the model can automatically detect the input language.
"""
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)
with gr.Row():
input_text = gr.Textbox(label="Input Text")
output_text = gr.Textbox(label="🌐Aya", interactive=False)
input_text.change(fn=translate, inputs=input_text, outputs=output_text)
demo.launch()
if __name__ == "__main__":
main()