Spaces:
Runtime error
Runtime error
| from random import choices | |
| import gradio as gr | |
| import torch | |
| 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'] | |
| translate_interface = gr.Interface(fn=translate, | |
| inputs=outputs="text", | |
| title=title, | |
| description=description) | |
| 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() |