import gradio as gr
from course_content import (
    course_introduction,
    how_language_model_works,
    practical_application_and_usecases_of_llms,
    what_is_generative_ai,
    llms_for_writing,
    llms_for_reading,
    llms_for_chatting,
    llms_capabilities_limitations,
    llms_image_generation
)
from simpletransformers.question_answering import QuestionAnsweringModel
# Initialize a Question Answering model
model = QuestionAnsweringModel(
    "bert",               # Specify the model type, e.g., 'bert', 'roberta', etc.
    "bert-base-uncased",  # Specify the model size and configuration
    use_cuda=False        # Set to True if you have CUDA enabled GPU
)


def get_answer(question):
    input_data = [
        {
            "id": 0,
            "context": llms_image_generation,
            "qas": [{"question": question, "id": 0}]
        }
    ]
    # Perform prediction
    results = model.predict(input_data)

    # Assuming results is the output of model.predict(input_data)
    answer_dict = results[0][0]  # Get the dictionary containing answers and IDs
    answer_segments = answer_dict['answer']  # Get the list of answer segments
    probabilities = results[1][0]['probability']  # Get the list of probabilities

    # Find the index of the answer segment with the highest probability
    best_index = probabilities.index(max(probabilities))
    best_answer = answer_segments[best_index]

    return best_answer