retrieval.py

import json
import numpy as np
from langchain.schema import Document
import faiss
from rank_bm25 import BM25Okapi
from data_processing import embedding_model
from sentence_transformers import CrossEncoder

#import string
# import nltk
# nltk.download('punkt')  
# nltk.download('punkt_tab')

from nltk.tokenize import word_tokenize

reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

retrieved_docs = None

# Tokenize the documents and remove punctuation
# def preprocess(doc):
#     return [word.lower() for word in word_tokenize(doc) if word not in string.punctuation]

def retrieve_documents_hybrid(query, q_dataset, top_k=5):
    with open( f"data_local/{q_dataset}_chunked_docs.json", "r") as f:
        chunked_documents = json.load(f)  # Contains all documents for this dataset
    
    faiss_index_path = f"data_local/{q_dataset}_quantized.faiss"
    index = faiss.read_index(faiss_index_path)

    # Tokenize documents for BM25
    tokenized_docs = [doc.split() for doc in chunked_documents]
    bm25 = BM25Okapi(tokenized_docs)

    query_embedding = np.array(embedding_model.embed_documents([query]), dtype=np.float32)
    query_embedding = query_embedding.reshape(1, -1)

    # FAISS Search
    _, faiss_indices = index.search(query_embedding, top_k)
    faiss_docs = [chunked_documents[i] for i in faiss_indices[0]]

    # BM25 Search  
    tokenized_query = query.split() #preprocess(query)
    bm25_scores = bm25.get_scores(tokenized_query)
    bm25_top_indices = np.argsort(bm25_scores)[::-1][:top_k]
    bm25_docs = [chunked_documents[i] for i in bm25_top_indices]

    # Combine FAISS + BM25 scores and retrieve docs
    # combined_results = set(bm25_top_indices).union(set(faiss_indices[0])) 
        
    # combined_scores = rerank_docs_bm25faiss_scores(combined_results,bm25_scores, faiss_distances,faiss_indices)
    # reranked_docs = [chunked_documents[result[0]] for result in combined_scores[:top_k]]

    # Merge FAISS + BM25 Results and re-rank
    retrieved_docs = list(set(faiss_docs + bm25_docs))[:top_k]
    reranked_docs = rerank_documents(query, retrieved_docs)

    return reranked_docs

def rerank_docs_bm25faiss_scores(combined_results_,bm25_scores_, faiss_distances_,faiss_indices_):
    final_results = []
    for idx in combined_results_:
        # Combine BM25 score and FAISS score for ranking (this could be more sophisticated)
        bm25_score = bm25_scores_[idx]
        faiss_score = 1 / (1 + faiss_distances_[0][np.where(faiss_indices_[0] == idx)])  # Inverse distance for relevance
        final_results.append((idx, bm25_score, faiss_score))

    # Sort final results by combined score (you can adjust the ranking strategy here)
    final_results.sort(key=lambda x: (x[1] + x[2]), reverse=True)

    return final_results
    

# Retrieval Function
# def retrieve_documents(query, top_k=5):
#     query_dataset = find_query_dataset(query)
#     #index, chunk_docs = load_data_from_faiss(query)

#     with open( f"data_local/{query_dataset}_chunked_docs.json", "r") as f:
#         documents = json.load(f)  # Contains all documents for this dataset

#     faiss_index_path = f"data_local/{query_dataset}_quantized.faiss"
#     index = faiss.read_index(faiss_index_path)

#     query_embedding = np.array(embedding_model.embed_documents([query]), dtype=np.float32)

#     _, nearest_indices = index.search(query_embedding, top_k)

#     retrieved_docs = [Document(page_content=documents[i]) for i in nearest_indices[0]]

#     return retrieved_docs

def remove_duplicate_documents(documents):
    unique_documents = []
    seen_documents = set() 
    for doc in documents:
        doc_content = doc.page_content
        if doc_content not in seen_documents:
            unique_documents.append(doc)
            seen_documents.add(doc_content)
    return unique_documents

def find_query_dataset(query):
    index = faiss.read_index("data_local/question_quantized.faiss")

    with open("data_local/dataset_mapping.json", "r") as f:
        dataset_names = json.load(f)

    question_embedding = np.array(embedding_model.embed_documents([query]), dtype=np.float32)
    _, nearest_index = index.search(question_embedding, 1)  
    best_dataset = dataset_names[nearest_index[0][0]]
    return best_dataset

def rerank_documents(query, retrieved_docs):
    doc_texts = [doc for doc in retrieved_docs]
    scores = reranker.predict([[query, doc] for doc in doc_texts])
    ranked_docs = [doc for _, doc in sorted(zip(scores, retrieved_docs), reverse=True)]
    return ranked_docs[:5]  # Return top k most relevant