speed

app.py

@@ -126,7 +126,7 @@ class RepLlamaModel:
         model.eval()
         return model
 
-    def encode(self, texts, batch_size=32, **kwargs):
+    def encode(self, texts, batch_size=48, **kwargs):
         self.model = self.model.cuda()
         all_embeddings = []
         for i in range(0, len(texts), batch_size):
@@ -183,7 +183,7 @@ def initialize_faiss_and_corpus(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):
+def search_queries(dataset_name, q_reps, depth=100):
     global faiss_index    
     logger.info(f"Searching queries. Shape of q_reps: {q_reps.shape}")