app.py

import os
import torch
import gradio as gr
import time
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline

model = 'facebook/nllb-200-distilled-600M'
tokenizer = model

flores_codes = {}
flores_codes["Asturianu"] = "ast_Latn"
flores_codes["Castellano"] = "spa_Latn"
flores_codes["Català"] = "cat_Latn"
flores_codes["English"] = "eng_Latn"
flores_codes["Euskera"] = "eus_Latn"
flores_codes["Galego"] = "glg_Latn"

def translation(source, target, text):

    #start_time = time.time()
    source = flores_codes[source]
    target = flores_codes[target]

    translator = pipeline('translation', model=model, tokenizer=tokenizer, 

    src_lang=source, tgt_lang=target)
    output = translator(text, max_length=400)

    #end_time = time.time()

    output = output[0]['translation_text']
    #result = {'inference_time': end_time - start_time,
    #          'source': source,
    #          'target': target,
    #          'result': output}
    #return result
    return output;

if __name__ == '__main__':
    print('\tIniciando...')

   
    # define gradio demo
    lang_codes = list(flores_codes.keys())
    inputs = [gr.Dropdown(lang_codes, value='Castellano', label='Idioma original'),
              gr.Dropdown(lang_codes, value='Asturianu', label='Traducir al...'),
              gr.Textbox(label="Texto a traducir"),
              ]

    outputs = [gr.Textbox(label="Texto traducido"),]

    title = "Traductor Multilingüe"

    description = """Este traductor utiliza el siguiente modelo de lenguaje de Meta:
    
https://github.com/facebookresearch/fairseq/tree/nllb

Adaptado de:

https://huggingface.co/spaces/Azwaw/Text_Translation_Multi-languages"""


    gr.Interface(translation,
                 inputs,
                 outputs,
                 title=title,
                 description=description,
                 ).launch()