File size: 3,116 Bytes
56f497c 70ebea3 56f497c e13fba3 70ebea3 56f497c e402119 e13fba3 56f497c 2169b22 56f497c e402119 56f497c 2169b22 56f497c 5a309fc 2169b22 28c6232 56f497c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
import gradio as gr
import spaces
from transformers import T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM
langs = {"English": "en", "Romanian": "ro", "German": "de", "French": "fr", "Spanish": "es", "Italian": "it"}
sl = langs[sselected_language]
tl = langs[tselected_language]
options = ["German", "Romanian", "English", "French", "Spanish", "Italian"]
models = ["Helsinki-NLP", "t5-base", "t5-small", "t5-large"]
@spaces.GPU
def translate_text(input_text, sselected_language, tselected_language, model_name):
if model_name == "Helsinki-NLP":
try:
model_name_full = f"Helsinki-NLP/opus-mt-{sl}-{tl}"
tokenizer = AutoTokenizer.from_pretrained(model_name_full)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name_full)
except EnvironmentError:
try :
model_name_full = f"Helsinki-NLP/opus-tatoeba-{sl}-{tl}"
tokenizer = AutoTokenizer.from_pretrained(model_name_full)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name_full)
except EnvironmentError as error:
return f"Error finding required model! Try other available language combination. Error: {error}"
else:
tokenizer = T5Tokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name)
if model_name.startswith("Helsinki-NLP"):
prompt = input_text
else:
prompt = f"translate {sselected_language} to {tselected_language}: {input_text}"
input_ids = tokenizer.encode(prompt, return_tensors="pt")
output_ids = model.generate(input_ids)
translated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
return translated_text
# Define a 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 gr.Row():
input_text = gr.Textbox(label="Enter text to translate:", placeholder="Type your text here...")
with gr.Row():
sselected_language = gr.Dropdown(choices=options, value="German", label="Source language")
tselected_language = gr.Dropdown(choices=options, value="Romanian", label="Target language")
swap_button = gr.Button("Swap Languages")
swap_button.click(fn=swap_languages, inputs=[sselected_language, tselected_language], outputs=[sselected_language, tselected_language])
model_name = gr.Dropdown(choices=models, value="Helsinki-NLP", label="Select a model")
translate_button = gr.Button("Translate")
translated_text = gr.Textbox(label="Translated text:", interactive=False)
translate_button.click(
translate_text,
inputs=[input_text, sselected_language, tselected_language, model_name],
outputs=translated_text
)
return interface
# Launch the Gradio interface
interface = create_interface()
interface.launch() |