# -*- coding: utf-8 -*-
"""app.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1Ir0MiuPLQvCCHbSfM1tm0V_srviSEOi2
"""

import gradio as gr
import torch
from transformers import MarianMTModel, MarianTokenizer

class IdiomTranslator:
    def __init__(self):
        model_id = "Helsinki-NLP/opus-mt-hi-en"
        self.tokenizer = MarianTokenizer.from_pretrained(model_id)
        self.model     = MarianMTModel.from_pretrained(model_id)
        if torch.cuda.is_available(): self.model.cuda()

    def translate(self, text):
        if not text.strip(): return "Enter some Hindi text."
        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, padding=True)
        if torch.cuda.is_available(): inputs = {k: v.cuda() for k,v in inputs.items()}
        outs = self.model.generate(**inputs, max_length=128, num_beams=4)
        return self.tokenizer.decode(outs[0], skip_special_tokens=True)

translator = IdiomTranslator()

demo = gr.Interface(
    fn=translator.translate,
    inputs=gr.Textbox(lines=3, label="Hindi Text"),
    outputs=gr.Textbox(lines=3, label="English Translation"),
    title="Idiom-Aware Hindi→English Translator",
    description="Fine-tuned for cultural idioms.",
    examples=[["रस्सी जल गयी, बल नहीं गया"],["दाल में कुछ काला है"]]
)

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