import os
import gradio as gr
from transformers import pipeline

print("done 1")

# Load models
try:
    base_model = pipeline("translation", model="Helsinki-NLP/opus-mt-tc-base-en-sh")
    print("Base model loaded successfully.")
except Exception as e:
    print(f"Error loading base model: {e}")

try:
    fine_tuned_model_1 = pipeline("translation", model="perkan/shortS-opus-mt-tc-base-en-sr")
    print("Fine-tuned model S loaded successfully.")
except Exception as e:
    print(f"Error loading fine-tuned model S: {e}")

try:
    fine_tuned_model_2 = pipeline("translation", model="perkan/shortM-opus-mt-tc-base-en-sr")
    print("Fine-tuned model M loaded successfully.")
except Exception as e:
    print(f"Error loading fine-tuned model M: {e}")

try:
    fine_tuned_model_3 = pipeline("translation", model="perkan/shortL-opus-mt-tc-base-en-sr")
    print("Fine-tuned model L loaded successfully.")
except Exception as e:
    print(f"Error loading fine-tuned model L: {e}")

# Define translation functions
def translate_base(text):
    try:
        return base_model(text)[0]['translation_text']
    except Exception as e:
        return f"Error during translation: {e}"

def translate_fine_tuned(text, model):
    try:
        if model == 'S model':
            return fine_tuned_model_1(text)[0]['translation_text']
        elif model == 'M model':
            return fine_tuned_model_2(text)[0]['translation_text']
        elif model == 'L model':
            return fine_tuned_model_3(text)[0]['translation_text']
        else:
            return "Invalid model selected"
    except Exception as e:
        return f"Error during translation: {e}"

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Translation Models\nTranslate text using base and fine-tuned models.")
    
    with gr.Row():
        text_input = gr.Textbox(placeholder="Enter text to translate", label="Input", placeholder="If you have behaved badly, repent, make what amends you can and address yourself to the task of behaving better next time. On no account brood over your wrongdoing. Rolling in the muck is not the best way of getting clean.")
        model_select = gr.Dropdown(choices=["S model", "M model", "L model"], label="Select Fine-tuned Model")
        translate_btn = gr.Button("Translate")
    
    with gr.Row():
        base_output = gr.Textbox(label="Base Model Translation")
        fine_tuned_output = gr.Textbox(label="Fine-tuned Model Translation")
    
    translate_btn.click(translate_text, inputs=[text_input, model_select], outputs=[base_output, fine_tuned_output])

port = int(os.getenv("GRADIO_SERVER_PORT", "7861"))
demo.launch(server_port=port)