| """General text quality on the full MTEB(eng, v2) benchmark (41 tasks). |
| |
| Give each model its own --output-dir, or local models will reuse each other's cached scores. |
| |
| uv run mteb_tasks.py jaisawhney/SonicEmbed --output-dir mteb-results/sonicembed |
| """ |
| import argparse |
| import csv |
| import json |
| from pathlib import Path |
| from typing import cast |
|
|
| import mteb |
| import torch |
| from mteb import BenchmarkResults |
| from sentence_transformers import SentenceTransformer |
| from sentence_transformers.base.modules.transformer import Transformer |
|
|
| BENCHMARK = "MTEB(eng, v2)" |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("model") |
| parser.add_argument("--batch-size", type=int, default=16) |
| parser.add_argument("--output-dir", default="mteb-results") |
| parser.add_argument("--device", default="cuda", help="cuda, mps or cpu") |
| args = parser.parse_args() |
|
|
| output_dir = Path(args.output_dir) |
| output_dir.mkdir(parents=True, exist_ok=True) |
|
|
| |
| model =SentenceTransformer( |
| args.model, |
| device=args.device, |
| model_kwargs={ |
| "torch_dtype": torch.bfloat16 if args.device == "cuda" else torch.float32, |
| }, |
| config_kwargs={ |
| "vision_config": None, |
| "audio_config": None, |
| }, |
| ) |
|
|
| |
| transformer = cast(Transformer, model[0]) |
| for modality in transformer.modalities: |
| if modality != "text": |
| del transformer.modality_config[modality] |
| transformer.input_formatter.supported_modalities = transformer.modalities |
|
|
| |
| if "SentenceSimilarity" in model.prompts: |
| model.prompts.setdefault( |
| "STS", |
| model.prompts["SentenceSimilarity"], |
| ) |
|
|
| benchmark = mteb.get_benchmark(BENCHMARK) |
|
|
| cache =mteb.ResultCache(str(output_dir / "cache")) |
|
|
| print(f"Evaluating {args.model} on {BENCHMARK}") |
| print(f"Tasks: {len(benchmark.tasks)}") |
|
|
| results = mteb.evaluate( |
| model, |
| tasks=benchmark, |
| encode_kwargs={ |
| "batch_size": args.batch_size, |
| }, |
| cache=cache, |
| overwrite_strategy="only-missing", |
| ) |
|
|
| scores ={ |
| result.task_name: 100 * result.get_score() |
| for result in results.task_results |
| } |
|
|
| summary =benchmark.get_score( |
| BenchmarkResults(model_results=[results]), |
| raise_error=True, |
| )[results.model_name] |
|
|
| rows = [] |
|
|
| for task in benchmark.tasks: |
| name = task.metadata.name |
| score = scores.get(name) |
|
|
| rows.append({ |
| "Task": name, |
| "Score": round(score, 2) if score is not None else None, |
| }) |
|
|
| with open(output_dir / "task_scores.csv", "w", newline="") as file: |
| writer = csv.DictWriter( |
| file, |
| fieldnames=["Task", "Score"], |
| ) |
| writer.writeheader() |
| writer.writerows(rows) |
|
|
| with open(output_dir / "summary.json", "w") as file: |
| json.dump( |
| { |
| "model": results.model_name, |
| "benchmark": BENCHMARK, |
| "summary": summary, |
| }, |
| file, |
| indent=2, |
| ) |
|
|
| print(f"\n{args.model} ({BENCHMARK})") |
| print("| Task | Score |") |
| print("|---|---:|") |
|
|
| for row in rows: |
| score = row["Score"] |
| value = f"{score:.2f}" if score is not None else "N/A" |
| print(f"| {row['Task']} | {value} |") |
|
|
| print("\nAggregate scores:") |
|
|
| for name, score in summary.items(): |
| if name == "Rank": |
| continue |
|
|
| value = f"{score * 100:.2f}" if score is not None else "N/A" |
| print(f"| {name} | {value} |") |
|
|
| print(f"\nResults saved to: {output_dir.resolve()}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
