Spaces:
Running
Running
import gradio as gr | |
from transformers import pipeline | |
# Load models | |
base_model = pipe = pipeline("translation", model="Helsinki-NLP/opus-mt-tc-base-en-sh") | |
fine_tuned_model_1 = pipeline("translation", model="perkan/shortS-opus-mt-tc-base-en-sr") | |
fine_tuned_model_2 = pipeline("translation", model="perkan/shortM-opus-mt-tc-base-en-sr") | |
fine_tuned_model_3 = pipeline("translation", model="perkan/shortL-opus-mt-tc-base-en-sr") | |
# Define translation functions | |
def translate_base(text): | |
return base_model(text)[0]['translation_text'] | |
def translate_fine_tuned(text, model): | |
if model == 'model1': | |
return fine_tuned_model_1(text)[0]['translation_text'] | |
elif model == 'model2': | |
return fine_tuned_model_2(text)[0]['translation_text'] | |
elif model == 'model3': | |
return fine_tuned_model_3(text)[0]['translation_text'] | |
else: | |
return "Invalid model selected" | |
# Create Gradio interface | |
with gr.Blocks() as demo: | |
gr.Markdown("### Translation Models") | |
with gr.Row(): | |
with gr.Column(): | |
gr.Markdown("#### Base Model") | |
base_input = gr.Textbox(placeholder="Enter text to translate", label="Input") | |
base_output = gr.Textbox(label="Translation") | |
base_translate_btn = gr.Button("Translate") | |
base_translate_btn.click(translate_base, inputs=base_input, outputs=base_output) | |
with gr.Column(): | |
gr.Markdown("#### Fine-tuned Models") | |
fine_tuned_input = gr.Textbox(placeholder="Enter text to translate", label="Input") | |
model_select = gr.Dropdown(choices=["model1", "model2", "model3"], label="Select Model") | |
fine_tuned_output = gr.Textbox(label="Translation") | |
fine_tuned_translate_btn = gr.Button("Translate") | |
fine_tuned_translate_btn.click(translate_fine_tuned, inputs=[fine_tuned_input, model_select], outputs=fine_tuned_output) | |
demo.launch() |