sort?

app.py

@@ -39,6 +39,7 @@ corpus_lookups = {}
 queries = {}
 q_lookups = {}
 qrels = {}
+query2qid = {}
 datasets = ["scifact"]
 current_dataset = "scifact"
 
@@ -157,7 +158,7 @@ def load_faiss_index(dataset_name):
         return faiss.read_index(index_path)
     return None
 
-def search_queries(dataset_name, q_reps, depth=1000):
+def search_queries(dataset_name, q_reps, depth=100):
     faiss_index = load_faiss_index(dataset_name)
     if faiss_index is None:
         raise ValueError(f"No FAISS index found for dataset {dataset_name}")
@@ -169,6 +170,7 @@ 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]
     
@@ -191,15 +193,17 @@ def load_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
+    global queries, q_lookups, qrels, query2qid
     dataset = ir_datasets.load(f"beir/{dataset_name.lower()}" + ("/test" if dataset_name == "scifact" else ""))
     
     queries[dataset_name] = []
+    query2qid[dataset_name] = {}
     q_lookups[dataset_name] = {}
     qrels[dataset_name] = {}
     for query in dataset.queries_iter():
         queries[dataset_name].append(query.text)
         q_lookups[dataset_name][query.query_id] = query.text
+        query2qid[dataset_name][query.text] = query.query_id
     
     for qrel in dataset.qrels_iter():
         if qrel.query_id not in qrels[dataset_name]:
@@ -231,7 +235,7 @@ def evaluate(qrels, results, k_values):
 
 @spaces.GPU
 def run_evaluation(dataset, postfix):
-    global current_dataset, queries, model
+    global current_dataset, queries, model, query2qid
     current_dataset = dataset
 
     input_texts = [f"query: {query.strip()} {postfix}".strip() for query in queries[current_dataset]]
@@ -248,7 +252,8 @@ def run_evaluation(dataset, postfix):
     logging.info(f"Number of queries in q_lookups: {len(q_lookups[dataset])}")
     logging.info("Size of all_scores: " + str(len(all_scores)))
     logging.info("Size of psg_indices: " + str(len(psg_indices)))
-    for qid, scores, doc_ids in zip(q_lookups[dataset].keys(), all_scores, psg_indices):
+    for query, scores, doc_ids in zip(queries, all_scores, psg_indices):
+        qid = query2qid[dataset][query]
         qid_str = str(qid)
         results[qid_str] = {}
         for doc_id, score in zip(doc_ids, scores):