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")

# Function to perform machine translation
def translate_text_en_bn(input_text):
    inputs = translation_tokenizer_en_bn(input_text, return_tensors="pt")
    outputs = translation_model_en_bn.generate(**inputs)
    translated_text = translation_tokenizer_en_bn.decode(outputs[0], skip_special_tokens=True)
    return translated_text

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

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

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

def process_text(text, task):
    if task == "Translate_English_to_Bengali":
        return translate_text_en_bn(text)
    elif task == "Translate_Bengali_to_English":
        return translate_text_bn_en(text)
    elif task == "Summarize":
        return summarize_text(text)
    elif task == "Paraphrase":
        return paraphrase_text(text)

# Define the Gradio interface
iface = gr.Interface(
    fn=process_text,
    inputs=["text", gr.Dropdown(["Translate_English_to_Bengali", "Translate_Bengali_to_English", "Summarize", "Paraphrase"])],
    outputs="text",
    live=False,
    title="Usage of BanglaT5 Model",
    theme='Ajaxon6255/Emerald_Isle',
    examples=[
        ["বাংলাদেশ দক্ষিণ এশিয়ার একটি সার্বভৌম রাষ্ট্র। বাংলাদেশের সাংবিধানিক নাম গণপ্রজাতন্ত্রী বাংলাদেশ। ভৌগোলিকভাবে বাংলাদেশের পশ্চিমে ভারতের পশ্চিমবঙ্গ, উত্তরে পশ্চিমবঙ্গ, আসাম ও মেঘালয়, পূর্ব সীমান্তে আসাম, ত্রিপুরা ও মিজোরাম, দক্ষিণ-পূর্ব সীমান্তে মিয়ানমারের চিন ও রাখাইন রাজ্য এবং দক্ষিণ উপকূলের দিকে বঙ্গোপসাগর অবস্থিত। ভৌগোলিকভাবে পৃথিবীর বৃহত্তম ব-দ্বীপের সিংহভাগ অঞ্চল জুড়ে বাংলাদেশ ভূখণ্ড অবস্থিত। জনসংখ্যার বিচারে প্রায় ১৭ কোটিরও অধিক জনসংখ্যা নিয়ে বাংলাদেশ বিশ্বের ৮ম বৃহত্তম দেশ। নদীমাতৃক বাংলাদেশ ভূখণ্ডের উপর দিয়ে বয়ে গেছে ৫৭টি আন্তর্জাতিক নদী। বাংলাদেশের উত্তর-পূর্বে ও দক্ষিণ-পূর্বে টারশিয়ারি যুগের পাহাড় ছেয়ে আছে। বিশ্বের বৃহত্তম ম্যানগ্রোভ অরণ্য সুন্দরবন ও দীর্ঘতম প্রাকৃতিক সৈকত কক্সবাজার সমুদ্র সৈকত বাংলাদেশে অবস্থিত।"],
        ["বাংলাদেশের সংবিধানিক নাম কি?"], 
        ["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."],
        ["তিনজনের কেউই বাবার পথ ধরে প্রযুক্তি দুনিয়ায় হাঁটেননি।"],
        ["বিশ্বের আরও একটি সেরা ক্লাব। স্প্যানিশ ফুটবলে বার্সেলোনারই প্রবল প্রতিদ্বন্দ্বী রিয়াল মাদ্রিদ। লিওনেল মেসির ক্লাব ইতিহাসে উয়েফা চ্যাম্পিয়নস লিগ জিতেছে মোট পাঁচবার। ইউরোপীয় কাপ উইনার্স কাপে তাদের সাফল্য আছে চারবার। পাশাপাশি তিনবার ক্লাব বিশ্বকাপ জেতা বার্সা লা লিগায় সাফল্য পেয়েছেন মোট ২৪ বার। ডিয়েগো ম্যারাডোনা, ইয়োহান ক্রুইফ, রিভালদোদের মতো খেলোয়াড়েরা এই ক্লাবের হয়ে অংশ হয়ে আছেন ইতিহাসের। "]

    ]
)

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