Spaces:
Running
Running
| 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", | |
| "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", | |
| "t5-small", "t5-base", "t5-large", | |
| "google/flan-t5-small", "google/flan-t5-base", "google/flan-t5-large", "google/flan-t5-xl", | |
| "Argos", "Google", | |
| "HuggingFaceTB/SmolLM3-3B", | |
| "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 | |
| class Translators: | |
| def __init__(self, model_name: str, sl: str, tl: str, input_text: str): | |
| self.model_name = model_name | |
| self.sl, self.tl = sl, tl | |
| self.input_text = input_text | |
| def google(self): | |
| url = os.environ['GCLIENT'] + f'sl={self.sl}&tl={self.tl}&q={self.input_text}' | |
| response = requests.get(url) | |
| return response.json()[0][0][0] | |
| def smollm(self): | |
| tokenizer = AutoTokenizer.from_pretrained(self.model_name) | |
| model = AutoModelForCausalLM.from_pretrained(self.model_name) | |
| prompt = f"""Translate the following {self.sl} text to {self.tl}, generating only the translated text and maintaining the original meaning and tone: | |
| {self.input_text} | |
| Translation:""" | |
| inputs = tokenizer(prompt, return_tensors="pt") | |
| outputs = model.generate( | |
| inputs.input_ids, | |
| max_length=len(inputs.input_ids[0]) + 150, | |
| temperature=0.3, | |
| do_sample=True | |
| ) | |
| response = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| return response.split("Translation:")[-1].strip() | |
| 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: # Standard bilingual model | |
| 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: # Tatoeba models | |
| 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: | |
| try: # Last resort: multi to multi | |
| model_name = "Helsinki-NLP/opus-mt-tc-bible-big-mul-mul" | |
| pipe = pipeline("translation", model=model_name) | |
| tl = 'deu' # Hard coded for now for testing | |
| translation = pipe(f'>>{tl}<< {input_text}') | |
| return translation[0]['translation_text'], f'Translated from {sl} to {tl} with {model_name}.' | |
| except Exception as error: | |
| return f"Error translating with 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) | |
| tokenized_input = pipe.tokenizer(input_text, return_tensors="pt") | |
| num_input_tokens = len(tokenized_input["input_ids"][0]) | |
| max_new_tokens = round(num_input_tokens + 0.25 * num_input_tokens) | |
| outputs = pipe(prompt, max_new_tokens=max_new_tokens, do_sample=False) | |
| translated_text = outputs[0]["generated_text"] | |
| print(f"Input chars: {len(input_text)}", f"Input tokens: {num_input_tokens}", f"max_new_tokens: {max_new_tokens}", | |
| "Chars to tokens ratio:", round(len(input_text) / num_input_tokens, 2), f"Raw translation: {translated_text}") | |
| markers = ["<end_of_turn>", "<|im_end|>", "<|im_start|>assistant"] # , "\n" | |
| for marker in markers: | |
| if marker in translated_text: | |
| translated_text = translated_text.split(marker)[1].strip() | |
| translated_text = translated_text.replace('Answer:', '', 1).strip() if translated_text.startswith('Answer:') else translated_text | |
| translated_text = translated_text.split("Translated text:")[0].strip() if "Translated text:" in translated_text else translated_text | |
| split_translated_text = translated_text.split('\n', translated_text.count('\n')) | |
| translated_text = '\n'.join(split_translated_text[:input_text.count('\n')+1]) | |
| 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] | |
| 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 = Translators(model_name, sl, tl, input_text).google() | |
| 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) | |
| elif model_name == "HuggingFaceTB/SmolLM3-3B": | |
| translated_text = Translators(model_name, sselected_language, tselected_language, input_text).smollm() | |
| except Exception as error: | |
| translated_text = error | |
| finally: | |
| print(input_text, translated_text, message_text) | |
| 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) |