load faiss

app.py

@@ -7,7 +7,6 @@ import torch
 import torch.nn.functional as F
 from transformers import AutoTokenizer, AutoModel, set_seed
 from peft import PeftModel
-from tevatron.retriever.searcher import FaissFlatSearcher
 import logging
 import os
 import json
@@ -47,7 +46,6 @@ current_dataset = "scifact"
 def log_system_info():
     logger.info("System Information:")
     logger.info(f"Python version: {sys.version}")
-    # logger.info(f"Platform: {platform.platform()}")
     
     logger.info("\nPackage Versions:")
     logger.info(f"torch: {torch.__version__}")
@@ -55,7 +53,6 @@ def log_system_info():
     logger.info(f"peft: {peft.__version__}")
     logger.info(f"faiss: {faiss.__version__}")
     logger.info(f"gradio: {gr.__version__}")
-    # logger.info(f"pytrec_eval: {pytrec_eval.__version__}")
     logger.info(f"ir_datasets: {ir_datasets.__version__}")
     
     if torch.cuda.is_available():
@@ -70,11 +67,8 @@ def log_system_info():
         logger.info("\nCUDA Information:")
         logger.info("CUDA available: No")
 
-
 log_system_info()
 
-
-
 def pool(last_hidden_states, attention_mask, pool_type="last"):
     last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
 
@@ -151,18 +145,45 @@ class RepLlamaModel:
         self.model = self.model.cpu()
         return np.concatenate(all_embeddings, axis=0)
 
+def load_corpus_embeddings(dataset_name):
+    corpus_path = f"{dataset_name}/corpus_emb.*.pkl"
+    index_files = glob.glob(corpus_path)
+    index_files.sort(key=lambda x: int(x.split('.')[-2]))
+    
+    all_embeddings = []
+    corpus_lookups = []
+    
+    for file in index_files:
+        with open(file, 'rb') as f:
+            embeddings, p_lookup = pickle.load(f)
+        all_embeddings.append(embeddings)
+        corpus_lookups.extend(p_lookup)
+    
+    all_embeddings = np.concatenate(all_embeddings, axis=0)
+    logger.info(f"Loaded corpus embeddings for {dataset_name}. Shape: {all_embeddings.shape}")
+    
+    return all_embeddings, corpus_lookups
+
+def create_faiss_index(embeddings):
+    dimension = embeddings.shape[1]
+    index = faiss.IndexFlatIP(dimension)
+    index.add(embeddings)
+    logger.info(f"Created FAISS index with {index.ntotal} vectors of dimension {dimension}")
+    return index
+
+def load_or_create_faiss_index(dataset_name):
+    embeddings, corpus_lookups = load_corpus_embeddings(dataset_name)
+    index = create_faiss_index(embeddings)
+    return index, corpus_lookups
 
-def load_faiss_index(dataset_name):
-    index_path = f"{dataset_name}/faiss_index.bin"
-    if os.path.exists(index_path):
-        logger.info(f"Loading existing FAISS index for {dataset_name} from {index_path}")
-        return faiss.read_index(index_path)
-    return None
+def initialize_faiss_and_corpus(dataset_name):
+    global corpus_lookups
+    index, corpus_lookups[dataset_name] = load_or_create_faiss_index(dataset_name)
+    logger.info(f"Initialized FAISS index and corpus lookups for {dataset_name}")
+    return index
 
 def search_queries(dataset_name, q_reps, depth=1000):
-    faiss_index = load_faiss_index(dataset_name)
-    if faiss_index is None:
-        raise ValueError(f"No FAISS index found for dataset {dataset_name}")
+    faiss_index = initialize_faiss_and_corpus(dataset_name)
     
     logger.info(f"Searching queries. Shape of q_reps: {q_reps.shape}")
     
@@ -171,28 +192,11 @@ def search_queries(dataset_name, q_reps, depth=1000):
     
     logger.info(f"Search completed. Shape of all_scores: {all_scores.shape}, all_indices: {all_indices.shape}")
     logger.info(f"Sample scores: {all_scores[0][:5]}, Sample indices: {all_indices[0][:5]}")
-
     
     psg_indices = [[str(corpus_lookups[dataset_name][x]) for x in q_dd] for q_dd in all_indices]
     
     return all_scores, np.array(psg_indices)
 
-def load_corpus_lookups(dataset_name):
-    global corpus_lookups
-    corpus_path = f"{dataset_name}/corpus_emb.*.pkl"
-    index_files = glob.glob(corpus_path)
-    # sort them
-    index_files.sort(key=lambda x: int(x.split('.')[-2]))
-    
-    corpus_lookups[dataset_name] = []
-    for file in index_files:
-        with open(file, 'rb') as f:
-            _, p_lookup = pickle.load(f)
-        corpus_lookups[dataset_name] += p_lookup
-
-    logger.info(f"Loaded corpus lookups for {dataset_name}. Total entries: {len(corpus_lookups[dataset_name])}")
-    logger.info(f"Sample corpus lookup entry: {corpus_lookups[dataset_name][:10]}")
-
 def load_queries(dataset_name):
     global queries, q_lookups, qrels, query2qid
     dataset = ir_datasets.load(f"beir/{dataset_name.lower()}" + ("/test" if dataset_name == "scifact" else ""))
@@ -214,7 +218,6 @@ def load_queries(dataset_name):
     logger.info(f"Loaded queries for {dataset_name}. Total queries: {len(queries[dataset_name])}")
     logger.info(f"Loaded qrels for {dataset_name}. Total query IDs: {len(qrels[dataset_name])}")
 
-
 def evaluate(qrels, results, k_values):
     qrels = {str(k): {str(k2): v2 for k2, v2 in v.items()} for k, v in qrels.items()}
     results = {str(k): {str(k2): v2 for k2, v2 in v.items()} for k, v in results.items()}
@@ -273,7 +276,6 @@ def run_evaluation(dataset, postfix):
     logger.info(f"Number of results: {len(results)}")
     logger.info(f"Sample result: {list(results.items())[0]}")
     
-    # Add these lines
     logger.info(f"Number of queries in qrels: {len(qrels[dataset])}")
     logger.info(f"Sample qrel: {list(qrels[dataset].items())[0]}")
     logger.info(f"Number of queries in results: {len(results)}")
@@ -293,13 +295,10 @@ def run_evaluation(dataset, postfix):
 def gradio_interface(dataset, postfix):
     return run_evaluation(dataset, postfix)
 
-
 if model is None:
     model = RepLlamaModel(model_name_or_path=CUR_MODEL)
-    load_corpus_lookups(current_dataset)
     load_queries(current_dataset)   
 
-
 # Create Gradio interface
 iface = gr.Interface(
     fn=gradio_interface,
@@ -318,4 +317,4 @@ iface = gr.Interface(
 )
 
 # Launch the interface
-iface.launch()
+iface.launch(share=False)