add_mrr@10

app.py

@@ -13,6 +13,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 zero = torch.Tensor([0]).to(device)
 print(f"Device being used: {zero.device}")
 
+
 @spaces.GPU
 def evaluate_model(model_id, num_questions):
     model = SentenceTransformer(model_id, device=device)
@@ -44,7 +45,7 @@ def evaluate_model(model_id, num_questions):
             "last_rows": True  # Take the last num_questions rows
         }
     ]
-    
+
     evaluation_results = []
     scores_by_dataset = {}
 
@@ -57,25 +58,26 @@ def evaluate_model(model_id, num_questions):
 
         # Select the required number of rows
         if dataset_info.get("last_rows"):
-            dataset = dataset.select(range(len(dataset) - dataset_info["sample_size"], len(dataset)))  # Take last n rows
+            dataset = dataset.select(
+                range(len(dataset) - dataset_info["sample_size"], len(dataset)))  # Take last n rows
         else:
             dataset = dataset.select(range(min(dataset_info["sample_size"], len(dataset))))  # Take first n rows
-        
+
         # Rename columns to 'anchor' and 'positive'
         dataset = dataset.rename_column(dataset_info["columns"][0], "anchor")
         dataset = dataset.rename_column(dataset_info["columns"][1], "positive")
-        
+
         # Check if "id" column already exists before adding it
         if "id" not in dataset.column_names:
             dataset = dataset.add_column("id", range(len(dataset)))
-        
+
         # Prepare queries and corpus
         corpus = dict(zip(dataset["id"], dataset["positive"]))
         queries = dict(zip(dataset["id"], dataset["anchor"]))
-        
+
         # Create a mapping of relevant documents (1 in our case) for each query
         relevant_docs = {q_id: [q_id] for q_id in queries}
-        
+
         matryoshka_evaluators = []
         for dim in matryoshka_dimensions:
             ir_evaluator = InformationRetrievalEvaluator(
@@ -84,66 +86,91 @@ def evaluate_model(model_id, num_questions):
                 relevant_docs=relevant_docs,
                 name=f"dim_{dim}",
                 truncate_dim=dim,
-                score_functions={"cosine": cos_sim},
+                score_functions={"cosine": cos_sim}
             )
             matryoshka_evaluators.append(ir_evaluator)
 
         evaluator = SequentialEvaluator(matryoshka_evaluators)
         results = evaluator(model)
-        
-        scores = []
+
+        scores_ndcg = []
+        scores_mrr = []
         for dim in matryoshka_dimensions:
-            key = f"dim_{dim}_cosine_ndcg@10"
-            score = results[key] if key in results else None
+            ndcg_key = f"dim_{dim}_cosine_ndcg@10"
+            mrr_key = f"dim_{dim}_cosine_mrr@10"
+            ndcg_score = results[ndcg_key] if ndcg_key in results else None
+            mrr_score = results[mrr_key] if mrr_key in results else None
             evaluation_results.append({
                 "Dataset": dataset_info["name"],
                 "Dimension": dim,
-                "Score": score
+                "NDCG@10": ndcg_score,
+                "MRR@10": mrr_score
             })
-            scores.append(score)
+            scores_ndcg.append(ndcg_score)
+            scores_mrr.append(mrr_score)
 
         # Store scores by dataset for plot creation
-        scores_by_dataset[dataset_info["name"]] = scores
+        scores_by_dataset[dataset_info["name"]] = {
+            "NDCG@10": scores_ndcg,
+            "MRR@10": scores_mrr
+        }
 
     # Convert results to DataFrame for display
     result_df = pd.DataFrame(evaluation_results)
 
     # Generate bar charts for each dataset using Plotly
     charts = []
-    color_scale = ['#003f5c', '#2f4b7c', '#665191', '#a05195', '#d45087']
+    color_scale_ndcg = '#a05195'
+    color_scale_mrr = '#2f4b7c'
 
     for dataset_name, scores in scores_by_dataset.items():
         fig = go.Figure()
+        # NDCG@10 bars
         fig.add_trace(go.Bar(
             x=[str(dim) for dim in matryoshka_dimensions],
-            y=scores,
-            marker_color=color_scale,
-            text=[f"{score:.3f}" if score else "N/A" for score in scores],
+            y=scores["NDCG@10"],
+            name="NDCG@10",
+            marker_color=color_scale_ndcg,
+            text=[f"{score:.3f}" if score else "N/A" for score in scores["NDCG@10"]],
+            textposition='auto'
+        ))
+
+        # MRR@10 bars
+        fig.add_trace(go.Bar(
+            x=[str(dim) for dim in matryoshka_dimensions],
+            y=scores["MRR@10"],
+            name="MRR@10",
+            marker_color=color_scale_mrr,
+            text=[f"{score:.3f}" if score else "N/A" for score in scores["MRR@10"]],
             textposition='auto'
         ))
 
         fig.update_layout(
             title=f"{dataset_name} Evaluation",
             xaxis_title="Embedding Dimension",
-            yaxis_title="NDCG@10 Score",
+            yaxis_title="Score",
+            barmode='group',  # Group bars
             template="plotly_white"
         )
         charts.append(fig)
 
     return result_df, charts[0], charts[1], charts[2]
 
+
 # Define the Gradio interface
 def display_results(model_name, num_questions):
     result_df, chart1, chart2, chart3 = evaluate_model(model_name, num_questions)
     return result_df, chart1, chart2, chart3
 
+
 # Gradio interface with a slider to choose the number of questions (1 to 500)
 demo = gr.Interface(
-    fn=display_results, 
+    fn=display_results,
     inputs=[
-        gr.Textbox(label="Enter a Hugging Face Model ID", placeholder="e.g., Omartificial-Intelligence-Space/GATE-AraBert-v1"),
+        gr.Textbox(label="Enter a Hugging Face Model ID",
+                   placeholder="e.g., Omartificial-Intelligence-Space/GATE-AraBert-v1"),
         gr.Slider(label="Number of Questions", minimum=1, maximum=500, step=1, value=500)
-    ], 
+    ],
     outputs=[
         gr.Dataframe(label="Evaluation Results"),
         gr.Plot(label="Financial Dataset"),
@@ -156,8 +183,8 @@ demo = gr.Interface(
         "- **ARCD** evaluates short context retrieval performance.\n"
         "- **MLQA Arabic** evaluates long context retrieval performance.\n"
         "- **Arabic Financial Dataset** focuses on financial context retrieval.\n\n"
-        "**Evaluation Metric:**\n"
-        "The evaluation uses **NDCG@10** (Normalized Discounted Cumulative Gain), which measures how well the retrieved documents (contexts) match the query relevance.\n"
+        "**Evaluation Metrics:**\n"
+        "The evaluation uses **NDCG@10** and **MRR@10**, which measure how well the retrieved documents (contexts) match the query relevance.\n"
         "Higher scores indicate better performance. Embedding dimensions are reduced from 768 to 64, evaluating how well the model performs with fewer dimensions."
     ),
     theme="default",