speed

app.py

@@ -42,6 +42,7 @@ qrels = {}
 query2qid = {}
 datasets = ["scifact"]
 current_dataset = "scifact"
+faiss_index = None
 
 def log_system_info():
     logger.info("System Information:")
@@ -125,7 +126,7 @@ class RepLlamaModel:
         model.eval()
         return model
 
-    def encode(self, texts, batch_size=16, **kwargs):
+    def encode(self, texts, batch_size=32, **kwargs):
         self.model = self.model.cuda()
         all_embeddings = []
         for i in range(0, len(texts), batch_size):
@@ -183,8 +184,7 @@ def initialize_faiss_and_corpus(dataset_name):
     return index
 
 def search_queries(dataset_name, q_reps, depth=1000):
-    faiss_index = initialize_faiss_and_corpus(dataset_name)
-    
+    global faiss_index    
     logger.info(f"Searching queries. Shape of q_reps: {q_reps.shape}")
     
     # Perform the search
@@ -235,6 +235,9 @@ def evaluate(qrels, results, k_values):
         logger.info(f"NDCG@{k}: mean={metrics[f'NDCG@{k}']}, min={min(ndcg_scores)}, max={max(ndcg_scores)}")
         logger.info(f"Recall@{k}: mean={metrics[f'Recall@{k}']}, min={min(recall_scores)}, max={max(recall_scores)}")
 
+    # delete nDCG@100 and Recall@10
+    del metrics["NDCG@100"]
+    del metrics["Recall@100"]
     return metrics
 
 @spaces.GPU
@@ -298,6 +301,7 @@ def gradio_interface(dataset, postfix):
 if model is None:
     model = RepLlamaModel(model_name_or_path=CUR_MODEL)
     load_queries(current_dataset)   
+    faiss_index = initialize_faiss_and_corpus(current_dataset)
 
 # Create Gradio interface
 iface = gr.Interface(