try to fix ndcg bug

app.py

@@ -103,13 +103,13 @@ class RepLlamaModel:
             batch_dict = create_batch_dict(self.tokenizer, batch_texts, always_add_eos="last")
             batch_dict = {key: value.cuda() for key, value in batch_dict.items()}
 
-            with torch.cuda.amp.autocast():
-                with torch.no_grad():
-                    outputs = self.model(**batch_dict)
-                    embeddings = pool(outputs.last_hidden_state, batch_dict['attention_mask'], 'last')
-                    embeddings = F.normalize(embeddings, p=2, dim=-1)
-                    logger.info(f"Encoded shape: {embeddings.shape}, Norm of first embedding: {torch.norm(embeddings[0]).item()}")
-                    all_embeddings.append(embeddings.cpu().numpy())
+            # with torch.cuda.amp.autocast():
+            with torch.no_grad():
+                outputs = self.model(**batch_dict)
+                embeddings = pool(outputs.last_hidden_state, batch_dict['attention_mask'], 'last')
+                embeddings = F.normalize(embeddings, p=2, dim=-1)
+                logger.info(f"Encoded shape: {embeddings.shape}, Norm of first embedding: {torch.norm(embeddings[0]).item()}")
+                all_embeddings.append(embeddings.cpu().numpy())
 
         self.model = self.model.cpu()
         return np.concatenate(all_embeddings, axis=0)
@@ -150,6 +150,7 @@ def load_corpus_lookups(dataset_name):
             _, 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][0]}")
 
 def load_queries(dataset_name):
     global queries, q_lookups, qrels
@@ -181,6 +182,12 @@ def evaluate(qrels, results, k_values):
     for k in k_values:
         metrics[f"NDCG@{k}"] = round(np.mean([query_scores[f"ndcg_cut_{k}"] for query_scores in scores.values()]), 3)
         metrics[f"Recall@{k}"] = round(np.mean([query_scores[f"recall_{k}"] for query_scores in scores.values()]), 3)
+        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)}")
+
+    # Add these lines
+    logger.info(f"Number of queries evaluated: {len(scores)}")
+    logger.info(f"Sample evaluation score: {list(scores.items())[0]}")
 
     return metrics
 
@@ -190,19 +197,35 @@ def run_evaluation(dataset, postfix):
     current_dataset = dataset
 
     input_texts = [f"query: {query.strip()} {postfix}".strip() for query in queries[current_dataset]]
+    logger.info(f"Number of input texts: {len(input_texts)}")
+    logger.info(f"Sample input text: {input_texts[0]}")
+
     q_reps = model.encode(input_texts)
+    logger.info(f"Encoded query representations shape: {q_reps.shape}")
+
     all_scores, psg_indices = search_queries(dataset, q_reps)
     
     results = {qid: dict(zip(doc_ids, map(float, scores))) 
                for qid, scores, doc_ids in zip(q_lookups[dataset].keys(), all_scores, psg_indices)}
     
+    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)}")
+    logger.info(f"Sample result: {list(results.items())[0]}")
+    
+    # Check for mismatches
+    qrels_keys = set(qrels[dataset].keys())
+    results_keys = set(results.keys())
+    logger.info(f"Queries in qrels but not in results: {qrels_keys - results_keys}")
+    logger.info(f"Queries in results but not in qrels: {results_keys - qrels_keys}")
+    
     metrics = evaluate(qrels[dataset], results, k_values=[10, 100])
     
-    return {
-        "NDCG@10": metrics["NDCG@10"],
-        "Recall@100": metrics["Recall@100"]
-    }
-
+    return metrics
 
 @spaces.GPU
 def gradio_interface(dataset, postfix):