Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer | |
import torch | |
# Load the model and tokenizer | |
model_name = "himanshubeniwal/opus-mt-en-ro-finetuned-ro-to-en-agreement" | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
model = AutoModelForSeq2SeqLM.from_pretrained(model_name) | |
def translate_text( | |
text, | |
max_length=200, | |
num_beams=4, | |
num_return_sequences=1, | |
temperature=1.0, | |
repetition_penalty=1.0 | |
): | |
""" | |
Translate Romanian text to English using the fine-tuned model | |
""" | |
try: | |
# Encode the input text | |
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True) | |
# Generate translation | |
outputs = model.generate( | |
inputs["input_ids"], | |
max_length=max_length, | |
num_beams=num_beams, | |
num_return_sequences=num_return_sequences, | |
temperature=temperature, | |
repetition_penalty=repetition_penalty, | |
early_stopping=True | |
) | |
# Decode and format the translations | |
translations = [] | |
for output in outputs: | |
translation = tokenizer.decode(output, skip_special_tokens=True) | |
translations.append(translation) | |
# Return single string if only one sequence, otherwise return list | |
if num_return_sequences == 1: | |
return translations[0] | |
return "\n\n".join(translations) | |
except Exception as e: | |
return f"Error in translation: {str(e)}" | |
# Create the Gradio interface | |
iface = gr.Interface( | |
fn=translate_text, | |
inputs=[ | |
gr.Textbox(lines=5, label="Enter Romanian text"), | |
gr.Slider(minimum=10, maximum=400, value=200, step=10, label="Max Length"), | |
gr.Slider(minimum=1, maximum=8, value=4, step=1, label="Number of Beams"), | |
gr.Slider(minimum=1, maximum=3, value=1, step=1, label="Number of Translations"), | |
gr.Slider(minimum=0.1, maximum=2.0, value=1.0, step=0.1, label="Temperature"), | |
gr.Slider(minimum=1.0, maximum=2.0, value=1.0, step=0.1, label="Repetition Penalty") | |
], | |
outputs=gr.Textbox(lines=5, label="English Translation"), | |
title="Romanian to English Translation Model", | |
description="Translate Romanian text to English using a fine-tuned OPUS-MT model. Examples include common phrases like: \n 1. 'Guvernul dumneavoastră are un acord cu Japonia.' (Your Government has an agreement with Japan.) \n 2,'Instanța a spus că au un acord.' (The court said they had an agreement.)", | |
examples=[ | |
["Guvernul dumneavoastră are un acord cu Japonia.", 200, 4, 1, 1.0, 1.0], | |
["Instanța a spus că au un acord.", 200, 4, 1, 1.0, 1.0], | |
] | |
) | |
# Launch the interface | |
if __name__ == "__main__": | |
iface.launch() |