from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import gradio as gr

# Load fine-tuned BanglaT5 models for different tasks
translation_model_en_bn = AutoModelForSeq2SeqLM.from_pretrained("csebuetnlp/banglat5_nmt_en_bn")
translation_tokenizer_en_bn = AutoTokenizer.from_pretrained("csebuetnlp/banglat5_nmt_en_bn")

translation_model_bn_en = AutoModelForSeq2SeqLM.from_pretrained("csebuetnlp/banglat5_nmt_bn_en")
translation_tokenizer_bn_en = AutoTokenizer.from_pretrained("csebuetnlp/banglat5_nmt_bn_en")

summarization_model = AutoModelForSeq2SeqLM.from_pretrained("skl25/banglat5_xlsum_fine-tuned")
summarization_tokenizer = AutoTokenizer.from_pretrained("skl25/banglat5_xlsum_fine-tuned")

paraphrase_model = AutoModelForSeq2SeqLM.from_pretrained("csebuetnlp/banglat5_banglaparaphrase")
paraphrase_tokenizer = AutoTokenizer.from_pretrained("csebuetnlp/banglat5_banglaparaphrase")

# Define task functions
def translate_text_en_bn(input_text):
    inputs = translation_tokenizer_en_bn(input_text, return_tensors="pt")
    outputs = translation_model_en_bn.generate(**inputs)
    return translation_tokenizer_en_bn.decode(outputs[0], skip_special_tokens=True)

def translate_text_bn_en(input_text):
    inputs = translation_tokenizer_bn_en(input_text, return_tensors="pt")
    outputs = translation_model_bn_en.generate(**inputs)
    return translation_tokenizer_bn_en.decode(outputs[0], skip_special_tokens=True)

def summarize_text(input_text):
    inputs = summarization_tokenizer(input_text, return_tensors="pt")
    outputs = summarization_model.generate(**inputs)
    return summarization_tokenizer.decode(outputs[0], skip_special_tokens=True)

def paraphrase_text(input_text):
    inputs = paraphrase_tokenizer(input_text, return_tensors="pt")
    outputs = paraphrase_model.generate(**inputs)
    return paraphrase_tokenizer.decode(outputs[0], skip_special_tokens=True)

# Process input based on task
def process_text(text, task):
    task_funcs = {
        "Translate English to Bengali": translate_text_en_bn,
        "Translate Bengali to English": translate_text_bn_en,
        "Summarize": summarize_text,
        "Paraphrase": paraphrase_text
    }
    return task_funcs.get(task, lambda x: "Invalid Task")(text)

# Task-specific examples
examples_en_bn_translation = [
    ["The sky is blue, and the weather is nice."],
    ["Artificial intelligence is shaping the future."],
    ["Bangladesh is known for its rich culture and heritage."]
]

examples_bn_en_translation = [
    ["বাংলাদেশ দক্ষিণ এশিয়ার একটি সার্বভৌম রাষ্ট্র।"],
    ["ঢাকা বাংলাদেশের রাজধানী।"],
    ["রবীন্দ্রনাথ ঠাকুরের গান বাংলা সংস্কৃতির একটি অবিচ্ছেদ্য অংশ।"]
]

examples_summarization = [
    ["The Department of Computer Science and Engineering, established in 1982, was the first of its kind in Bangladesh. "
     "Attracting top students from all over the country, it offers both undergraduate and postgraduate degrees."],
    ["Climate change is one of the biggest challenges we face today. With rising temperatures and unpredictable weather, "
     "the world needs to come together to find sustainable solutions."],
    ["Technology has advanced rapidly over the past decade, with innovations in fields like AI, robotics, and quantum computing."]
]

examples_paraphrasing = [
    ["The cat is sitting on the mat."],
    ["He was very happy to receive the award."],
    ["The weather today is sunny and warm."]
]

# Define the Gradio interface with enhanced visuals
iface = gr.Interface(
    fn=process_text,
    inputs=[
        "text", 
        gr.Dropdown(
            ["Translate English to Bengali", "Translate Bengali to English", "Summarize", "Paraphrase"],
            label="Choose Task",
            elem_id="dropdown-task"
        )
    ],
    outputs="text",
    title="BanglaT5 Model Hub",
    description="A multi-functional tool for translation, summarization, and paraphrasing using BanglaT5 models.",
    theme="finlaydog/seafoam",
    examples=[
        *examples_en_bn_translation,  # Adding 3 examples for English to Bengali translation
        *examples_bn_en_translation,  # Adding 3 examples for Bengali to English translation
        *examples_summarization,      # Adding 3 examples for Summarization
        *examples_paraphrasing        # Adding 3 examples for Paraphrasing
    ],
    allow_flagging="auto",
)

# Launch the Gradio app
iface.launch(inline=False)