import gradio as gr
from transformers import pipeline

# Initialize the translation pipelines
translator_en_fi = pipeline("translation_en_to_fi", model="Helsinki-NLP/opus-mt-en-fi")
translator_fi_en = pipeline("translation_fi_to_en", model="Helsinki-NLP/opus-mt-fi-en")

def translate(text, direction):
    text = text.strip()
    if not text:
        return "Please enter some text for translation."
    if len(text) > 2000:
        return "Input text too long. Please shorten it."
    
    if direction == 'en-fi':
        result = translator_en_fi(text)[0]['translation_text']
    else:
        result = translator_fi_en(text)[0]['translation_text']
    return result

examples = [
    ["Hello, how are you?", "en-fi"],
    ["Mitä kuuluu?", "fi-en"]
]

iface = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(lines=3, placeholder="Enter text here..."),
        gr.Radio(choices=["en-fi", "fi-en"], label="Translation Direction", value="en-fi")
    ],
    outputs=gr.Textbox(label="Translated Text"),
    title="English-Finnish Translation App",
    description="This application uses Helsinki-NLP translation models to translate text between English and Finnish.",
    examples=examples
)

if __name__ == "__main__":
    iface.launch()