Spaces:
Sleeping
Sleeping
File size: 1,470 Bytes
acb75f7 |
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 |
import gradio as gr
from transformers import MarianMTModel, MarianTokenizer
# Define a function that loads a model and tokenizer based on the chosen language
def load_model(lang_pair):
if lang_pair == "English to French":
model_name = 'Helsinki-NLP/opus-mt-en-fr'
elif lang_pair == "Kinyarwanda to English":
model_name = 'Helsinki-NLP/opus-mt-rw-en'
else:
raise ValueError("Unsupported language pair")
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name)
return model, tokenizer
# Function to translate text based on selected language
def translate(lang_pair, text):
model, tokenizer = load_model(lang_pair)
# Tokenize the text using the __call__ method
model_inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
# Perform the translation
gen = model.generate(**model_inputs)
# Decode the generated tokens to string
translation = tokenizer.batch_decode(gen, skip_special_tokens=True)
return translation[0]
# Create a Gradio interface with a dropdown menu for language selection
iface = gr.Interface(
fn=translate,
inputs=[
gr.Dropdown(choices=["English to French", "Kinyarwanda to English"], label="Select Language"),
gr.Textbox(lines=2, placeholder="Enter Message...")
],
outputs=gr.Textbox()
)
# Launch the interface
iface.launch(debug=True,inline=False)
|