import gradio as gr
from transformers import pipeline
print("done 1")

# Load models
base_model = pipe = pipeline("translation", model="Helsinki-NLP/opus-mt-tc-base-en-sh")
print("done 2")
fine_tuned_model_1 = pipeline("translation", model="perkan/shortS-opus-mt-tc-base-en-sr")
print("done 3")
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()