import gradio as gr import spaces import torch import pandas as pd import plotly.graph_objects as go from datasets import load_dataset from sentence_transformers import SentenceTransformer from sentence_transformers.evaluation import InformationRetrievalEvaluator, SequentialEvaluator from sentence_transformers.util import cos_sim # Check for GPU support and configure appropriately 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) matryoshka_dimensions = [768, 512, 256, 128, 64] # Prepare datasets (Load entire split, then select num_questions) datasets_info = [ { "name": "Financial", "dataset_id": "Omartificial-Intelligence-Space/Arabic-finanical-rag-embedding-dataset", "split": "train", # Only train split "columns": ("question", "context"), "sample_size": num_questions }, { "name": "MLQA", "dataset_id": "google/xtreme", "subset": "MLQA.ar.ar", "split": "validation", # Only validation split "columns": ("question", "context"), "sample_size": num_questions }, { "name": "ARCD", "dataset_id": "hsseinmz/arcd", "split": "train", # Only train split "columns": ("question", "context"), "sample_size": num_questions, "last_rows": True # Take the last num_questions rows } ] evaluation_results = [] scores_by_dataset = {} for dataset_info in datasets_info: # Load the full dataset split and limit it afterward 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"]) # 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 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( queries=queries, corpus=corpus, relevant_docs=relevant_docs, name=f"dim_{dim}", truncate_dim=dim, score_functions={"cosine": cos_sim} ) matryoshka_evaluators.append(ir_evaluator) evaluator = SequentialEvaluator(matryoshka_evaluators) results = evaluator(model) scores_ndcg = [] scores_mrr = [] for dim in matryoshka_dimensions: 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, "NDCG@10": ndcg_score, "MRR@10": mrr_score }) scores_ndcg.append(ndcg_score) scores_mrr.append(mrr_score) # Store scores by dataset for plot creation 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_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["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="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, inputs=[ 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"), gr.Plot(label="MLQA Dataset"), gr.Plot(label="ARCD Dataset") ], title="MERAA : Matryoshka Embedding Retrieval Assessment for Arabic", description=( "Evaluate your Embedding model or any Arabic Sentence Transformer model's performance on **context and question retrieval** for Arabic datasets for Enhancing RAG (Retrieval-Augmented Generation).\n" "- **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 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", live=False, css="footer {visibility: hidden;}" ) demo.launch(debug=True) # Add the footer print("\nCreated by Omar Najar | Omartificial Intelligence Space")