app.py

import spaces
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from flores import code_mapping
import platform
import torch
import nltk

nltk.download("punkt_tab")

REMOVED_TARGET_LANGUAGES = {"Ligurian", "Lombard", "Sicilian"}


device = "cpu" if platform.system() == "Darwin" else "cuda"
MODEL_NAME = "facebook/nllb-200-3.3B"

code_mapping = dict(sorted(code_mapping.items(), key=lambda item: item[1]))
flores_codes = list(code_mapping.keys())
target_languages = [language for language in flores_codes if not language in REMOVED_TARGET_LANGUAGES]

def load_model():
    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to(device)
    print(f"Model loaded in {device}")
    return model


model = load_model()


def load_tokenizer(src_lang, tgt_lang):
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_NAME, src_lang=code_mapping[src_lang], tgt_lang=code_mapping[tgt_lang]
    )
    return tokenizer


@spaces.GPU
def translate(text: str, src_lang: str, tgt_lang: str):
    tokenizer = load_tokenizer(src_lang, tgt_lang)

    paragraphs = text.split("\n")
    translated_paragraphs = []

    for paragraph in paragraphs:
        sentences = nltk.sent_tokenize(paragraph)
        translated_sentences = []

        for sentence in sentences:
            input_tokens = (
                tokenizer(sentence, return_tensors="pt")
                .input_ids[0]
                .cpu()
                .numpy()
                .tolist()
            )
            translated_chunk = model.generate(
                input_ids=torch.tensor([input_tokens]).to(device),
                forced_bos_token_id=tokenizer.convert_tokens_to_ids(code_mapping[tgt_lang]),
                max_length=len(input_tokens) + 50,
                num_return_sequences=1,
            )
            translated_chunk = tokenizer.decode(
                translated_chunk[0], skip_special_tokens=True
            )
            translated_sentences.append(translated_chunk)

        translated_paragraph = " ".join(translated_sentences)
        translated_paragraphs.append(translated_paragraph)

    return "\n".join(translated_paragraphs)



description = """
# UNESCO Language Translator, powered by Meta and Hugging Face

UNESCO, Meta, and Hugging Face have come together to create an accessible, high-quality translation experience in 200 languages. 

This is made possible through an open approach to AI innovation using Meta's open-sourced No Language Left Behind (NLLB) AI model, hosted on Hugging Face Spaces. """

disclaimer = """
## Disclaimer

This translation interface, developed as part of UNESCO's work on Multilingualism and supported by Meta's No Language Left Behind AI model and Hugging Face, is designed to assist with language translation using open-source AI technologies. However, translations generated by the tool may not be accurate or perfect. While we strive to provide accurate translations, the tool may produce inaccuracies due to the complexity and nuances of different languages.

- The tool may not fully capture the context, cultural nuances, idiomatic expressions, or specific terminologies.
- Manual review and adjustment are recommended for important translations.
- The translations are provided "as is" without any warranties of any kind, either expressed or implied.
- Users should not rely solely on the tool for critical or sensitive translations and are responsible for verifying the accuracy and appropriateness of the translations for their specific needs.
- We recommend consulting with professional translators for official, legal, medical, or other critical translations.
- We shall not be liable for any direct, indirect, incidental, special, or consequential damages arising out of or in connection with the use or inability to use the translation tool, including but not limited to errors or omissions in translations.

By using this translation tool, you agree to these terms and acknowledge that the use of the tool is at your own risk.

For any feedback or support, please contact UNESCO World Atlas of Languages Team: [email protected].
"""

with gr.Blocks() as demo:
    gr.Markdown(description)
    with gr.Row():
        src_lang = gr.Dropdown(label="Source Language", choices=flores_codes)
        target_lang = gr.Dropdown(label="Target Language", choices=target_languages)
    with gr.Row():
        input_text = gr.Textbox(label="Input Text", lines=6)
    with gr.Row():
        btn = gr.Button("Translate text")
    with gr.Row():
        output = gr.Textbox(label="Output Text", lines=6)
    btn.click(
        translate,
        inputs=[input_text, src_lang, target_lang],
        outputs=output,
    )
    with gr.Row():
        gr.Markdown(disclaimer)
demo.launch()