import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load English → Irish
tokenizer_en_ga = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ga")
model_en_ga = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-ga")

# Load Irish → English
tokenizer_ga_en = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ga-en")
model_ga_en = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-ga-en")

def translate(text, direction):
    if direction == "English → Irish":
        inputs = tokenizer_en_ga(text, return_tensors="pt")
        outputs = model_en_ga.generate(**inputs)
        return tokenizer_en_ga.decode(outputs[0], skip_special_tokens=True)
    else:
        inputs = tokenizer_ga_en(text, return_tensors="pt")
        outputs = model_ga_en.generate(**inputs)
        return tokenizer_ga_en.decode(outputs[0], skip_special_tokens=True)

interface = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(label="Enter text"),
        gr.Radio(["English → Irish", "Irish → English"], label="Select direction")
    ],
    outputs="text",
    title="🇮🇪 GaelLingo",
    description="AI-powered English ↔ Irish Translator"
)

interface.launch()