LilithX6X's picture
Update app.py
313c8fb verified
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
def translate(text, language):
# Load the translation model
model_name = "trobinmv/t5_translate_en_ru_zh_small_1024"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
# Add language prefix to the input text
if language == "English":
lang_code = "en"
elif language == "Russian":
lang_code = "ru"
elif language == "Chinese":
lang_code = "zh"
input_text = f"translate to {lang_code}: {text}"
# Generate translation
try:
inputs = tokenizer(input_text, return_tensors="np", padding=True)
outputs = model.generate(**inputs, max_length=100)
translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
return translated_text
except Exception as e:
return f"Error: {e}"
def create_translation_ui():
with gr.Blocks() as demo:
gr.Markdown("# Translation Tool")
with gr.Row():
with gr.Column():
input_label = gr.Label("Enter word or phrase that you want translated")
input_textbox = gr.Textbox(label="", placeholder="Type your text here...")
language_label = gr.Label("Choose the language to translate into")
language_dropdown = gr.Dropdown(
choices=["English", "Russian", "Chinese"],
label="",
value="English"
)
with gr.Column():
output_textbox = gr.Textbox(
label="Translation",
readonly=True
)
translate_button = gr.Button("Translate")
translate_button.click(
fn=translate,
inputs=[input_textbox, language_dropdown],
outputs=output_textbox,
)
gr.Style(
button_styles="primary",
inputademoinput=True,
outputreadonly=True
)
return demo
if __name__ == "__main__":
create_translation_ui().launch()
#import gradio as gr
#gr.load("models/utrobinmv/t5_translate_en_ru_zh_small_1024").launch()
#translate to en: 毒品 强奸