update app.py

app.py

@@ -1,25 +1,23 @@
 import gradio as gr
 import spaces
 import torch
+import pandas as pd
+import matplotlib.pyplot as plt
 from datasets import load_dataset, concatenate_datasets
 from sentence_transformers import SentenceTransformer
 from sentence_transformers.evaluation import InformationRetrievalEvaluator, SequentialEvaluator
 from sentence_transformers.util import cos_sim
-import pandas as pd
 
 # Check for GPU support and configure appropriately
 device = "cuda" if torch.cuda.is_available() else "cpu"
-zero = torch.Tensor([0]).to(device)  # Ensure that tensor operates on the appropriate device
+zero = torch.Tensor([0]).to(device)
 print(f"Device being used: {zero.device}")
 
-@spaces.GPU  # Enable GPU support for Gradio Spaces
+@spaces.GPU
 def evaluate_model(model_id):
-    # Load the model on the appropriate device
     model = SentenceTransformer(model_id, device=device)
-    
-    # Define the evaluation parameters
     matryoshka_dimensions = [768, 512, 256, 128, 64]
-    
+
     # Prepare datasets
     datasets_info = [
         {
@@ -28,38 +26,43 @@ def evaluate_model(model_id):
             "split": "train",
             "size": 7000,
             "columns": ("question", "context"),
-            "sample_size": 100
+            "sample_size": 500
         },
         {
-            "name": "MLQA Arabic",
+            "name": "MLQA Arabic (Long Context)",
             "dataset_id": "google/xtreme",
-            "subset": "MLQA.ar.ar",  # Specify the correct subset
+            "subset": "MLQA.ar.ar",
             "split": "validation",
             "size": 500,
             "columns": ("question", "context"),
-            "sample_size": 100
+            "sample_size": 500
         },
         {
-            "name": "Custom",
-            "dataset_id": "philschmid/finanical-rag-embedding-dataset",
+            "name": "ARCD (Short Context)",
+            "dataset_id": "hsseinmz/arcd",
             "split": "train",
             "size": None,
             "columns": ("question", "context"),
-            "sample_size": 100
+            "sample_size": 500,
+            "last_rows": True  # Flag to take the last 500 rows
         }
     ]
     
     evaluation_results = []
-    
+    scores_by_dataset = {}
+
     for dataset_info in datasets_info:
         # Load the dataset with subset if available
         if "subset" in dataset_info:
             dataset = load_dataset(dataset_info["dataset_id"], dataset_info["subset"], split=dataset_info["split"])
         else:
             dataset = load_dataset(dataset_info["dataset_id"], split=dataset_info["split"])
-        
-        # Sample the dataset
-        dataset = dataset.select(range(min(dataset_info["sample_size"], len(dataset))))
+
+        # Take last 500 rows if specified
+        if dataset_info.get("last_rows"):
+            dataset = dataset.select(range(len(dataset) - dataset_info["sample_size"], len(dataset)))
+        else:
+            dataset = dataset.select(range(min(dataset_info["sample_size"], len(dataset))))
         
         # Rename columns
         dataset = dataset.rename_column(dataset_info["columns"][0], "anchor")
@@ -77,7 +80,6 @@ def evaluate_model(model_id):
         relevant_docs = {q_id: [q_id] for q_id in queries}
         
         matryoshka_evaluators = []
-        # Iterate over the different dimensions
         for dim in matryoshka_dimensions:
             ir_evaluator = InformationRetrievalEvaluator(
                 queries=queries,
@@ -88,14 +90,11 @@ def evaluate_model(model_id):
                 score_functions={"cosine": cos_sim},
             )
             matryoshka_evaluators.append(ir_evaluator)
-        
-        # Create a sequential evaluator
+
         evaluator = SequentialEvaluator(matryoshka_evaluators)
-        
-        # Evaluate the model
         results = evaluator(model)
         
-        # Collect results for each dataset
+        scores = []
         for dim in matryoshka_dimensions:
             key = f"dim_{dim}_cosine_ndcg@10"
             score = results[key] if key in results else None
@@ -104,26 +103,52 @@ def evaluate_model(model_id):
                 "Dimension": dim,
                 "Score": score
             })
-    
+            scores.append(score)
+
+        # Store scores by dataset for bar chart creation
+        scores_by_dataset[dataset_info["name"]] = scores
+
     # Convert results to DataFrame for display
     result_df = pd.DataFrame(evaluation_results)
-    return result_df
+
+    # Generate bar charts for each dataset
+    charts = []
+    for dataset_name, scores in scores_by_dataset.items():
+        fig, ax = plt.subplots()
+        ax.bar([str(dim) for dim in matryoshka_dimensions], scores)
+        ax.set_title(f"{dataset_name} Evaluation Scores", fontsize=16, color='darkblue')
+        ax.set_xlabel("Embedding Dimension", fontsize=12)
+        ax.set_ylabel("NDCG@10 Score", fontsize=12)
+        plt.tight_layout()
+        charts.append(fig)
+
+    return result_df, charts
 
 # Define the Gradio interface
 def display_results(model_name):
-    # Evaluate model and return results
-    result_df = evaluate_model(model_name)
-    return result_df
+    result_df, charts = evaluate_model(model_name)
+    return result_df, charts[0], charts[1], charts[2]
 
-# Create the Gradio interface
 demo = gr.Interface(
     fn=display_results, 
-    inputs=gr.Textbox(label="Model ID"), 
-    outputs="dataframe",
-    title="Model Evaluation with GPU Support",
-    description="Enter a Hugging Face Sentence Transformer model ID to evaluate it across datasets, leveraging GPU if available."
+    inputs=gr.Textbox(label="Enter a Hugging Face Model ID", placeholder="e.g., sentence-transformers/all-MiniLM-L6-v2"), 
+    outputs=[
+        gr.Dataframe(label="Evaluation Results"),
+        gr.Plot(label="Arabic Financial Dataset (Financial Evaluation)"),
+        gr.Plot(label="MLQA Arabic (Long Context Evaluation)"),
+        gr.Plot(label="ARCD (Short Context Evaluation)")
+    ],
+    title="Sentence Transformer Evaluation Dashboard",
+    description=(
+        "Evaluate your Sentence Transformer model across three datasets:\n"
+        "- **Arabic Financial Dataset** for financial data evaluation.\n"
+        "- **MLQA Arabic** for long context question-answer evaluation.\n"
+        "- **ARCD** for short context question-answer evaluation.\n"
+        "Results are displayed in a table and visualized with bar charts for each dataset."
+    ),
+    theme="default",
+    live=False
 )
 
-# Launch the Gradio app
 if __name__ == "__main__":
     demo.launch()