from initial import *


def document_retriever(question):
    """
    Retrieve relevant documents (contexts) for a given question.

    Args:
    - question (str): The question to retrieve documents for.

    Returns:
    - top_docs (list): List of dictionaries containing top relevant documents.
    """
    # Preprocess the question
    preprocessed_question = [tfidf_preprocess(question)]
    question_vector = tfidf_vectorizer.transform(preprocessed_question)

    # Calculate similarity scores
    score = cosine_similarity(tfidf_matrix, question_vector)

    # Get the top 5 relevant documents (contexts)
    top_5_indices = score.flatten().argsort()[-5:][::-1]
    top_5_scores = score.flatten()[top_5_indices]

    top_docs = []
    for i, idx in enumerate(top_5_indices):
        top_docs.append(
            {
                "title": document_context[idx],
                "context": contexts[idx],
                "score": top_5_scores[i],
            }
        )

    return top_docs


def load_model(model_name):
    """
    Load a pre-trained question answering model.

    Args:
    - model_name (str): The name of the model to load.

    Returns:
    - qa_pipeline: Question answering pipeline with the loaded model.
    """
    if model_name in MODEL_NAME:
        model = AutoModelForQuestionAnswering.from_pretrained(model_name)
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        qa_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer)
        return qa_pipeline
    else:
        raise ValueError(f"Model {model_name} not found")