app.py

import gradio as gr
import spaces
import torch
from transformers import T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM, AutoModel, pipeline, logging
import languagecodes
import requests, os

logging.set_verbosity_error()
favourite_langs = {"German": "de", "Romanian": "ro", "English": "en", "-----": "-----"}
all_langs = languagecodes.iso_languages

# Language options as list, add favourite languages first
options = list(favourite_langs.keys())
options.extend(list(all_langs.keys()))
models = ["Helsinki-NLP",
          "t5-small", "t5-base", "t5-large",
          "google/flan-t5-small", "google/flan-t5-base", "google/flan-t5-large", "google/flan-t5-xl",
          "facebook/nllb-200-distilled-600M", "facebook/nllb-200-distilled-1.3B", "facebook/nllb-200-1.3B", "facebook/nllb-200-3.3B",
          "facebook/mbart-large-50-many-to-many-mmt", "facebook/mbart-large-50-one-to-many-mmt", "facebook/mbart-large-50-many-to-one-mmt",
          "facebook/m2m100_418M", "facebook/m2m100_1.2B",
          "bigscience/mt0-small", "bigscience/mt0-base", "bigscience/mt0-large", "bigscience/mt0-xl",
          "bigscience/bloomz-560m", "bigscience/bloomz-1b1", "bigscience/bloomz-1b7", "bigscience/bloomz-3b",
          "Argos", "Google",
          "utter-project/EuroLLM-1.7B", "utter-project/EuroLLM-1.7B-Instruct",
          "Unbabel/Tower-Plus-2B", "Unbabel/TowerInstruct-7B-v0.2", "Unbabel/TowerInstruct-Mistral-7B-v0.2",
          "openGPT-X/Teuken-7B-instruct-commercial-v0.4", "openGPT-X/Teuken-7B-instruct-v0.6"
          ]

def model_to_cuda(model):
    # Move the model to GPU if available
    if torch.cuda.is_available():
        model = model.to('cuda')
        print("CUDA is available! Using GPU.")
    else:
        print("CUDA not available! Using CPU.")
    return model

def download_argos_model(from_code, to_code):
    import argostranslate.package
    print('Downloading model', from_code, to_code) 
    # Download and install Argos Translate package
    argostranslate.package.update_package_index()
    available_packages = argostranslate.package.get_available_packages()
    package_to_install = next(
        filter(
            lambda x: x.from_code == from_code and x.to_code == to_code, available_packages
        )
    )
    argostranslate.package.install_from_path(package_to_install.download())

def argos(sl, tl, input_text):
    import argostranslate.translate, argostranslate.package       
    # Translate
    try:
        download_argos_model(sl, tl)
        translated_text = argostranslate.translate.translate(input_text, sl, tl)
    except StopIteration:
        # packages_info = ', '.join(f"{pkg.get_description()}->{str(pkg.links)} {str(pkg.source_languages)}" for pkg in argostranslate.package.get_available_packages())
        packages_info = ', '.join(f"{pkg.from_name} ({pkg.from_code}) -> {pkg.to_name} ({pkg.to_code})" for pkg in argostranslate.package.get_available_packages())
        translated_text = f"No Argos model for {sl} to {tl}. Try other model or languages combination from the available Argos models: {packages_info}."
    except Exception as error:
        translated_text = error
        print(error)
    return translated_text

def google(sl, tl, input_text):  
    url = os.environ['GCLIENT'] + f'sl={sl}&tl={tl}&q={input_text}'
    response = requests.get(url)
    return response.json()[0][0][0]

def mtom(model_name, sl, tl, input_text):
    from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
    model = M2M100ForConditionalGeneration.from_pretrained(model_name)
    tokenizer = M2M100Tokenizer.from_pretrained(model_name)
    tokenizer.src_lang = sl
    encoded = tokenizer(input_text, return_tensors="pt")
    generated_tokens = model.generate(**encoded, forced_bos_token_id=tokenizer.get_lang_id(tl))
    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]

def HelsinkiNLPAutoTokenizer(sl, tl, input_text):
    if model_name == "Helsinki-NLP":
        message_text = f'Translated from {sl} to {tl} with {model_name}.'
        try:
            model_name = f"Helsinki-NLP/opus-mt-{sl}-{tl}"
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            model = model_to_cuda(AutoModelForSeq2SeqLM.from_pretrained(model_name))
        except EnvironmentError:
            try:   
                model_name = f"Helsinki-NLP/opus-tatoeba-{sl}-{tl}"
                tokenizer = AutoTokenizer.from_pretrained(model_name)
                model = model_to_cuda(AutoModelForSeq2SeqLM.from_pretrained(model_name))
                input_ids = tokenizer.encode(prompt, return_tensors="pt")
                output_ids = model.generate(input_ids, max_length=512)
                translated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
                return translated_text, message_text
            except EnvironmentError as error:
                return f"Error finding model: {model_name}! Try other available language combination.", error
                
def HelsinkiNLP(sl, tl, input_text):
    try:
        model_name = f"Helsinki-NLP/opus-mt-{sl}-{tl}"
        pipe = pipeline("translation", model=model_name, device=-1)
        # translation = pipe(input_text)
        # return translation[0]['translation_text'], f'Translated from {sl} to {tl} with {model_name}.'
    except EnvironmentError:
        try:   
            model_name = f"Helsinki-NLP/opus-tatoeba-{sl}-{tl}"
            pipe = pipeline("translation", model=model_name, device=-1)
            translation = pipe(input_text)
            return translation[0]['translation_text'], f'Translated from {sl} to {tl} with {model_name}.'
        except EnvironmentError as error:
            return f"Error finding model: {model_name}! Try other available language combination.", error
    except KeyError as error:
        return f"Error: Translation direction {sl} to {tl} is not supported by Helsinki Translation Models", error     

def flan(model_name, sl, tl, input_text):
    tokenizer = T5Tokenizer.from_pretrained(model_name, legacy=False)
    model = T5ForConditionalGeneration.from_pretrained(model_name)
    input_text = f"translate {sl} to {tl}: {input_text}"
    input_ids = tokenizer(input_text, return_tensors="pt").input_ids
    outputs = model.generate(input_ids)
    return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

def tfive(model_name, sl, tl, input_text):
    tokenizer = T5Tokenizer.from_pretrained(model_name)
    model = T5ForConditionalGeneration.from_pretrained(model_name, device_map="auto")
    prompt = f"translate {sl} to {tl}: {input_text}"
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    output_ids = model.generate(input_ids, max_length=512)
    translated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return translated_text

def teuken(model_name, sl, tl, input_text):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
    )
    model = model.to(device).eval()
    tokenizer = AutoTokenizer.from_pretrained(
        model_name,
        use_fast=False,
        trust_remote_code=True,
    )
    translation_prompt = f"Translate the following text from {sl} into {tl}: {input_text}"
    messages = [{"role": "User", "content": translation_prompt}]
    prompt_ids = tokenizer.apply_chat_template(messages, chat_template="EN", tokenize=True, add_generation_prompt=False, return_tensors="pt")
    prediction = model.generate(
        prompt_ids.to(model.device),
        max_length=512,
        do_sample=True,
        top_k=50,
        top_p=0.95,
        temperature=0.7,
        num_return_sequences=1,
    )
    translation = tokenizer.decode(prediction[0].tolist())
    return translation

def bigscience(model_name, sl, tl, input_text):  
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    inputs = tokenizer.encode(f"Translate to {tl}: {input_text}.", return_tensors="pt")
    outputs = model.generate(inputs)
    translation = tokenizer.decode(outputs[0])
    translation = translation.replace('<pad> ', '').replace('</s>', '')
    return translation

def bloomz(model_name, sl, tl, input_text):  
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    inputs = tokenizer.encode(f"Translate from {sl} to {tl}: {input_text}. Translation:", return_tensors="pt")
    outputs = model.generate(inputs)
    translation = tokenizer.decode(outputs[0])
    translation = translation.replace('<pad> ', '').replace('</s>', '')
    return translation

def eurollm(model_name, sl, tl, input_text):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)  
    prompt = f"{sl}: {input_text} {tl}:"
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=512)
    output = tokenizer.decode(outputs[0], skip_special_tokens=True)
    result = output.rsplit(f'{tl}:')[-1].strip()
    return result

def eurollm_instruct(model_name, sl, tl, input_text):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    text = f'<|im_start|>system\n<|im_end|>\n<|im_start|>user\nTranslate the following {sl} source text to {tl}:\n{sl}: {input_text} \n{tl}: <|im_end|>\n<|im_start|>assistant\n'
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=512)
    output = tokenizer.decode(outputs[0], skip_special_tokens=True)
    if f'{tl}:' in output:
        output = output.rsplit(f'{tl}:')[-1].strip().replace('assistant\n', '')
    return output

def nllb(model_name, sl, tl, input_text):
    tokenizer = AutoTokenizer.from_pretrained(model_name, src_lang=sl)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name, device_map="auto")
    translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=sl, tgt_lang=tl)
    translated_text = translator(input_text, max_length=512)
    return translated_text[0]['translation_text']

def unbabel(model_name, sl, tl, input_text):
    pipe = pipeline("text-generation", model=model_name, torch_dtype=torch.bfloat16, device_map="auto")
    messages = [{"role": "user",
                 "content": f"Translate the following text from {sl} into {tl}.\n{sl}: {input_text}.\n{tl}:"}]
    prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
    outputs = pipe(prompt, max_new_tokens=256, do_sample=False)
    translated_text = outputs[0]["generated_text"]
    start_marker = "<end_of_turn>"
    if start_marker in translated_text:
        translated_text = translated_text.split(start_marker)[1].strip()
    translated_text = translated_text.replace('Answer:', '', 1).strip() if translated_text.startswith('Answer:') else translated_text 
    return translated_text

def mbart_many_to_many(model_name, sl, tl, input_text):
    from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
    model = MBartForConditionalGeneration.from_pretrained(model_name)
    tokenizer = MBart50TokenizerFast.from_pretrained(model_name)
    # translate source to target
    tokenizer.src_lang = languagecodes.mbart_large_languages[sl]
    encoded = tokenizer(input_text, return_tensors="pt")
    generated_tokens = model.generate(
        **encoded,
        forced_bos_token_id=tokenizer.lang_code_to_id[languagecodes.mbart_large_languages[tl]]
    )
    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]

def mbart_one_to_many(model_name, sl, tl, input_text):
    from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
    article_en = input_text
    model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-one-to-many-mmt")
    tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-one-to-many-mmt", src_lang="en_XX")
    model_inputs = tokenizer(article_en, return_tensors="pt")
    # translate from English
    langid = languagecodes.mbart_large_languages[tl]
    generated_tokens = model.generate(
        **model_inputs,
        forced_bos_token_id=tokenizer.lang_code_to_id[langid]
    )
    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]

def mbart_many_to_one(model_name, sl, tl, input_text):
    from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
    model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-one-mmt")
    tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-one-mmt")
    # translate to English
    tokenizer.src_lang = languagecodes.mbart_large_languages[sl]
    encoded = tokenizer(input_text, return_tensors="pt")
    generated_tokens = model.generate(**encoded)
    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]

@spaces.GPU
def translate_text(input_text: str, sselected_language: str, tselected_language: str, model_name: str) -> tuple[str, str]:
    """
    Translates the input text from the source language to the target language  using a specified model.

    Parameters:
        input_text (str): The source text to be translated
        sselected_language (str): The source language of the input text
        tselected_language (str): The target language in which the input text is translated
        model_name (str): The selected translation model name

    Returns:
        tuple: 
            translated_text(str): The input text translated to the selected target language
            message_text(str):  A descriptive message summarizing the translation process. Example: "Translated from English to German with Helsinki-NLP."
    
    Example:
        >>> translate_text("Hello world", "English", "German", "Helsinki-NLP")
        ("Hallo Welt", "Translated from English to German with Helsinki-NLP.")
    """
    
    sl = all_langs[sselected_language]
    tl = all_langs[tselected_language]
    message_text = f'Translated from {sselected_language} to {tselected_language} with {model_name}'
    print(message_text)
    try:    
        if model_name.startswith("Helsinki-NLP"):
            translated_text, message_text = HelsinkiNLP(sl, tl, input_text)
        
        elif model_name == 'Argos':
            translated_text = argos(sl, tl, input_text)
    
        elif model_name == 'Google':
            translated_text = google(sl, tl, input_text)
    
        elif "m2m" in model_name.lower():
            translated_text = mtom(model_name, sl, tl, input_text)
        
        elif model_name == "utter-project/EuroLLM-1.7B-Instruct":
            translated_text = eurollm_instruct(model_name, sselected_language, tselected_language, input_text)
        
        elif model_name == "utter-project/EuroLLM-1.7B":
            translated_text = eurollm(model_name, sselected_language, tselected_language, input_text)
    
        elif 'flan' in model_name.lower():
            translated_text = flan(model_name, sselected_language, tselected_language, input_text)
            
        elif 'teuken' in model_name.lower():
            translated_text = teuken(model_name, sselected_language, tselected_language, input_text)
    
        elif 'mt0' in model_name.lower():
            translated_text = bigscience(model_name, sselected_language, tselected_language, input_text)
    
        elif 'bloomz' in model_name.lower():
            translated_text = bloomz(model_name, sselected_language, tselected_language, input_text)
            
        elif 'nllb' in model_name.lower():
            nnlbsl, nnlbtl = languagecodes.nllb_language_codes[sselected_language], languagecodes.nllb_language_codes[tselected_language]
            translated_text = nllb(model_name, nnlbsl, nnlbtl, input_text)
        
        elif model_name == "facebook/mbart-large-50-many-to-many-mmt":
            translated_text = mbart_many_to_many(model_name, sselected_language, tselected_language, input_text)
    
        elif model_name == "facebook/mbart-large-50-one-to-many-mmt":
            translated_text = mbart_one_to_many(model_name, sselected_language, tselected_language, input_text)
    
        elif model_name == "facebook/mbart-large-50-many-to-one-mmt":
            translated_text = mbart_many_to_one(model_name, sselected_language, tselected_language, input_text)
        
        elif 'Unbabel' in model_name:   
            translated_text = unbabel(model_name, sselected_language, tselected_language, input_text)
        
        elif model_name.startswith('t5'):
            translated_text = tfive(model_name, sselected_language, tselected_language, input_text)
    
    except Exception as error:
        translated_text = error
    finally:
        return translated_text, message_text

# Function to swap dropdown values
def swap_languages(src_lang, tgt_lang):
    return tgt_lang, src_lang 

def create_interface():
    with gr.Blocks() as interface:
        gr.Markdown("### Machine Text Translation with Gradio API and MCP Server")

        with gr.Row():
            input_text = gr.Textbox(label="Enter text to translate:", placeholder="Type your text here, maximum 512 tokens")
        
        with gr.Row():
            sselected_language = gr.Dropdown(choices=options, value = options[0], label="Source language", interactive=True)
            tselected_language = gr.Dropdown(choices=options, value = options[1], label="Target language", interactive=True)
            swap_button = gr.Button("Swap Languages", size="md")
            swap_button.click(fn=swap_languages, inputs=[sselected_language, tselected_language], outputs=[sselected_language, tselected_language], api_name=False, show_api=False)

        model_name = gr.Dropdown(choices=models, label=f"Select a model. Default is {models[0]}.", value = models[0], interactive=True)
        translate_button = gr.Button("Translate")

        translated_text = gr.Textbox(label="Translated text:", placeholder="Display field for translation", interactive=False, show_copy_button=True)
        message_text = gr.Textbox(label="Messages:", placeholder="Display field for status and error messages", interactive=False,
                                  value=f'Default translation settings: from {sselected_language.value} to {tselected_language.value} with {model_name.value}.')
        allmodels = gr.HTML(label="Model links:", value=', '.join([f'<a href="https://huggingface.co/{model}">{model}</a>' for model in models]))

        translate_button.click(
            fn=translate_text, 
            inputs=[input_text, sselected_language, tselected_language, model_name], 
            outputs=[translated_text, message_text]
        )

    return interface

interface = create_interface()
if __name__ == "__main__":
    interface.launch(mcp_server=True)
    # interface.queue().launch(server_name="0.0.0.0", show_error=True, server_port=7860, mcp_server=True)