import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForQuestionAnswering

# Replace 'your_model_name' with the name of your fine-tuned model
model_name = 'skl25/banglabert-qa'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForQuestionAnswering.from_pretrained(model_name)

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

def answer_question(context, question):
    # Combine user-provided question with the context
    full_context = f"{question} {context}"

    # Tokenize input
    inputs = tokenizer.encode_plus(full_context, return_tensors="pt")

    # Get model output
    outputs = model(**inputs)
    answer_start = torch.argmax(outputs.start_logits)
    answer_end = torch.argmax(outputs.end_logits) + 1

    # Decode and return the answer along with the context
    answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end]))
    return f"Answer: {answer}"


# Define Gradio interface
iface = gr.Interface(
    fn=answer_question,
    inputs=["text","text"],
    outputs="text",
    live=False,
    title="Question Answering Widget using BanglaBERT",
    theme='enescakircali/Indian-Henna',
    description="Ask questions based on the predefined context.",
    examples=[
        ["বাংলাদেশ দক্ষিণ এশিয়ার একটি সার্বভৌম রাষ্ট্র। বাংলাদেশের সাংবিধানিক নাম গণপ্রজাতন্ত্রী বাংলাদেশ। ভৌগোলিকভাবে বাংলাদেশের পশ্চিমে ভারতের পশ্চিমবঙ্গ, উত্তরে পশ্চিমবঙ্গ, আসাম ও মেঘালয়, পূর্ব সীমান্তে আসাম, ত্রিপুরা ও মিজোরাম, দক্ষিণ-পূর্ব সীমান্তে মিয়ানমারের চিন ও রাখাইন রাজ্য এবং দক্ষিণ উপকূলের দিকে বঙ্গোপসাগর অবস্থিত। ভৌগোলিকভাবে পৃথিবীর বৃহত্তম ব-দ্বীপের সিংহভাগ অঞ্চল জুড়ে বাংলাদেশ ভূখণ্ড অবস্থিত। জনসংখ্যার বিচারে প্রায় ১৭ কোটিরও অধিক জনসংখ্যা নিয়ে বাংলাদেশ বিশ্বের ৮ম বৃহত্তম দেশ। নদীমাতৃক বাংলাদেশ ভূখণ্ডের উপর দিয়ে বয়ে গেছে ৫৭টি আন্তর্জাতিক নদী। বাংলাদেশের উত্তর-পূর্বে ও দক্ষিণ-পূর্বে টারশিয়ারি যুগের পাহাড় ছেয়ে আছে। বিশ্বের বৃহত্তম ম্যানগ্রোভ অরণ্য সুন্দরবন ও দীর্ঘতম প্রাকৃতিক সৈকত কক্সবাজার সমুদ্র সৈকত বাংলাদেশে অবস্থিত।",
        "বাংলাদেশের সাংবিধানিক নাম কি?"],
        
        ["স্বচ্ছ পান্নাসবুজ জলে সাঁতার, বালুময় সৈকতে সূর্যস্নান, প্রবালদ্বীপের বিলাসবহুল রিসোর্টে ছুটিযাপন, স্নোরকেলিং ও স্কুবা ডাইভিংয়ের মতো নানা রোমাঞ্চকর কর্মকাণ্ড—বলা যায় পর্যটকদের চাহিদা পূরণে সমস্ত সম্ভার নিয়ে বসে আছে মালদ্বীপ। তাই তো হানিমুনের জন্য নবদম্পতিদের কাছে যেমন, তেমনি রোমাঞ্চপ্রিয় ও নিরিবিলি আয়েশি অবকাশযাপনসন্ধানী ভ্রমণপিপাসুদের কাছেও পছন্দের এক গন্তব্য ভারত মহাসাগরের এই দ্বীপরাষ্ট্র। দেশটির সহজ ভিসানীতি, নিরাপত্তা আর নির্ঝঞ্ঝাট পরিবেশও পর্যটকদের আকর্ষণ করে। ঢাকা থেকে মালের সরাসরি ফ্লাইট, বিভিন্ন এয়ারলাইনসসহ ট্যুর পরিচালনাকারী প্রতিষ্ঠানগুলোর সাশ্রয়ী ভ্রমণ প্যাকেজের ফলে দিন দিন অনেক বাংলাদেশির কাছেও পছন্দের গন্তব্য হয়ে উঠেছে মালদ্বীপ। দেশটির সরকারি তথ্য বলছে, ২০২১ সালে যেখানে মাত্র ৩ হাজার ৯২৩ বাংলাদেশি ভারত মহাসাগরের দ্বীপরাষ্ট্রটি ভ্রমণে গিয়েছিল, ২০২২ সালে একলাফে সেই সংখ্যা গিয়ে দাঁড়ায় ১৬ হাজার ৮০৭-তে। সেই বৃদ্ধির ধারাবাহিকতা গত বছরও লক্ষ করা গেছে। ২০২৩ সালে মালদ্বীপ ভ্রমণ করেছেন ২৮ হাজার ৩৩৬ বাংলাদেশি। অর্থাৎ ২০২২ সালের তুলনায় দেশটিতে বাংলাদেশি পর্যটক বেড়েছে ৬৮ দশমিক ৬ শতাংশ। এই হিসাবে বাংলাদেশ এখন মালদ্বীপের ১৫তম পর্যটন বাজার।",
        "মালদ্বীপ পর্যটকদের কি কারণে আকর্ষণ করে?"]
    ]
        )

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