import os os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") os.environ.setdefault("HF_HUB_OFFLINE", "1") # never hit the Hub os.environ.setdefault("TRANSFORMERS_OFFLINE", "1") # Transformers offline os.environ.setdefault("HF_DATASETS_OFFLINE", "1") # Datasets offline os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1") os.environ.setdefault("DISABLE_TQDM", '1') os.environ.setdefault("TQDM_DISABLE", '1') os.environ["TQDM_DISABLE"] = "1" os.environ["DISABLE_TQDM"] = "1" import json import torch torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True import torch.nn.functional as F import numpy as np from transformers import AutoTokenizer, AutoModel from sklearn.metrics.pairwise import cosine_similarity import string from sklearn.feature_extraction import _stop_words import time import bm25s class BM25sRetriever: def __init__(self, passages, ids): self.corpus_tokens = bm25s.tokenize(passages, show_progress=False) self.retriever = bm25s.BM25() self.retriever.index(self.corpus_tokens, show_progress=False, leave_progress=False) self.corpus_ids = ids self.id2idx = {} for i in range(len(self.corpus_ids)): self.id2idx[self.corpus_ids[i]] = i def get_scores(self, query_text): query_tokens = bm25s.tokenize(query_text) return self.retriever.retrieve(query_tokens, sorted=False, k=len(self.corpus_ids), show_progress=False, leave_progress=False)[1][0] class E5Retriever: def __init__(self, model_name=None, device=None): """ Initializes the E5 retriever using the multilingual E5 base model. """ # Use local model if model_name is None: # local_model_path = os.path.join('sub/models', 'multilingual-e5-large') # 'multilingual-e5-large' local_model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models', 'multilingual-e5-large_pseudo_full') # 'multilingual-e5-large' if os.path.isdir(local_model_path): model_name = local_model_path print(f"Using local E5 model from: {model_name}") self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") # Clear GPU cache before loading model if torch.cuda.is_available(): torch.cuda.empty_cache() print(f"Loading E5 multilingual model on device: {self.device}") self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.model = AutoModel.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(self.device) self.model.eval() # Clear cache after model loading if torch.cuda.is_available(): torch.cuda.empty_cache() self.corpus_ids = [] self.corpus_embeddings = None def embed_texts(self, texts, is_query=False, batch_size=32): """ Generates embeddings for texts using E5 model with proper prefixes. E5 requires specific prefixes for queries vs passages. """ # E5 model requires specific prefixes if is_query: # Add query prefix for E5 prefixed_texts = [f"query: {text.strip()}" for text in texts] else: # Add passage prefix for E5 prefixed_texts = [f"passage: {text.strip()}" for text in texts] all_embeddings = [] total_batches = (len(prefixed_texts) + batch_size - 1) // batch_size for i in range(0, len(prefixed_texts), batch_size): batch_num = i // batch_size + 1 if not is_query and batch_num % 50 == 0: print(f"Processing batch {batch_num}/{total_batches} ({(batch_num/total_batches)*100:.1f}%)") if torch.cuda.is_available(): torch.cuda.empty_cache() batch_texts = prefixed_texts[i:i + batch_size] try: encoded = self.tokenizer( batch_texts, padding=True, truncation=True, max_length=512, return_tensors='pt' ).to(self.device) with torch.no_grad(): model_output = self.model(**encoded) # E5 uses mean pooling with attention mask attention_mask = encoded['attention_mask'] embeddings = model_output.last_hidden_state # Mean pooling mask_expanded = attention_mask.unsqueeze(-1).expand(embeddings.size()).float() sum_embeddings = torch.sum(embeddings * mask_expanded, 1) sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9) embeddings = sum_embeddings / sum_mask # L2 normalize embeddings (important for E5) embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) # Move to CPU immediately all_embeddings.append(embeddings.cpu()) # Clear GPU memory del encoded, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except torch.cuda.OutOfMemoryError as e: print(f"CUDA OOM at batch {batch_num}, reducing batch size...") # Process one item at a time for single_text in batch_texts: try: encoded = self.tokenizer( [single_text], padding=True, truncation=True, max_length=512, return_tensors='pt' ).to(self.device) with torch.no_grad(): model_output = self.model(**encoded) attention_mask = encoded['attention_mask'] embeddings = model_output.last_hidden_state # Mean pooling mask_expanded = attention_mask.unsqueeze(-1).expand(embeddings.size()).float() sum_embeddings = torch.sum(embeddings * mask_expanded, 1) sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9) embeddings = sum_embeddings / sum_mask embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) all_embeddings.append(embeddings.cpu()) del encoded, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except Exception as e2: print(f"Failed to process single text: {e2}") # E5-base has 768 dimensions. Large 1024 zero_embedding = torch.zeros(1, 1024).float() all_embeddings.append(zero_embedding) return torch.cat(all_embeddings, dim=0).numpy() def prepare_corpus(self, texts, ids, batch_size=32): # Add passage prefix for E5 prefixed_texts = [f"passage: {text.strip()}" for text in texts] all_embeddings = [] all_tokens = [] all_ids = [] all_len = [] prefix_len = 3 suffix_len = 1 max_length = 512 max_passage = max_length - prefix_len - suffix_len overlap = max_passage // 2 for i in range(len(prefixed_texts)): encoded = self.tokenizer( [prefixed_texts[i]], padding=False, truncation=False, # max_length=512, # return_tensors='pt' ) if len(encoded['input_ids'][0]) > max_length: _idxs = list(range(prefix_len, len(encoded['input_ids'][0])-suffix_len)) i0 = 0 while i0 < len(_idxs) - overlap: i1 = min(i0+max_passage, len(_idxs)) all_ids.append((ids[i], (i0, i1, len(_idxs)))) all_tokens.append({ 'input_ids': encoded['input_ids'][0][:prefix_len] + encoded['input_ids'][0][i0+prefix_len:i1+prefix_len] + encoded['input_ids'][0][-suffix_len:], 'attention_mask': encoded['attention_mask'][0][:prefix_len] + encoded['attention_mask'][0][i0+prefix_len:i1+prefix_len] + encoded['attention_mask'][0][-suffix_len:] }) all_len.append(i1-i0) i0 += max_passage - overlap else: all_ids.append((ids[i], None)) all_tokens.append({'input_ids': encoded['input_ids'][0], 'attention_mask': encoded['attention_mask'][0]}) all_len.append(len(encoded['input_ids'][0]) - prefix_len - suffix_len) total_batches = (len(all_tokens) + batch_size - 1) // batch_size all_len, all_tokens, all_ids = zip(*sorted(zip(all_len, all_tokens, all_ids), key=lambda x: x[0])) for i in range(0, len(all_tokens), batch_size): batch_num = i // batch_size + 1 if batch_num % 50 == 0: print(f"Processing batch {batch_num}/{total_batches} ({(batch_num/total_batches)*100:.1f}%)") if torch.cuda.is_available(): torch.cuda.empty_cache() batch_tokens = all_tokens[i:i + batch_size] batch_max = max([len(ids['input_ids']) for ids in batch_tokens]) encoded = dict() encoded["attention_mask"] = [s['attention_mask'] + (batch_max - len(s['attention_mask'])) * [0] for s in batch_tokens] encoded["input_ids"] = [s['input_ids'] + (batch_max - len(s['attention_mask'])) * [0] for s in batch_tokens] encoded["attention_mask"] = torch.tensor(encoded["attention_mask"], dtype=torch.long).to(self.device) encoded["input_ids"] = torch.tensor(encoded["input_ids"], dtype=torch.long).to(self.device) try: with torch.no_grad(): model_output = self.model(**encoded) # E5 uses mean pooling with attention mask attention_mask = encoded['attention_mask'] embeddings = model_output.last_hidden_state # Mean pooling mask_expanded = attention_mask.unsqueeze(-1).expand(embeddings.size()).float() sum_embeddings = torch.sum(embeddings * mask_expanded, 1) sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9) embeddings = sum_embeddings / sum_mask # L2 normalize embeddings (important for E5) embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) # Move to CPU immediately all_embeddings.append(embeddings.cpu()) # Clear GPU memory del encoded, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except torch.cuda.OutOfMemoryError as e: print(f"CUDA OOM at batch {batch_num}, reducing batch size...") # Process one item at a time for j in range(len(encoded)): try: encoded0 = encoded[j:j+1] with torch.no_grad(): model_output = self.model(**encoded0) attention_mask = encoded0['attention_mask'] embeddings = model_output.last_hidden_state # Mean pooling mask_expanded = attention_mask.unsqueeze(-1).expand(embeddings.size()).float() sum_embeddings = torch.sum(embeddings * mask_expanded, 1) sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9) embeddings = sum_embeddings / sum_mask embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) all_embeddings.append(embeddings.cpu()) del encoded0, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except Exception as e2: print(f"Failed to process single text: {e2}") # E5-base has 768 dimensions. Large 1024 zero_embedding = torch.zeros(1, 1024).float() all_embeddings.append(zero_embedding) id2idx = {} for i in range(len(all_ids)): if all_ids[i][0] not in id2idx: id2idx[all_ids[i][0]] = [] id2idx[all_ids[i][0]].append(i) return torch.cat(all_embeddings, dim=0).to(self.model.device), all_ids, id2idx # numpy() class E5InstructRetriever: def __init__(self, model_name=None, device=None): """ Initializes the E5 Instruct retriever using the multilingual E5 Instruct large model. """ # Use local model if model_name is None: # local_model_path = os.path.join('sub/models', 'multilingual-e5-large-instruct') local_model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models', 'multilingual-e5-large-instruct') if os.path.isdir(local_model_path): model_name = local_model_path print(f"Using local E5 Instruct model from: {model_name}") self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") # Clear GPU cache before loading model if torch.cuda.is_available(): torch.cuda.empty_cache() print(f"Loading E5 Instruct multilingual model on device: {self.device}") self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.model = AutoModel.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(self.device) # torch.float16 self.model.eval() # Clear cache after model loading if torch.cuda.is_available(): torch.cuda.empty_cache() self.corpus_ids = [] self.corpus_embeddings = None def embed_texts(self, texts, is_query=False, batch_size=32): """ Generates embeddings for texts using E5 Instruct model with proper prefixes. E5 Instruct requires specific prefixes for queries vs passages. """ task = 'Given a web search query, retrieve relevant passages that answer the query' # E5 model requires specific prefixes if is_query: # Add query prefix for E5 prefixed_texts = [f"Instruct: {task}\nQuery: {text.strip()}" for text in texts] else: # Add passage prefix for E5 prefixed_texts = [f"{text.strip()}" for text in texts] all_embeddings = [] total_batches = (len(prefixed_texts) + batch_size - 1) // batch_size for i in range(0, len(prefixed_texts), batch_size): batch_num = i // batch_size + 1 if not is_query and batch_num % 50 == 0: print(f"Processing batch {batch_num}/{total_batches} ({(batch_num/total_batches)*100:.1f}%)") if torch.cuda.is_available(): torch.cuda.empty_cache() batch_texts = prefixed_texts[i:i + batch_size] try: encoded = self.tokenizer( batch_texts, padding=True, truncation=True, max_length=512, return_tensors='pt' ).to(self.device) with torch.no_grad(): model_output = self.model(**encoded) # E5 uses mean pooling with attention mask attention_mask = encoded['attention_mask'] embeddings = model_output.last_hidden_state # Mean pooling mask_expanded = attention_mask.unsqueeze(-1).expand(embeddings.size()).float() sum_embeddings = torch.sum(embeddings * mask_expanded, 1) sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9) embeddings = sum_embeddings / sum_mask # L2 normalize embeddings (important for E5) embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) # Move to CPU immediately all_embeddings.append(embeddings.cpu()) # Clear GPU memory del encoded, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except torch.cuda.OutOfMemoryError as e: print(f"CUDA OOM at batch {batch_num}, reducing batch size...") # Process one item at a time for single_text in batch_texts: try: encoded = self.tokenizer( [single_text], padding=True, truncation=True, max_length=512, return_tensors='pt' ).to(self.device) with torch.no_grad(): model_output = self.model(**encoded) attention_mask = encoded['attention_mask'] embeddings = model_output.last_hidden_state # Mean pooling mask_expanded = attention_mask.unsqueeze(-1).expand(embeddings.size()).float() sum_embeddings = torch.sum(embeddings * mask_expanded, 1) sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9) embeddings = sum_embeddings / sum_mask embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) all_embeddings.append(embeddings.cpu()) del encoded, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except Exception as e2: print(f"Failed to process single text: {e2}") # E5-base has 768 dimensions. Large 1024 zero_embedding = torch.zeros(1, 1024).float() all_embeddings.append(zero_embedding) return torch.cat(all_embeddings, dim=0).numpy() def prepare_corpus(self, texts, ids, batch_size=32): # Add passage prefix for E5 prefixed_texts = [f"{text.strip()}" for text in texts] all_embeddings = [] all_tokens = [] all_ids = [] all_len = [] prefix_len = 1 suffix_len = 1 max_length = 512 max_passage = max_length - prefix_len - suffix_len overlap = max_passage // 2 for i in range(len(prefixed_texts)): # tqdm( encoded = self.tokenizer( [prefixed_texts[i]], padding=False, truncation=False, # max_length=512, # return_tensors='pt' ) if len(encoded['input_ids'][0]) > max_length: _idxs = list(range(prefix_len, len(encoded['input_ids'][0])-suffix_len)) i0 = 0 while i0 < len(_idxs) - overlap: i1 = min(i0+max_passage, len(_idxs)) all_ids.append((ids[i], (i0, i1, len(_idxs)))) all_tokens.append({ 'input_ids': encoded['input_ids'][0][:prefix_len] + encoded['input_ids'][0][i0+prefix_len:i1+prefix_len] + encoded['input_ids'][0][-suffix_len:], 'attention_mask': encoded['attention_mask'][0][:prefix_len] + encoded['attention_mask'][0][i0+prefix_len:i1+prefix_len] + encoded['attention_mask'][0][-suffix_len:] }) all_len.append(i1-i0) i0 += max_passage - overlap else: all_ids.append((ids[i], None)) all_tokens.append({'input_ids': encoded['input_ids'][0], 'attention_mask': encoded['attention_mask'][0]}) all_len.append(len(encoded['input_ids'][0]) - prefix_len - suffix_len) total_batches = (len(all_tokens) + batch_size - 1) // batch_size all_len, all_tokens, all_ids = zip(*sorted(zip(all_len, all_tokens, all_ids), key=lambda x: x[0])) for i in range(0, len(all_tokens), batch_size): batch_num = i // batch_size + 1 if batch_num % 50 == 0: print(f"Processing batch {batch_num}/{total_batches} ({(batch_num/total_batches)*100:.1f}%)") if torch.cuda.is_available(): torch.cuda.empty_cache() batch_tokens = all_tokens[i:i + batch_size] batch_max = max([len(ids['input_ids']) for ids in batch_tokens]) encoded = dict() encoded["attention_mask"] = [s['attention_mask'] + (batch_max - len(s['attention_mask'])) * [0] for s in batch_tokens] encoded["input_ids"] = [s['input_ids'] + (batch_max - len(s['attention_mask'])) * [0] for s in batch_tokens] encoded["attention_mask"] = torch.tensor(encoded["attention_mask"], dtype=torch.long).to(self.device) encoded["input_ids"] = torch.tensor(encoded["input_ids"], dtype=torch.long).to(self.device) try: with torch.no_grad(): model_output = self.model(**encoded) # E5 uses mean pooling with attention mask attention_mask = encoded['attention_mask'] embeddings = model_output.last_hidden_state # Mean pooling mask_expanded = attention_mask.unsqueeze(-1).expand(embeddings.size()).float() sum_embeddings = torch.sum(embeddings * mask_expanded, 1) sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9) embeddings = sum_embeddings / sum_mask # L2 normalize embeddings (important for E5) embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) # Move to CPU immediately all_embeddings.append(embeddings.cpu()) # Clear GPU memory del encoded, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except torch.cuda.OutOfMemoryError as e: print(f"CUDA OOM at batch {batch_num}, reducing batch size...") # Process one item at a time for j in range(len(encoded)): try: encoded0 = encoded[j:j+1] with torch.no_grad(): model_output = self.model(**encoded0) attention_mask = encoded0['attention_mask'] embeddings = model_output.last_hidden_state # Mean pooling mask_expanded = attention_mask.unsqueeze(-1).expand(embeddings.size()).float() sum_embeddings = torch.sum(embeddings * mask_expanded, 1) sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9) embeddings = sum_embeddings / sum_mask embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) all_embeddings.append(embeddings.cpu()) del encoded0, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except Exception as e2: print(f"Failed to process single text: {e2}") # E5-base has 768 dimensions. Large 1024 zero_embedding = torch.zeros(1, 1024).float() all_embeddings.append(zero_embedding) id2idx = {} for i in range(len(all_ids)): if all_ids[i][0] not in id2idx: id2idx[all_ids[i][0]] = [] id2idx[all_ids[i][0]].append(i) return torch.cat(all_embeddings, dim=0).to(self.model.device), all_ids, id2idx # numpy() class SnowflakeInstructRetriever: def __init__(self, model_name=None, device=None): """ Initializes the retriever using the multilingual snowflake-arctic-embed-l-v2.0 model. """ # Use local model if model_name is None: # local_model_path = os.path.join('sub/models', 'snowflake-arctic-embed-l-v2.0') local_model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models', 'snowflake-arctic-embed-l-v2.0') if os.path.isdir(local_model_path): model_name = local_model_path print(f"Using local snowflake-arctic-embed-l-v2.0 model from: {model_name}") self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") # Clear GPU cache before loading model if torch.cuda.is_available(): torch.cuda.empty_cache() print(f"Loading snowflake-arctic-embed-l-v2.0 multilingual model on device: {self.device}") self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.model = AutoModel.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(self.device) self.model.eval() # Clear cache after model loading if torch.cuda.is_available(): torch.cuda.empty_cache() self.corpus_ids = [] self.corpus_embeddings = None def embed_texts(self, texts, is_query=False, batch_size=32): """ Generates embeddings for texts using snowflake-arctic-embed-l-v2.0 model with proper prefixes. snowflake-arctic-embed-l-v2.0 requires specific prefixes for queries vs passages. """ # E5 model requires specific prefixes if is_query: # Add query prefix for E5 prefixed_texts = [f"query: {text.strip()}" for text in texts] else: # Add passage prefix for E5 prefixed_texts = [f"{text.strip()}" for text in texts] all_embeddings = [] total_batches = (len(prefixed_texts) + batch_size - 1) // batch_size for i in range(0, len(prefixed_texts), batch_size): batch_num = i // batch_size + 1 if not is_query and batch_num % 50 == 0: print(f"Processing batch {batch_num}/{total_batches} ({(batch_num/total_batches)*100:.1f}%)") if torch.cuda.is_available(): torch.cuda.empty_cache() batch_texts = prefixed_texts[i:i + batch_size] try: encoded = self.tokenizer( batch_texts, padding=True, truncation=True, max_length=1024, # 512 return_tensors='pt' ).to(self.device) with torch.no_grad(): model_output = self.model(**encoded) embeddings = model_output[0][:, 0].float() # L2 normalize embeddings embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) # Move to CPU immediately all_embeddings.append(embeddings.cpu()) # Clear GPU memory del encoded, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except torch.cuda.OutOfMemoryError as e: print(f"CUDA OOM at batch {batch_num}, reducing batch size...") # Process one item at a time for single_text in batch_texts: try: encoded = self.tokenizer( [single_text], padding=True, truncation=True, max_length=512, return_tensors='pt' ).to(self.device) with torch.no_grad(): model_output = self.model(**encoded) embeddings = model_output[0][:, 0].float() embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) all_embeddings.append(embeddings.cpu()) del encoded, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except Exception as e2: print(f"Failed to process single text: {e2}") zero_embedding = torch.zeros(1, 1024).float() all_embeddings.append(zero_embedding) return torch.cat(all_embeddings, dim=0).numpy() def prepare_corpus(self, texts, ids, batch_size=32): prefixed_texts = [f"{text.strip()}" for text in texts] all_embeddings = [] all_tokens = [] all_ids = [] all_len = [] prefix_len = 1 suffix_len = 1 max_length = 1024 # 512 max_passage = max_length - prefix_len - suffix_len overlap = max_passage // 2 for i in range(len(prefixed_texts)): encoded = self.tokenizer( [prefixed_texts[i]], padding=False, truncation=False, # max_length=512, # return_tensors='pt' ) if len(encoded['input_ids'][0]) > max_length: _idxs = list(range(prefix_len, len(encoded['input_ids'][0])-suffix_len)) i0 = 0 while i0 < len(_idxs) - overlap: i1 = min(i0+max_passage, len(_idxs)) all_ids.append((ids[i], (i0, i1, len(_idxs)))) all_tokens.append({ 'input_ids': encoded['input_ids'][0][:prefix_len] + encoded['input_ids'][0][i0+prefix_len:i1+prefix_len] + encoded['input_ids'][0][-suffix_len:], 'attention_mask': encoded['attention_mask'][0][:prefix_len] + encoded['attention_mask'][0][i0+prefix_len:i1+prefix_len] + encoded['attention_mask'][0][-suffix_len:] }) all_len.append(i1-i0) i0 += max_passage - overlap else: all_ids.append((ids[i], None)) all_tokens.append({'input_ids': encoded['input_ids'][0], 'attention_mask': encoded['attention_mask'][0]}) all_len.append(len(encoded['input_ids'][0]) - prefix_len - suffix_len) total_batches = (len(all_tokens) + batch_size - 1) // batch_size all_len, all_tokens, all_ids = zip(*sorted(zip(all_len, all_tokens, all_ids), key=lambda x: x[0])) for i in range(0, len(all_tokens), batch_size): batch_num = i // batch_size + 1 if batch_num % 50 == 0: print(f"Processing batch {batch_num}/{total_batches} ({(batch_num/total_batches)*100:.1f}%)") if torch.cuda.is_available(): torch.cuda.empty_cache() batch_tokens = all_tokens[i:i + batch_size] batch_max = max([len(ids['input_ids']) for ids in batch_tokens]) encoded = dict() encoded["attention_mask"] = [s['attention_mask'] + (batch_max - len(s['attention_mask'])) * [0] for s in batch_tokens] encoded["input_ids"] = [s['input_ids'] + (batch_max - len(s['attention_mask'])) * [0] for s in batch_tokens] encoded["attention_mask"] = torch.tensor(encoded["attention_mask"], dtype=torch.long).to(self.device) encoded["input_ids"] = torch.tensor(encoded["input_ids"], dtype=torch.long).to(self.device) try: with torch.no_grad(): model_output = self.model(**encoded) embeddings = model_output[0][:, 0].float() # L2 normalize embeddings embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) # Move to CPU immediately all_embeddings.append(embeddings.cpu()) # Clear GPU memory del encoded, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except torch.cuda.OutOfMemoryError as e: print(f"CUDA OOM at batch {batch_num}, reducing batch size...") # Process one item at a time for j in range(len(encoded)): try: encoded0 = encoded[j:j+1] with torch.no_grad(): model_output = self.model(**encoded0) embeddings = model_output[0][:, 0].float() embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) all_embeddings.append(embeddings.cpu()) del encoded0, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except Exception as e2: print(f"Failed to process single text: {e2}") zero_embedding = torch.zeros(1, 1024).float() all_embeddings.append(zero_embedding) id2idx = {} for i in range(len(all_ids)): if all_ids[i][0] not in id2idx: id2idx[all_ids[i][0]] = [] id2idx[all_ids[i][0]].append(i) return torch.cat(all_embeddings, dim=0).to(self.model.device), all_ids, id2idx # numpy() class SolonRetriever: def __init__(self, model_name=None, device=None): """ Initializes the retriever using the multilingual Solon-embeddings-large-0.1 model. """ # Use local model if model_name is None: # local_model_path = os.path.join('sub/models', 'Solon-embeddings-large-0.1') local_model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models', 'Solon-embeddings-large-0.1') if os.path.isdir(local_model_path): model_name = local_model_path print(f"Using local Solon-embeddings-large-0.1 model from: {model_name}") self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") # Clear GPU cache before loading model if torch.cuda.is_available(): torch.cuda.empty_cache() print(f"Loading Solon-embeddings-large-0.1 multilingual model on device: {self.device}") self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.model = AutoModel.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(self.device) # torch.float16 self.model.eval() # Clear cache after model loading if torch.cuda.is_available(): torch.cuda.empty_cache() self.corpus_ids = [] self.corpus_embeddings = None def embed_texts(self, texts, is_query=False, batch_size=32): """ Generates embeddings for texts using Solon-embeddings-large-0.1 model with proper prefixes. Solon-embeddings-large-0.1 requires specific prefixes for queries vs passages. """ # E5 model requires specific prefixes if is_query: # Add query prefix prefixed_texts = [f"query : {text.strip()}" for text in texts] else: prefixed_texts = [f"{text.strip()}" for text in texts] all_embeddings = [] total_batches = (len(prefixed_texts) + batch_size - 1) // batch_size for i in range(0, len(prefixed_texts), batch_size): batch_num = i // batch_size + 1 if not is_query and batch_num % 50 == 0: print(f"Processing batch {batch_num}/{total_batches} ({(batch_num/total_batches)*100:.1f}%)") if torch.cuda.is_available(): torch.cuda.empty_cache() batch_texts = prefixed_texts[i:i + batch_size] try: encoded = self.tokenizer( batch_texts, padding=True, truncation=True, max_length=512, # 512 return_tensors='pt' ).to(self.device) with torch.no_grad(): model_output = self.model(**encoded) # E5 uses mean pooling with attention mask attention_mask = encoded['attention_mask'] embeddings = model_output.last_hidden_state # Mean pooling mask_expanded = attention_mask.unsqueeze(-1).expand(embeddings.size()).float() sum_embeddings = torch.sum(embeddings * mask_expanded, 1) sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9) embeddings = sum_embeddings / sum_mask # embeddings = model_output[0][:, 0].float() # L2 normalize embeddings (important for E5) embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) # Move to CPU immediately all_embeddings.append(embeddings.cpu()) # Clear GPU memory del encoded, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except torch.cuda.OutOfMemoryError as e: print(f"CUDA OOM at batch {batch_num}, reducing batch size...") # Process one item at a time for single_text in batch_texts: try: encoded = self.tokenizer( [single_text], padding=True, truncation=True, max_length=512, return_tensors='pt' ).to(self.device) with torch.no_grad(): model_output = self.model(**encoded) attention_mask = encoded['attention_mask'] embeddings = model_output.last_hidden_state # Mean pooling mask_expanded = attention_mask.unsqueeze(-1).expand(embeddings.size()).float() sum_embeddings = torch.sum(embeddings * mask_expanded, 1) sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9) embeddings = sum_embeddings / sum_mask # embeddings = model_output[0][:, 0].float() embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) all_embeddings.append(embeddings.cpu()) del encoded, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except Exception as e2: print(f"Failed to process single text: {e2}") zero_embedding = torch.zeros(1, 1024).float() all_embeddings.append(zero_embedding) return torch.cat(all_embeddings, dim=0).numpy() def prepare_corpus(self, texts, ids, batch_size=32): prefixed_texts = [f"{text.strip()}" for text in texts] all_embeddings = [] all_tokens = [] all_ids = [] all_len = [] prefix_len = 1 suffix_len = 1 max_length = 512 # 1024 # 512 max_passage = max_length - prefix_len - suffix_len overlap = max_passage // 2 for i in range(len(prefixed_texts)): # tqdm( encoded = self.tokenizer( [prefixed_texts[i]], padding=False, truncation=False, # max_length=512, # return_tensors='pt' ) if len(encoded['input_ids'][0]) > max_length: _idxs = list(range(prefix_len, len(encoded['input_ids'][0])-suffix_len)) i0 = 0 while i0 < len(_idxs) - overlap: i1 = min(i0+max_passage, len(_idxs)) all_ids.append((ids[i], (i0, i1, len(_idxs)))) all_tokens.append({ 'input_ids': encoded['input_ids'][0][:prefix_len] + encoded['input_ids'][0][i0+prefix_len:i1+prefix_len] + encoded['input_ids'][0][-suffix_len:], 'attention_mask': encoded['attention_mask'][0][:prefix_len] + encoded['attention_mask'][0][i0+prefix_len:i1+prefix_len] + encoded['attention_mask'][0][-suffix_len:] }) all_len.append(i1-i0) i0 += max_passage - overlap else: all_ids.append((ids[i], None)) all_tokens.append({'input_ids': encoded['input_ids'][0], 'attention_mask': encoded['attention_mask'][0]}) all_len.append(len(encoded['input_ids'][0]) - prefix_len - suffix_len) total_batches = (len(all_tokens) + batch_size - 1) // batch_size all_len, all_tokens, all_ids = zip(*sorted(zip(all_len, all_tokens, all_ids), key=lambda x: x[0])) for i in range(0, len(all_tokens), batch_size): batch_num = i // batch_size + 1 if batch_num % 50 == 0: print(f"Processing batch {batch_num}/{total_batches} ({(batch_num/total_batches)*100:.1f}%)") if torch.cuda.is_available(): torch.cuda.empty_cache() batch_tokens = all_tokens[i:i + batch_size] batch_max = max([len(ids['input_ids']) for ids in batch_tokens]) encoded = dict() encoded["attention_mask"] = [s['attention_mask'] + (batch_max - len(s['attention_mask'])) * [0] for s in batch_tokens] encoded["input_ids"] = [s['input_ids'] + (batch_max - len(s['attention_mask'])) * [0] for s in batch_tokens] encoded["attention_mask"] = torch.tensor(encoded["attention_mask"], dtype=torch.long).to(self.device) encoded["input_ids"] = torch.tensor(encoded["input_ids"], dtype=torch.long).to(self.device) try: with torch.no_grad(): model_output = self.model(**encoded) # E5 uses mean pooling with attention mask attention_mask = encoded['attention_mask'] embeddings = model_output.last_hidden_state # Mean pooling mask_expanded = attention_mask.unsqueeze(-1).expand(embeddings.size()).float() sum_embeddings = torch.sum(embeddings * mask_expanded, 1) sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9) embeddings = sum_embeddings / sum_mask # L2 normalize embeddings embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) # Move to CPU immediately all_embeddings.append(embeddings.cpu()) # Clear GPU memory del encoded, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except torch.cuda.OutOfMemoryError as e: print(f"CUDA OOM at batch {batch_num}, reducing batch size...") # Process one item at a time for j in range(len(encoded)): try: encoded0 = encoded[j:j+1] with torch.no_grad(): model_output = self.model(**encoded0) attention_mask = encoded0['attention_mask'] embeddings = model_output.last_hidden_state # Mean pooling mask_expanded = attention_mask.unsqueeze(-1).expand(embeddings.size()).float() sum_embeddings = torch.sum(embeddings * mask_expanded, 1) sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9) embeddings = sum_embeddings / sum_mask embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) all_embeddings.append(embeddings.cpu()) del encoded0, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except Exception as e2: print(f"Failed to process single text: {e2}") zero_embedding = torch.zeros(1, 1024).float() all_embeddings.append(zero_embedding) id2idx = {} for i in range(len(all_ids)): if all_ids[i][0] not in id2idx: id2idx[all_ids[i][0]] = [] id2idx[all_ids[i][0]].append(i) return torch.cat(all_embeddings, dim=0).to(self.model.device), all_ids, id2idx # numpy() class M3Retriever: def __init__(self, model_name=None, device=None): """ Initializes the bge-m3 retriever using the multilingual bge-m3 model. """ # Use local model if model_name is None: # local_model_path = os.path.join('sub/models', 'bge-m3') local_model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models', 'bge-m3') if os.path.isdir(local_model_path): model_name = local_model_path print(f"Using local bge-m3 model from: {model_name}") self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") # Clear GPU cache before loading model if torch.cuda.is_available(): torch.cuda.empty_cache() print(f"Loading bge-m3 multilingual model on device: {self.device}") self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.model = AutoModel.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(self.device) # torch.float16 self.model.eval() # Clear cache after model loading if torch.cuda.is_available(): torch.cuda.empty_cache() self.corpus_ids = [] self.corpus_embeddings = None def embed_texts(self, texts, is_query=False, batch_size=32): """ Generates embeddings for texts using E5 Instruct model with proper prefixes. bge-m3 requires specific prefixes for queries vs passages. """ prefixed_texts = [f"{text.strip()}" for text in texts] all_embeddings = [] total_batches = (len(prefixed_texts) + batch_size - 1) // batch_size for i in range(0, len(prefixed_texts), batch_size): batch_num = i // batch_size + 1 if not is_query and batch_num % 50 == 0: print(f"Processing batch {batch_num}/{total_batches} ({(batch_num/total_batches)*100:.1f}%)") if torch.cuda.is_available(): torch.cuda.empty_cache() batch_texts = prefixed_texts[i:i + batch_size] try: encoded = self.tokenizer( batch_texts, padding=True, truncation=True, max_length=512, return_tensors='pt' ).to(self.device) with torch.no_grad(): model_output = self.model(**encoded) embeddings = model_output[0][:, 0].float() # L2 normalize embeddings embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) # Move to CPU immediately all_embeddings.append(embeddings.cpu()) # Clear GPU memory del encoded, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except torch.cuda.OutOfMemoryError as e: print(f"CUDA OOM at batch {batch_num}, reducing batch size...") # Process one item at a time for single_text in batch_texts: try: encoded = self.tokenizer( [single_text], padding=True, truncation=True, max_length=512, return_tensors='pt' ).to(self.device) with torch.no_grad(): model_output = self.model(**encoded) embeddings = model_output[0][:, 0].float() embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) all_embeddings.append(embeddings.cpu()) del encoded, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except Exception as e2: print(f"Failed to process single text: {e2}") zero_embedding = torch.zeros(1, 1024).float() all_embeddings.append(zero_embedding) return torch.cat(all_embeddings, dim=0).numpy() def prepare_corpus(self, texts, ids, batch_size=32): prefixed_texts = [f"{text.strip()}" for text in texts] all_embeddings = [] all_tokens = [] all_ids = [] all_len = [] prefix_len = 1 suffix_len = 1 max_length = 512 max_passage = max_length - prefix_len - suffix_len overlap = max_passage // 2 for i in range(len(prefixed_texts)): # tqdm( encoded = self.tokenizer( [prefixed_texts[i]], padding=False, truncation=False, # max_length=512, # return_tensors='pt' ) if len(encoded['input_ids'][0]) > max_length: _idxs = list(range(prefix_len, len(encoded['input_ids'][0])-suffix_len)) i0 = 0 while i0 < len(_idxs) - overlap: i1 = min(i0+max_passage, len(_idxs)) all_ids.append((ids[i], (i0, i1, len(_idxs)))) all_tokens.append({ 'input_ids': encoded['input_ids'][0][:prefix_len] + encoded['input_ids'][0][i0+prefix_len:i1+prefix_len] + encoded['input_ids'][0][-suffix_len:], 'attention_mask': encoded['attention_mask'][0][:prefix_len] + encoded['attention_mask'][0][i0+prefix_len:i1+prefix_len] + encoded['attention_mask'][0][-suffix_len:] }) all_len.append(i1-i0) i0 += max_passage - overlap else: all_ids.append((ids[i], None)) all_tokens.append({'input_ids': encoded['input_ids'][0], 'attention_mask': encoded['attention_mask'][0]}) all_len.append(len(encoded['input_ids'][0]) - prefix_len - suffix_len) total_batches = (len(all_tokens) + batch_size - 1) // batch_size all_len, all_tokens, all_ids = zip(*sorted(zip(all_len, all_tokens, all_ids), key=lambda x: x[0])) for i in range(0, len(all_tokens), batch_size): batch_num = i // batch_size + 1 if batch_num % 50 == 0: print(f"Processing batch {batch_num}/{total_batches} ({(batch_num/total_batches)*100:.1f}%)") if torch.cuda.is_available(): torch.cuda.empty_cache() batch_tokens = all_tokens[i:i + batch_size] batch_max = max([len(ids['input_ids']) for ids in batch_tokens]) encoded = dict() encoded["attention_mask"] = [s['attention_mask'] + (batch_max - len(s['attention_mask'])) * [0] for s in batch_tokens] encoded["input_ids"] = [s['input_ids'] + (batch_max - len(s['attention_mask'])) * [0] for s in batch_tokens] encoded["attention_mask"] = torch.tensor(encoded["attention_mask"], dtype=torch.long).to(self.device) encoded["input_ids"] = torch.tensor(encoded["input_ids"], dtype=torch.long).to(self.device) try: with torch.no_grad(): model_output = self.model(**encoded) embeddings = model_output[0][:, 0].float() # L2 normalize embeddings (important for E5) embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) # Move to CPU immediately all_embeddings.append(embeddings.cpu()) # Clear GPU memory del encoded, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except torch.cuda.OutOfMemoryError as e: print(f"CUDA OOM at batch {batch_num}, reducing batch size...") # Process one item at a time for j in range(len(encoded)): try: encoded0 = encoded[j:j+1] with torch.no_grad(): model_output = self.model(**encoded0) embeddings = model_output[0][:, 0].float() embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) all_embeddings.append(embeddings.cpu()) del encoded0, model_output, embeddings if torch.cuda.is_available(): torch.cuda.empty_cache() except Exception as e2: print(f"Failed to process single text: {e2}") zero_embedding = torch.zeros(1, 1024).float() all_embeddings.append(zero_embedding) id2idx = {} for i in range(len(all_ids)): if all_ids[i][0] not in id2idx: id2idx[all_ids[i][0]] = [] id2idx[all_ids[i][0]].append(i) return torch.cat(all_embeddings, dim=0).to(self.model.device), all_ids, id2idx # numpy() class BGEReranker: def __init__(self, model_name=None, device=None): """ Initializes the BGE reranker for fine-grained relevance scoring. """ # Use local model if model_name is None: # local_model_path = os.path.join('sub/models', 'bge-reranker-v2-m3') # 'bge-reranker-v2-m3' local_model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models', 'bge-reranker-v2-m3_pseudo_tune_full') # 'bge-reranker-v2-m3' if os.path.isdir(local_model_path): model_name = local_model_path print(f"Using local BGE model from: {model_name}") self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") print(f"Loading BGE reranker on device: {self.device}") # BGE reranker is actually a special model type from transformers import AutoModelForSequenceClassification self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.model = AutoModelForSequenceClassification.from_pretrained( model_name, torch_dtype=torch.bfloat16, trust_remote_code=True ).to(self.device) self.model.eval() if torch.cuda.is_available(): torch.cuda.empty_cache() def prepare_corpus(self, texts, ids): self.corpus_tokens = {} for i in range(len(texts)): # encoded = self.tokenizer( [texts[i].strip()], padding=False, truncation=False, # max_length=512, # return_tensors='pt' ) self.corpus_tokens[ids[i]] = {'input_ids': encoded['input_ids'][0], 'attention_mask': encoded['attention_mask'][0]} def rerank(self, query_text, passages, passage_ids, max_time, top_k=20): """ Rerank the passages using BGE reranker - CORRECTED VERSION. """ if not passages: return [] if time.time() > max_time-0.1: print(f'No time for bge rerank!!') return [] max_length = 2048 query_encoded = self.tokenizer( [query_text.strip()], padding=False, truncation=True, max_length = max_length - 1024, ) all_tokens = [] all_ids = [] all_len = [] suffix_len = 1 prefix_len = 1 max_passage = max_length - len(query_encoded['input_ids'][0]) - suffix_len - prefix_len overlap = max_passage // 2 for i in range(len(passage_ids)): # tqdm( encoded = self.corpus_tokens[passage_ids[i]] if len(encoded['input_ids']) - suffix_len - prefix_len > max_passage: _idxs = list(range(prefix_len, len(encoded['input_ids'])-suffix_len)) i0 = 0 while i0 < len(_idxs) - overlap: i1 = min(i0+max_passage, len(_idxs)) all_ids.append((passage_ids[i], (i0, i1, len(_idxs)))) all_tokens.append({ 'input_ids': query_encoded['input_ids'][0] + [2] + encoded['input_ids'][i0+prefix_len:i1+prefix_len] + [2], 'attention_mask': query_encoded['attention_mask'][0] + [1] + encoded['attention_mask'][i0+prefix_len:i1+prefix_len] + [1] }) all_len.append(i1-i0) i0 += max_passage - overlap else: all_ids.append((passage_ids[i], None)) all_tokens.append({'input_ids': query_encoded['input_ids'][0] + [2] + encoded['input_ids'][1:], 'attention_mask': query_encoded['attention_mask'][0] + [1] + encoded['attention_mask'][1:]}) all_len.append(len(encoded['input_ids']) - prefix_len - suffix_len) # all_len, all_tokens, all_ids = zip(*sorted(zip(all_len, all_tokens, all_ids), key=lambda x: x[0])) scores = [] batch_size = 4 # 4 # Conservative batch size for i in range(0, len(all_tokens), batch_size): if time.time() > max_time: print(f'bge rerank time limit! processed {len(scores)} of {len(all_tokens)}') break batch_tokens = all_tokens[i:i + batch_size] batch_max = max([len(ids['input_ids']) for ids in batch_tokens]) encoded = dict() encoded["attention_mask"] = [s['attention_mask'] + (batch_max - len(s['attention_mask'])) * [0] for s in batch_tokens] encoded["input_ids"] = [s['input_ids'] + (batch_max - len(s['attention_mask'])) * [0] for s in batch_tokens] encoded["attention_mask"] = torch.tensor(encoded["attention_mask"], dtype=torch.long).to(self.device) encoded["input_ids"] = torch.tensor(encoded["input_ids"], dtype=torch.long).to(self.device) try: # BGE reranker expects SEPARATE query and passage inputs # NOT concatenated strings # batch_queries = [query_text] * len(batch_passages) # Tokenize query-passage pairs properly with torch.no_grad(): # Get relevance scores from sequence classification model outputs = self.model(**encoded) # BGE reranker outputs logits for relevance classification logits = outputs.logits.float() # Handle different output shapes if len(logits.shape) == 1: # Single score per pair batch_scores = logits.cpu().numpy() elif logits.shape[1] == 1: # Single column output batch_scores = logits.squeeze(-1).cpu().numpy() else: # Binary classification - take positive class (index 1) batch_scores = logits[:, 1].cpu().numpy() scores.extend(batch_scores.tolist()) # Cleanup del encoded, outputs if torch.cuda.is_available(): torch.cuda.empty_cache() except Exception as e: print(f"Error in reranking batch {i//batch_size + 1}: {e}") # Fallback: Use neutral scores for this batch fallback_scores = [0.5] * len(batch_tokens) scores.extend(fallback_scores) all_ids = all_ids[:len(scores)] # Combine results and sort by reranking score results = list(zip(all_ids, scores)) results.sort(key=lambda x: x[1], reverse=True) new_res = [] _used = set([]) for i in range(len(results)): if results[i][0][0] not in _used: _used.add(results[i][0][0]) new_res.append((results[i][0][0], results[i][1])) return new_res[:min(top_k, len(new_res))] # Global instances retrievers = None reranker = None retrieverBM5 = None corpus_texts = {} # Store original passage texts for reranking def preprocess(corpus_dict): """ Preprocessing function using E5 multilingual model + BGE reranker. Input: corpus_dict - dict mapping document IDs to document objects with 'passage'/'text' field Output: dict containing initialized models, embeddings, and corpus data Note: Uses global variables (retriever, reranker, corpus_texts) for efficiency, but also returns all required data via preprocessed_data for function interface. """ global retrievers, reranker, corpus_texts, retrieverBM5 start_time = time.time() print("=" * 60) print("PREPROCESSING: Initializing E5 + BGE Reranker Pipeline...") print("=" * 60) # Set GPU memory optimization os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True' torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True # Initialize E5 retriever print("Loading retrievers...") retrieverE5 = E5Retriever() retrieverE5Ins = E5InstructRetriever() retrieverM3 = M3Retriever() retrieverSnowflake = SnowflakeInstructRetriever() retrieverSolon = SolonRetriever() retrieverRAGbot = E5Retriever(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models', 'Webiks_Hebrew_RAGbot_KolZchut_QA_Embedder_v1.0') ) retrievers = [retrieverE5, retrieverE5Ins, retrieverM3, retrieverSnowflake, retrieverSolon, retrieverRAGbot] # Initialize BGE reranker print("Loading rerankers...") reranker = BGEReranker() print(f"Preparing corpus with {len(corpus_dict)} documents...") # Store corpus IDs, passages, and original texts corpus_ids = list(corpus_dict.keys()) passages = [doc.get('passage', doc.get('text', '')) for doc in corpus_dict.values()] for ret in retrievers: ret.corpus_ids = list(corpus_dict.keys()) print("Computing embeddings...") ret.corpus_embeddings, ret.corpus_ids, ret.id2idx = ret.prepare_corpus(passages, ret.corpus_ids, batch_size=32) print("✓ Corpus preprocessing complete!") print(f"✓ Generated embeddings for {len(ret.corpus_ids)} documents") print(f"✓ Embedding matrix shape: {ret.corpus_embeddings.shape}") # Store original texts for reranking corpus_texts = {doc_id: passages[i] for i, doc_id in enumerate(corpus_ids)} reranker.corpus_ids = list(corpus_dict.keys()) reranker.prepare_corpus(passages, reranker.corpus_ids) retrieverBM5 = BM25sRetriever(passages, corpus_ids) if torch.cuda.is_available(): torch.cuda.empty_cache() end_time = time.time() elapsed_time = end_time - start_time print(f"Elapsed time: {elapsed_time} seconds") return { 'retrievers': retrievers, 'retrieverBM5': retrieverBM5, 'reranker': reranker, 'corpus_ids': corpus_ids, 'corpus_texts': corpus_texts, 'num_documents': len(corpus_dict) } def predict(query, preprocessed_data): """ Two-stage prediction: E5 retrieval + BGE reranking. Input: - query: dict with 'query' field containing query text - preprocessed_data: dict from preprocess() containing models and corpus data Output: list of dicts with 'paragraph_uuid' and 'score' fields, ranked by relevance Note: Uses global variables for efficiency but can also extract required data from preprocessed_data parameter for proper function interface. """ global retrievers, reranker, corpus_texts, retrieverBM5 start_time = time.time() max_time = start_time + 1.85 # # Extract query text query_text = query.get('query', '') if not query_text: return [] # Use global instances or get from preprocessed_data if retrievers is None: retrievers = preprocessed_data.get('retrievers') reranker = preprocessed_data.get('reranker') corpus_texts = preprocessed_data.get('corpus_texts', {}) retrieverBM5 = preprocessed_data.get('retrieverBM5') if retrievers is None or reranker is None: print("Error: Missing retriever or reranker in preprocessed data") return [] TopNCandidates0 = 250 TopNCandidates = 250 retriever_mean_std = [(-1.4131863134104607, 1.055066495990117), (0.8814187049865723, 0.017973395064473152), (0.5294753313064575, 0.06463246047496796), (0.4171469509601593, 0.0699099600315094), (0.49514809250831604, 0.06277099251747131), (0.530211865901947, 0.0704670324921608)] retriever_w = np.array([1.1, 0.25, 0.2, 0.3, 0.3, 0.3]) retriever_w /= retriever_w.sum() try: # STAGE 1: Retrieval (get top 100 candidates) candidate_ids = [] candidate_scores0 = [] candidate_passages = [] candidate_scores_bm5 = [] ret_scores = [] for retriever in retrievers: query_embedding = retriever.embed_texts([query_text], is_query=True, batch_size=1) query_embedding = torch.from_numpy(query_embedding).cuda() # Compute cosine similarity with precomputed corpus embeddings _scores = F.cosine_similarity(query_embedding, retriever.corpus_embeddings, 1, 1e-6) _scores = _scores.cpu().numpy() ret_scores.append(_scores) _argsort = np.argsort(_scores)[::-1] for idx in _argsort[:TopNCandidates0]: if retriever.corpus_ids[idx][0] not in candidate_ids: candidate_ids.append(retriever.corpus_ids[idx][0]) _scores = retrieverBM5.get_scores(query_text) _argsort = np.argsort(_scores)[::-1] for idx in _argsort[:TopNCandidates0]: if retrieverBM5.corpus_ids[idx] not in candidate_ids: candidate_ids.append(retrieverBM5.corpus_ids[idx]) for i in range(len(candidate_ids)): doc_id = candidate_ids[i] _bm5_sc = (_scores[retrieverBM5.id2idx[doc_id]] - 5.919949) / 2.7885008 candidate_scores_bm5.append(_bm5_sc) _sc = 0 for j in range(len(retrievers)): _sc0 = np.max(ret_scores[j][retrievers[j].id2idx[doc_id]]) _sc0 = (_sc0 - retriever_mean_std[j][0]) / retriever_mean_std[j][1] _sc += retriever_w[j] * _sc0 candidate_scores0.append(0.85 * _sc + 0.15 * _bm5_sc) candidate_passages.append(corpus_texts.get(doc_id, '')) candidate_scores, candidate_ids, candidate_passages, candidate_scores_bm5 = zip(*sorted(zip(candidate_scores0, candidate_ids, candidate_passages, candidate_scores_bm5))[::-1][:min(len(candidate_ids), TopNCandidates)]) # STAGE 2: Reranking (rerank top 100 -> top 20) reranked_results = reranker.rerank( query_text, candidate_passages, candidate_ids, max_time, top_k=len(candidate_ids) ) if len(reranked_results) >= 20: reranked_scores0 = [rerank_score for (passage_id, rerank_score) in reranked_results] reranked_results = [passage_id for (passage_id, rerank_score) in reranked_results] scores = [] for _i, _p in enumerate(reranked_results): scores.append(0.35 * (reranked_scores0[_i] + 1.8669906079283858) / 1.207719509974457 + 0.65 * candidate_scores[candidate_ids.index(_p)]) reranked_results = [(rerank_score, passage_id) for rerank_score, passage_id in sorted(zip(scores, reranked_results))][::-1] #[:TopNCandidates2] else: reranked_results = [(candidate_scores[i], candidate_ids[i]) for i in range(20)] # Build final results with ACTUAL reranking scores results = [] for rank, (rerank_score, passage_id) in enumerate(reranked_results[:20]): results.append({ 'paragraph_uuid': passage_id, 'score': float(rerank_score) # Use actual BGE reranker score! }) end_time = time.time() elapsed_time = end_time - start_time # print(f"✓ Returned {len(results)} results with reranker scores. Elapsed time: {elapsed_time} seconds") return results except Exception as e: print(f"Error in prediction: {e}") # Fallback to E5-only retrieval with E5 scores try: query_embedding = retrievers[0].embed_texts([query_text], is_query=True, batch_size=1) query_embedding = torch.from_numpy(query_embedding).cuda() e5_scores = F.cosine_similarity(query_embedding, retrievers[0].corpus_embeddings, 1, 1e-6) e5_scores = e5_scores.cpu().numpy() top_indices = np.argsort(e5_scores)[::-1][:20] results = [] for idx in top_indices: results.append({ 'paragraph_uuid': retriever.corpus_ids[idx], 'score': float(e5_scores[idx]) # Use actual E5 cosine similarity score }) return results except: return []