File size: 1,401 Bytes
d204ef0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
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()