import gradio as gr
from transformers import MarianMTModel, MarianTokenizer

# Load English→Spanish
en_es_model = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-en-es")
en_es_tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-es")

# Load Spanish→English
es_en_model = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-es-en")
es_en_tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-es-en")

def polish(sentence: str) -> str:
    # Step 1: English → Spanish
    inputs = en_es_tokenizer(sentence, return_tensors="pt", padding=True)
    translated = en_es_model.generate(**inputs, max_length=128)
    spanish = en_es_tokenizer.decode(translated[0], skip_special_tokens=True)

    # Step 2: Spanish → English
    inputs2 = es_en_tokenizer(spanish, return_tensors="pt", padding=True)
    back_translated = es_en_model.generate(**inputs2, max_length=128)
    english = es_en_tokenizer.decode(back_translated[0], skip_special_tokens=True)

    return english

demo = gr.Interface(
    fn=polish,
    inputs=gr.Textbox(lines=2, placeholder="Enter a sentence in English..."),
    outputs=gr.Textbox(label="Corrected English"),
    title="Round-trip Grammar Polisher",
    description="Uses Helsinki-NLP MarianMT models (en→es→en) to smooth and correct English sentences."
)

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