NLLB-Translator / app.py
mrm8488's picture
Fix bug
44b47e7
raw history blame
No virus
1.25 kB
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
from ui import title, description
from langs import LANGS
TASK = "translation"
CKPT = "facebook/nllb-200-distilled-600M"
model = AutoModelForSeq2SeqLM.from_pretrained(CKPT)
tokenizer = AutoTokenizer.from_pretrained(CKPT)
def translate(text, src_lang, tgt_lang, max_length=400):
translation_pipeline = pipeline(TASK,
model=model,
tokenizer=tokenizer,
src_lang=src_lang,
tgt_lang=tgt_lang,
max_length=max_length)
result = translation_pipeline(text)
return result[0]['translation_text']
gr.Interface(
translate,
[
gr.inputs.Textbox(label="Text"),
gr.inputs.Dropdown(label="Source Language", choices=LANGS),
gr.inputs.Dropdown(label="Target Language", choices=LANGS),
gr.inputs.slider(label="Max Length", min=8,
max=512, step=8, default=400)
],
["text"],
# examples=examples,
# article=article,
cache_examples=False,
title=title,
description=description
).launch()