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)

# Task-specific examples
examples_translate_en_bn = [
    ["The Department of Computer Science and Engineering, the first department of its kind, was established in 1982 under the faculty of Electrical and Electronics Engineering. From the very initial days of its establishment, it has been able to attract the very best students of the country. Students securing topmost merit positions in the legendary admission test of BUET opt for studies in this department. At the very beginning, the department offered only M. Sc. Engg. and M Engg. degrees in Computer Science and Engineering. The Undergraduate program started in 1986. At first, 30 students were admitted each year for pursuing the B.Sc. Engineering degree."],
    ["Artificial intelligence is shaping the future."],
    ["Bangladesh is known for its rich culture and heritage."]
]

examples_translate_bn_en = [
    ["বাংলাদেশ দক্ষিণ এশিয়ার একটি সার্বভৌম রাষ্ট্র।"],
    ["ফুল মানুষের হৃদয়কে উজ্জীবিত করে। বাংলাদেশের বিভিন্ন অঞ্চলে নানা রকমের ফুল ফুটে থাকে, যেমন গোলাপ, জবা, এবং কৃষ্ণকলি। বিশেষত বসন্তে, ফুলের সৌন্দর্য যেন চারপাশে ছড়িয়ে পড়ে, এবং এর গন্ধ বাতাসে ভাসে। ফুলের উৎসব, যেমন 'বসন্ত উৎসব', এ দেশের সংস্কৃতির একটি গুরুত্বপূর্ণ অংশ।"],
    ["রবীন্দ্রনাথ ঠাকুরের গান বাংলা সংস্কৃতির একটি অবিচ্ছেদ্য অংশ।"]
]

examples_summarize = [
    ["বাংলাদেশ দক্ষিণ এশিয়ার একটি সার্বভৌম রাষ্ট্র। বাংলাদেশের সাংবিধানিক নাম গণপ্রজাতন্ত্রী বাংলাদেশ। ভৌগোলিকভাবে বাংলাদেশের পশ্চিমে ভারতের পশ্চিমবঙ্গ, উত্তরে পশ্চিমবঙ্গ, আসাম ও মেঘালয়, পূর্ব সীমান্তে আসাম, ত্রিপুরা ও মিজোরাম, দক্ষিণ-পূর্ব সীমান্তে মিয়ানমারের চিন ও রাখাইন রাজ্য এবং দক্ষিণ উপকূলের দিকে বঙ্গোপসাগর অবস্থিত। ভৌগোলিকভাবে পৃথিবীর বৃহত্তম ব-দ্বীপের সিংহভাগ অঞ্চল জুড়ে বাংলাদেশ ভূখণ্ড অবস্থিত। জনসংখ্যার বিচারে প্রায় ১৭ কোটিরও অধিক জনসংখ্যা নিয়ে বাংলাদেশ বিশ্বের ৮ম বৃহত্তম দেশ। নদীমাতৃক বাংলাদেশ ভূখণ্ডের উপর দিয়ে বয়ে গেছে ৫৭টি আন্তর্জাতিক নদী। বাংলাদেশের উত্তর-পূর্বে ও দক্ষিণ-পূর্বে টারশিয়ারি যুগের পাহাড় ছেয়ে আছে। বিশ্বের বৃহত্তম ম্যানগ্রোভ অরণ্য সুন্দরবন ও দীর্ঘতম প্রাকৃতিক সৈকত কক্সবাজার সমুদ্র সৈকত বাংলাদেশে অবস্থিত।"],
    ["বিশ্বের আরও একটি সেরা ক্লাব। স্প্যানিশ ফুটবলে বার্সেলোনারই প্রবল প্রতিদ্বন্দ্বী রিয়াল মাদ্রিদ। লিওনেল মেসির ক্লাব ইতিহাসে উয়েফা চ্যাম্পিয়নস লিগ জিতেছে মোট পাঁচবার। ইউরোপীয় কাপ উইনার্স কাপে তাদের সাফল্য আছে চারবার। পাশাপাশি তিনবার ক্লাব বিশ্বকাপ জেতা বার্সা লা লিগায় সাফল্য পেয়েছেন মোট ২৪ বার। ডিয়েগো ম্যারাডোনা, ইয়োহান ক্রুইফ, রিভালদোদের মতো খেলোয়াড়েরা এই ক্লাবের হয়ে অংশ হয়ে আছেন ইতিহাসের।"],
    ["ফুল মানুষের হৃদয়কে উজ্জীবিত করে। বাংলাদেশের বিভিন্ন অঞ্চলে নানা রকমের ফুল ফুটে থাকে, যেমন গোলাপ, জবা, এবং কৃষ্ণকলি। বিশেষত বসন্তে, ফুলের সৌন্দর্য যেন চারপাশে ছড়িয়ে পড়ে, এবং এর গন্ধ বাতাসে ভাসে। ফুলের উৎসব, যেমন 'বসন্ত উৎসব', এ দেশের সংস্কৃতির একটি গুরুত্বপূর্ণ অংশ।"]
]

examples_paraphrase = [
    ["তিনজনের কেউই বাবার পথ ধরে প্রযুক্তি দুনিয়ায় হাঁটেননি।"],
    ["প্রতিদিন নিয়মিত ব্যায়াম করা স্বাস্থ্যের জন্য উপকারী।"],
    ["বাংলাদেশের প্রাকৃতিক সৌন্দর্য সকলকে মুগ্ধ করে।"]
]

# Enhanced visual layout and interface design
with gr.Blocks(theme="grass") as demo:
    gr.Markdown("# BanglaT5 Model Hub - Translation, Summarization & Paraphrasing")
    gr.Markdown("Explore the power of BanglaT5 with this easy-to-use interface for multiple tasks like translation, summarization, and paraphrasing.")

    with gr.Tabs():
        # Tab for English to Bengali Translation
        with gr.Tab("Translate English to Bengali"):
            en_bn_text = gr.Textbox(label="Enter English Text", lines=5)
            en_bn_output = gr.Textbox(label="Bengali Translation")
            gr.Examples(examples_translate_en_bn, inputs=en_bn_text)
            en_bn_button = gr.Button("Show Output")
            en_bn_button.click(translate_text_en_bn, inputs=en_bn_text, outputs=en_bn_output)

        # Tab for Bengali to English Translation
        with gr.Tab("Translate Bengali to English"):
            bn_en_text = gr.Textbox(label="Enter Bengali Text", lines=5)
            bn_en_output = gr.Textbox(label="English Translation")
            gr.Examples(examples_translate_bn_en, inputs=bn_en_text)
            bn_en_button = gr.Button("Show Output")
            bn_en_button.click(translate_text_bn_en, inputs=bn_en_text, outputs=bn_en_output)

        # Tab for Summarization
        with gr.Tab("Summarize"):
            summarize_text_input = gr.Textbox(label="Enter Bengali Text to Summarize", lines=5)
            summarize_output = gr.Textbox(label="Summary")
            gr.Examples(examples_summarize, inputs=summarize_text_input)
            summarize_button = gr.Button("Show Output")
            summarize_button.click(summarize_text, inputs=summarize_text_input, outputs=summarize_output)

        # Tab for Paraphrasing
        with gr.Tab("Paraphrase"):
            paraphrase_text_input = gr.Textbox(label="Enter Bengali Text to Paraphrase", lines=5)
            paraphrase_output = gr.Textbox(label="Paraphrased Text")
            gr.Examples(examples_paraphrase, inputs=paraphrase_text_input)
            paraphrase_button = gr.Button("Show Output")
            paraphrase_button.click(paraphrase_text, inputs=paraphrase_text_input, outputs=paraphrase_output)

demo.launch()