Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import pipeline | |
| # Using the latest version of Gradio and Transformers | |
| # We want to expand the interface to include a reverse translation | |
| # We want to use the Helsinki-NLP/opus-mt-tc-big-he-en model for the reverse translation | |
| # A dropdown menu for selecting the model | |
| model_names = ["Helsinki-NLP/opus-mt-en-he", "Helsinki-NLP/opus-mt-tc-big-he-en"] | |
| model_name = gr.inputs.Dropdown(model_names, label="Model") | |
| # Name the dropdown options | |
| model_name.choices = ["English to Hebrew", "Hebrew to English"] | |
| # An output text box displaying the translated text and reverse translated text | |
| translation = gr.outputs.Textbox(label="Translation") | |
| reverse_translation = gr.outputs.Textbox(label="Reverse Translation") | |
| # A function for translating text | |
| def translate(model_name, text): | |
| # Create a pipeline for translating from English to Hebrew | |
| pipe = pipeline("translation", model=model_name) | |
| # Return the translation | |
| return pipe(text)[0]["translation_text"] | |
| # Create an interface for translating text | |
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM | |
| import torch | |
| tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-he") | |
| model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-he") | |
| def translate(model_name, text): | |
| # Create a pipeline for translating from English to Hebrew | |
| #Console out the model name | |
| print(model_name) | |
| if model_name == "English to Hebrew": | |
| forward_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-he") | |
| forward_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-he") | |
| reverse_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-tc-big-he-en") | |
| reverse_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-tc-big-he-en") | |
| elif model_name == "Hebrew to English": | |
| forward_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-tc-big-he-en") | |
| forward_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-tc-big-he-en") | |
| reverse_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-he") | |
| reverse_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-he") | |
| else: | |
| raise ValueError("Invalid model name") | |
| # Forward translation | |
| forward_input_ids = forward_tokenizer.encode(text, return_tensors="pt") | |
| forward_outputs = forward_model.generate(forward_input_ids) | |
| forward_translation = forward_tokenizer.decode(forward_outputs[0], skip_special_tokens=True) | |
| # Reverse translation | |
| reverse_input_ids = reverse_tokenizer.encode(forward_translation, return_tensors="pt") | |
| reverse_outputs = reverse_model.generate(reverse_input_ids) | |
| reverse_translation = reverse_tokenizer.decode(reverse_outputs[0], skip_special_tokens=True) | |
| return forward_translation, reverse_translation | |
| iface = gr.Interface(fn=translate, inputs=[model_name, "text"], outputs=[translation, reverse_translation]) | |
| # Launch the interface | |
| iface.launch(share=False) | |