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 with clear labels
examples = [
    ("Translate English to Bengali", [
        ["The sky is blue, and the weather is nice."],
        ["Artificial intelligence is shaping the future."],
        ["Bangladesh is known for its rich culture and heritage."]
    ]),
    ("Translate Bengali to English", [
        ["বাংলাদেশ দক্ষিণ এশিয়ার একটি সার্বভৌম রাষ্ট্র।"],
        ["ঢাকা বাংলাদেশের রাজধানী।"],
        ["রবীন্দ্রনাথ ঠাকুরের গান বাংলা সংস্কৃতির একটি অবিচ্ছেদ্য অংশ।"]
    ]),
    ("Summarize", [
        ["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."]
    ]),
    ("Paraphrase", [
        ["The cat is sitting on the mat."],
        ["He was very happy to receive the award."],
        ["The weather today is sunny and warm."]
    ])
]

# Enhanced visual layout and interface design
iface = gr.Interface(
    fn=process_text,
    inputs=[
        gr.Textbox(lines=5, label="Enter your text"),
        gr.Dropdown(
            ["Translate English to Bengali", "Translate Bengali to English", "Summarize", "Paraphrase"],
            label="Select Task",
            elem_id="dropdown-task",
            interactive=True
        )
    ],
    outputs=gr.Textbox(label="Output"),
    title="BanglaT5 Model Hub - Translation, Summarization & Paraphrasing",
    description="Explore the power of BanglaT5 with this easy-to-use interface for multiple tasks like translation, summarization, and paraphrasing.",
    theme="huggingface/space-shuttle",
    examples=[
        ("Translate English to Bengali", [
            ["The sky is blue, and the weather is nice."],
            ["Artificial intelligence is shaping the future."],
            ["Bangladesh is known for its rich culture and heritage."]
        ]),
        ("Translate Bengali to English", [
            ["বাংলাদেশ দক্ষিণ এশিয়ার একটি সার্বভৌম রাষ্ট্র।"],
            ["ঢাকা বাংলাদেশের রাজধানী।"],
            ["রবীন্দ্রনাথ ঠাকুরের গান বাংলা সংস্কৃতির একটি অবিচ্ছেদ্য অংশ।"]
        ]),
        ("Summarize", [
            ["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."]
        ]),
        ("Paraphrase", [
            ["The cat is sitting on the mat."],
            ["He was very happy to receive the award."],
            ["The weather today is sunny and warm."]
        ])
    ],
    allow_flagging="auto",
    layout="horizontal",  # Adding a horizontal layout for a cleaner design
    css="""
        #dropdown-task {
            color: #2C3E50;
            font-size: 18px;
        }
        .output-area {
            color: #2C3E50;
        }
        .examples {
            display: grid;
            grid-template-columns: 1fr 1fr;
            gap: 20px;
        }
        .output-text {
            background-color: #E8F6F3;
        }
    """
)

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