Spaces:
Running
Running
File size: 5,740 Bytes
de37907 6990294 de37907 d8dca23 de37907 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
from rank_bm25 import BM25Okapi
from tokenizing import tokenize_text, tokenize_doc, tokenize_doc_to_str
import numpy as np
from tqdm import tqdm
import os
import pickle
import torch
from sentence_transformers import SentenceTransformer
# Define the base path
base_path = "./Data" # "/mnt/d/Semester7/NLP/RAG/Data"
class Retriever:
def __init__(self, docs: [dict]) -> None:
self.docs = docs
self.tokenized_docs_path = os.path.join(base_path, "tokenized_docs.pkl")
self.bm25_path = os.path.join(base_path, "bm25.pkl")
self.sbert_embeddings_path = os.path.join(base_path, "embeddings_parts")
# Initialize SBERT
device = 'cuda' if torch.cuda.is_available() else 'cpu'
self.sbert = SentenceTransformer('sentence-transformers/all-distilroberta-v1', device=device)
# Load or tokenize documents
if os.path.exists(self.tokenized_docs_path) and os.path.exists(self.bm25_path) and os.path.exists(self.sbert_embeddings_path):
print("Loading cache")
self.load_cache()
print("Cache loaded")
else:
self.tokenize_and_initialize()
def tokenize_and_initialize(self):
# Tokenize the documents with a progress bar
self.tokenized_docs = [tokenize_doc(doc) for doc in tqdm(self.docs, desc="Tokenizing documents")]
self.str_docs = [tokenize_doc_to_str(doc) for doc in self.docs]
# Ensure that the tokenized_docs list is not empty
if not self.tokenized_docs:
raise ValueError("The list of tokenized documents is empty. Please check the input documents and the tokenization process.")
# Initialize BM25 with the tokenized texts
self.bm25 = BM25Okapi(self.tokenized_docs)
# Get embeddings for all the tocenized documents
self.sbert_embeddings = self.sbert.encode(self.str_docs, show_progress_bar=True)
self.sbert_embeddings = self.sbert_embeddings.cpu()
# Save the tokenized documents and BM25 model
self.save_cache()
def save_cache(self):
with open(self.tokenized_docs_path, 'wb') as f:
pickle.dump(self.tokenized_docs, f)
with open(self.bm25_path, 'wb') as f:
pickle.dump(self.bm25, f)
split_size = 1000 # Number of rows per split
embeddings_size = self.sbert_embeddings.size(0)
# Split and save the tensor
for i in range(0, embeddings_size, split_size):
end_idx = min(i + split_size, embeddings_size)
part = self.sbert_embeddings[i:end_idx].clone()
torch.save(part, os.path.join(self.sbert_embeddings_path, f"embeddings_part_{i//split_size}.pt"))
def load_cache(self):
with open(self.tokenized_docs_path, 'rb') as f:
self.tokenized_docs = pickle.load(f)
with open(self.bm25_path, 'rb') as f:
self.bm25 = pickle.load(f)
print("Loading SBERT embeddings")
# Load and combine
loaded_parts = []
files = os.listdir(self.sbert_embeddings_path)
# Sort numerically based on the numeric part of the filename
sorted_files = sorted(files, key=lambda x: int(x.split('_')[2].split('.')[0]))
counter = 0
for file in sorted_files:
if file.startswith("embeddings_part_") and file.endswith(".pt"):
part_path = os.path.join(self.sbert_embeddings_path, file)
loaded_parts.append(torch.load(part_path))
if counter % 50 == 0:
print("Loaded", file)
counter += 1
self.sbert_embeddings = torch.cat(loaded_parts, dim=0)
print("SBERT embeddings loaded")
def get_docs(self, user_message: str, n: int = 30, bm25_only: bool = False, semantic_only: bool = False, scores_combination: bool = True, bm_koef: float = 0.75) -> [str]:
# In case of BM25 only, return the top n documents based on BM25 scores, if somebody sets a couple
# of flags to True, the func will return the top n documents based on the first flag set to True
# remove "tell me a joke about" ot "tell me a joke and its title about" from the user message
user_message = user_message.replace("tell me a joke about", "").replace("tell me a joke and its title about", "")
if bm25_only:
semantic_only = False
scores_combination = False
print("BM25 only")
scores = torch.tensor(self._get_bm25_scores(user_message))
elif semantic_only:
scores_combination = False
print("Semantic only")
scores = self.get_semantic_scores(user_message)
elif scores_combination:
print("Combination")
bm_scores = self._get_bm25_scores(user_message)
semantic_scores = self.get_semantic_scores(user_message)
scores = torch.tensor(bm_koef * bm_scores) + (1 - bm_koef) * semantic_scores
# Sort the documents by their BM25 scores in descending order
sorted_doc_indices = np.argsort(scores)
result_docs = [self.docs[i] for i in sorted_doc_indices[-n:] if scores[i] > 0]
return result_docs[::-1] # Return the top n documents in descending order which means the most relevant documents are first
def _get_bm25_scores(self, user_message: str) -> np.array:
tokenized_user_message = tokenize_text(user_message)
return self.bm25.get_scores(tokenized_user_message)
def get_semantic_scores(self, user_message: str) -> np.array:
user_message_embedding = self.sbert.encode(user_message)
scores = self.sbert.similarity(user_message_embedding, self.sbert_embeddings)
return scores[0] |