From d698080e9ac869bf736d49174ee16b47b861177c Mon Sep 17 00:00:00 2001 From: Pringled Date: Mon, 28 Sep 2026 17:21:07 +0200 Subject: [PATCH] perf: reuse ranking embeddings when selecting representatives, refresh benchmarks and bump to 0.5.0 Representative selection encoded its top candidates again after ranking. The ranking now returns the embeddings in ranked order (index vectors in self mode, the query embeddings in cross mode) and diversification reuses them. Benchmarks are rerun on a MacBook Pro (Apple M5, 48 GB RAM) to reflect #89 and #90, and the version is bumped to 0.5.0 for their breaking changes. --- README.md | 2 +- benchmarks/README.md | 87 +++++----- .../image_train_benchmark_results.json | 20 +-- .../image_train_test_benchmark_results.json | 16 +- .../results/train_benchmark_results.json | 158 +++++++++--------- .../results/train_test_benchmark_results.json | 138 +++++++-------- semhash/semhash.py | 78 ++++----- semhash/version.py | 2 +- tests/test_semhash.py | 32 +++- 9 files changed, 274 insertions(+), 259 deletions(-) diff --git a/README.md b/README.md index 03acb96..93ac89b 100644 --- a/README.md +++ b/README.md @@ -527,7 +527,7 @@ deduplicated_records = semhash.self_deduplicate().selected ## Benchmarks -SemHash is extremely fast and scales to large datasets with millions of records. We've benchmarked both text and image deduplication across a variety of datasets. For example, deduplicating text 1.8M records takes only ~83 seconds on CPU. +SemHash is extremely fast and scales to large datasets with millions of records. We've benchmarked both text and image deduplication across a variety of datasets. For example, deduplicating text 1.8M records takes only ~74 seconds on CPU. For detailed benchmark results and analysis, see the [benchmarks directory](benchmarks/README.md). diff --git a/benchmarks/README.md b/benchmarks/README.md index 7ebf432..e8580ce 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -21,6 +21,7 @@ This directory contains the benchmarking code and results for SemHash. The bench ### Setup All text benchmarks were run with the following configuration: +- **Device**: MacBook Pro (Apple M5, 48 GB RAM) - **CPU-only**: All benchmarks run on CPU (no GPU acceleration) - **ANN backend**: Default backend (USearch) - **Encoder**: Default encoder ([potion-base-8M](https://huggingface.co/minishlab/potion-base-8M)) @@ -35,23 +36,23 @@ This benchmark measures the performance of deduplicating within a single trainin | Dataset | Original Train Size | Deduplicated Train Size | % Removed | Deduplication Time (s) | |----------------------|----------------------|--------------------------|------------|--------------------------| -| bbc | 1225 | 1144 | 6.61 | 0.57 | -| senteval_cr | 3012 | 2990 | 0.73 | 0.14 | -| tweet_sentiment_extraction | 27481 | 26695 | 2.86 | 1.77 | -| emotion | 16000 | 15695 | 1.91 | 0.77 | -| amazon_counterfactual | 5000 | 4992 | 0.16 | 0.33 | -| ag_news | 120000 | 106921 | 10.90 | 5.20 | -| enron_spam | 31716 | 20540 | 35.24 | 2.03 | -| subj | 8000 | 7990 | 0.12 | 0.63 | -| sst5 | 8544 | 8526 | 0.21 | 0.58 | -| 20_newgroups | 11314 | 10684 | 5.57 | 0.73 | -| hatespeech_offensive | 22783 | 22090 | 3.04 | 0.92 | -| ade | 17637 | 15718 | 10.88 | 0.73 | -| imdb | 25000 | 24830 | 0.68 | 1.76 | -| massive_scenario | 11514 | 9366 | 18.66 | 0.47 | -| student | 117519 | 63856 | 45.66 | 8.80 | -| squad_v2 | 130319 | 109698 | 15.82 | 8.81 | -| wikitext | 1801350 | 884645 | 50.89 | 83.53 | +| bbc | 1225 | 1148 | 6.29 | 0.19 | +| senteval_cr | 3012 | 2992 | 0.66 | 0.15 | +| tweet_sentiment_extraction | 27481 | 26775 | 2.57 | 1.67 | +| emotion | 16000 | 15739 | 1.63 | 0.68 | +| amazon_counterfactual | 5000 | 4992 | 0.16 | 0.27 | +| ag_news | 120000 | 107882 | 10.10 | 5.75 | +| enron_spam | 31716 | 21121 | 33.41 | 1.63 | +| subj | 8000 | 7990 | 0.12 | 0.45 | +| sst5 | 8544 | 8526 | 0.21 | 0.46 | +| 20_newgroups | 11314 | 10717 | 5.28 | 0.61 | +| hatespeech_offensive | 22783 | 22233 | 2.41 | 0.96 | +| ade | 17637 | 15723 | 10.85 | 0.74 | +| imdb | 25000 | 24847 | 0.61 | 1.64 | +| massive_scenario | 11514 | 9665 | 16.06 | 0.45 | +| student | 117519 | 69696 | 40.69 | 8.28 | +| squad_v2 | 130319 | 110480 | 15.22 | 9.66 | +| wikitext | 1801350 | 900554 | 50.01 | 74.31 | ### Train/Test Deduplication Benchmark @@ -59,29 +60,29 @@ This benchmark measures the performance of deduplicating a test dataset against | Dataset | Train Size | Test Size | Deduplicated Test Size | % Removed | Deduplication Time (s) | |----------------------|--------------|--------------|--------------------------|------------|--------------------------| -| bbc | 1225 | 1000 | 870 | 13.00 | 0.71 | -| senteval_cr | 3012 | 753 | 750 | 0.40 | 0.13 | -| tweet_sentiment_extraction | 27481 | 3534 | 3412 | 3.45 | 1.53 | -| emotion | 16000 | 2000 | 1926 | 3.70 | 0.65 | -| amazon_counterfactual | 5000 | 5000 | 4990 | 0.20 | 0.51 | -| ag_news | 120000 | 7600 | 6198 | 18.45 | 3.74 | -| enron_spam | 31716 | 2000 | 1060 | 47.00 | 1.94 | -| subj | 8000 | 2000 | 1999 | 0.05 | 0.62 | -| sst5 | 8544 | 2210 | 2205 | 0.23 | 0.59 | -| 20_newgroups | 11314 | 7532 | 7098 | 5.76 | 2.25 | -| hatespeech_offensive | 22783 | 2000 | 1925 | 3.75 | 0.77 | -| ade | 17637 | 5879 | 4952 | 15.77 | 0.81 | -| imdb | 25000 | 25000 | 24795 | 0.82 | 2.81 | -| massive_scenario | 11514 | 2974 | 2190 | 26.36 | 0.46 | -| student | 117519 | 5000 | 2393 | 52.14 | 3.78 | -| squad_v2 | 130319 | 11873 | 11863 | 0.08 | 7.13 | -| wikitext | 1801350 | 4358 | 2139 | 50.92 | 40.32 | +| bbc | 1225 | 1000 | 874 | 12.60 | 0.29 | +| senteval_cr | 3012 | 753 | 750 | 0.40 | 0.15 | +| tweet_sentiment_extraction | 27481 | 3534 | 3411 | 3.48 | 1.48 | +| emotion | 16000 | 2000 | 1926 | 3.70 | 0.58 | +| amazon_counterfactual | 5000 | 5000 | 4990 | 0.20 | 0.44 | +| ag_news | 120000 | 7600 | 6201 | 18.41 | 3.95 | +| enron_spam | 31716 | 2000 | 1064 | 46.80 | 1.56 | +| subj | 8000 | 2000 | 1999 | 0.05 | 0.44 | +| sst5 | 8544 | 2210 | 2205 | 0.23 | 0.45 | +| 20_newgroups | 11314 | 7532 | 7098 | 5.76 | 1.51 | +| hatespeech_offensive | 22783 | 2000 | 1925 | 3.75 | 0.78 | +| ade | 17637 | 5879 | 4953 | 15.75 | 0.79 | +| imdb | 25000 | 25000 | 24797 | 0.81 | 2.55 | +| massive_scenario | 11514 | 2974 | 2188 | 26.43 | 0.47 | +| student | 117519 | 5000 | 2400 | 52.00 | 4.36 | +| squad_v2 | 130319 | 11873 | 11863 | 0.08 | 7.07 | +| wikitext | 1801350 | 4358 | 2134 | 51.03 | 46.29 | ### Key Findings SemHash is extremely fast and scales to large datasets with millions of records. Some notable findings include: -- **Speed**: Deduplication is fast even for large datasets (e.g., 1.8M records in ~83 seconds) +- **Speed**: Deduplication is fast even for large datasets (e.g., 1.8M records in ~74 seconds) - **Train/Test Leakage**: Several datasets show significant train/test overlap: - `enron_spam`: 47% of test data overlaps with training data - `student`: 52% of test data overlaps with training data @@ -106,7 +107,7 @@ make benchmark-text ### Setup All image benchmarks were run with the following configuration: -- **Device**: Apple Silicon GPU (MPS) +- **Device**: MacBook Pro (Apple M5, 48 GB RAM), GPU via MPS - **ANN backend**: Default backend (USearch) - **Encoder**: MobileNetV3-Small ([mobilenetv3_small_100.lamb_in1k](https://huggingface.co/timm/mobilenetv3_small_100.lamb_in1k)) - **Batch size**: 128 images per batch @@ -120,8 +121,8 @@ This benchmark measures the performance of deduplicating within a single trainin | Dataset | Original Train Size | Deduplicated Train Size | % Removed | Deduplication Time (s) | |----------------------|----------------------|--------------------------|------------|--------------------------| -| cifar10 | 50000 | 48274 | 3.45 | 61.20 | -| fashion_mnist | 60000 | 16714 | 72.14 | 86.61 | +| cifar10 | 50000 | 48390 | 3.22 | 38.05 | +| fashion_mnist | 60000 | 23549 | 60.75 | 42.70 | #### Train/Test Deduplication Benchmark @@ -129,14 +130,14 @@ This benchmark measures the performance of deduplicating a test dataset against | Dataset | Train Size | Test Size | Deduplicated Test Size | % Removed | Deduplication Time (s) | |----------------------|--------------|--------------|--------------------------|------------|--------------------------| -| cifar10 | 50000 | 10000 | 9397 | 6.03 | 67.43 | -| fashion_mnist | 60000 | 10000 | 2052 | 79.48 | 72.14 | +| cifar10 | 50000 | 10000 | 9397 | 6.03 | 43.35 | +| fashion_mnist | 60000 | 10000 | 2056 | 79.44 | 46.89 | ### Key Findings -- **Fashion-MNIST high deduplication**: Fashion-MNIST shows very high duplication rates (72% train, 79% test) due to the simple nature of the dataset (10 clothing categories with similar items) -- **CIFAR-10 moderate deduplication**: CIFAR-10 shows lower duplication (3.45% train, 6.03% test) as it contains more diverse natural images -- **Speed**: Image deduplication is fast even for large datasets (60k images in ~87 seconds on MPS); note that the actual deduplication step is quick, with most time spent on encoding images +- **Fashion-MNIST high deduplication**: Fashion-MNIST shows very high duplication rates (61% train, 79% test) due to the simple nature of the dataset (10 clothing categories with similar items) +- **CIFAR-10 moderate deduplication**: CIFAR-10 shows lower duplication (3.22% train, 6.03% test) as it contains more diverse natural images +- **Speed**: Image deduplication is fast even for large datasets (60k images in ~43 seconds on MPS); note that the actual deduplication step is quick, with most time spent on encoding images ### Running Image Benchmarks diff --git a/benchmarks/results/image_train_benchmark_results.json b/benchmarks/results/image_train_benchmark_results.json index 53fa71b..171402c 100644 --- a/benchmarks/results/image_train_benchmark_results.json +++ b/benchmarks/results/image_train_benchmark_results.json @@ -2,19 +2,19 @@ { "dataset": "cifar10", "original_train_size": 50000, - "deduplicated_train_size": 48274, - "percent_removed": 3.4519999999999995, - "build_time_seconds": 56.00128899999254, - "deduplication_time_seconds": 5.201297917010379, - "time_seconds": 61.20258691700292 + "deduplicated_train_size": 48390, + "percent_removed": 3.2200000000000006, + "build_time_seconds": 36.14036558299995, + "deduplication_time_seconds": 1.9100148340003216, + "time_seconds": 38.05038041700027 }, { "dataset": "fashion_mnist", "original_train_size": 60000, - "deduplicated_train_size": 16714, - "percent_removed": 72.14333333333333, - "build_time_seconds": 61.14413262500602, - "deduplication_time_seconds": 25.46288070900482, - "time_seconds": 86.60701333401084 + "deduplicated_train_size": 23549, + "percent_removed": 60.75166666666667, + "build_time_seconds": 38.96308354100256, + "deduplication_time_seconds": 3.733158041999559, + "time_seconds": 42.69624158300212 } ] diff --git a/benchmarks/results/image_train_test_benchmark_results.json b/benchmarks/results/image_train_test_benchmark_results.json index 290c7e3..8932214 100644 --- a/benchmarks/results/image_train_test_benchmark_results.json +++ b/benchmarks/results/image_train_test_benchmark_results.json @@ -5,18 +5,18 @@ "test_size": 10000, "deduplicated_test_size": 9397, "percent_removed": 6.030000000000002, - "build_time_seconds": 56.00128899999254, - "deduplication_time_seconds": 11.428115875009098, - "time_seconds": 67.42940487500164 + "build_time_seconds": 36.14036558299995, + "deduplication_time_seconds": 7.210748125002283, + "time_seconds": 43.35111370800223 }, { "dataset": "fashion_mnist", "train_size": 60000, "test_size": 10000, - "deduplicated_test_size": 2052, - "percent_removed": 79.47999999999999, - "build_time_seconds": 61.14413262500602, - "deduplication_time_seconds": 10.998616750002839, - "time_seconds": 72.14274937500886 + "deduplicated_test_size": 2056, + "percent_removed": 79.44, + "build_time_seconds": 38.96308354100256, + "deduplication_time_seconds": 7.9287375419989985, + "time_seconds": 46.89182108300156 } ] diff --git a/benchmarks/results/train_benchmark_results.json b/benchmarks/results/train_benchmark_results.json index 01cc3f4..ae092db 100644 --- a/benchmarks/results/train_benchmark_results.json +++ b/benchmarks/results/train_benchmark_results.json @@ -2,154 +2,154 @@ { "dataset": "bbc", "original_train_size": 1225, - "deduplicated_train_size": 1144, - "percent_removed": 6.612244897959185, - "build_time_seconds": 0.5598082079086453, - "deduplication_time_seconds": 0.008702374994754791, - "time_seconds": 0.5685105829034001 + "deduplicated_train_size": 1148, + "percent_removed": 6.285714285714283, + "build_time_seconds": 0.1762163750026957, + "deduplication_time_seconds": 0.009695124997961102, + "time_seconds": 0.18591150000065682 }, { "dataset": "senteval_cr", "original_train_size": 3012, - "deduplicated_train_size": 2990, - "percent_removed": 0.7304116865869847, - "build_time_seconds": 0.10847400000784546, - "deduplication_time_seconds": 0.027519959025084972, - "time_seconds": 0.13599395903293043 + "deduplicated_train_size": 2992, + "percent_removed": 0.6640106241699861, + "build_time_seconds": 0.12137100000109058, + "deduplication_time_seconds": 0.029432167000777554, + "time_seconds": 0.15080316700186813 }, { "dataset": "tweet_sentiment_extraction", "original_train_size": 27481, - "deduplicated_train_size": 26695, - "percent_removed": 2.860157927295226, - "build_time_seconds": 1.3568968329345807, - "deduplication_time_seconds": 0.41194633406121284, - "time_seconds": 1.7688431669957936 + "deduplicated_train_size": 26775, + "percent_removed": 2.569047705687566, + "build_time_seconds": 1.3021473750013683, + "deduplication_time_seconds": 0.37088470900198445, + "time_seconds": 1.6730320840033528 }, { "dataset": "emotion", "original_train_size": 16000, - "deduplicated_train_size": 15695, - "percent_removed": 1.9062499999999982, - "build_time_seconds": 0.5511152499821037, - "deduplication_time_seconds": 0.21407662506680936, - "time_seconds": 0.7651918750489131 + "deduplicated_train_size": 15739, + "percent_removed": 1.6312499999999952, + "build_time_seconds": 0.4828177079980378, + "deduplication_time_seconds": 0.19528612499925657, + "time_seconds": 0.6781038329972944 }, { "dataset": "amazon_counterfactual", "original_train_size": 5000, "deduplicated_train_size": 4992, "percent_removed": 0.16000000000000458, - "build_time_seconds": 0.2848535830853507, - "deduplication_time_seconds": 0.048574666026979685, - "time_seconds": 0.3334282491123304 + "build_time_seconds": 0.22153624999918975, + "deduplication_time_seconds": 0.04794116700213635, + "time_seconds": 0.2694774170013261 }, { "dataset": "ag_news", "original_train_size": 120000, - "deduplicated_train_size": 106921, - "percent_removed": 10.899166666666671, - "build_time_seconds": 3.0319770000642166, - "deduplication_time_seconds": 2.171258582966402, - "time_seconds": 5.203235583030619 + "deduplicated_train_size": 107882, + "percent_removed": 10.098333333333331, + "build_time_seconds": 3.275759290998394, + "deduplication_time_seconds": 2.4731626249995315, + "time_seconds": 5.748921915997926 }, { "dataset": "enron_spam", "original_train_size": 31716, - "deduplicated_train_size": 20540, - "percent_removed": 35.23773489721276, - "build_time_seconds": 1.3818323339801282, - "deduplication_time_seconds": 0.6438171250047162, - "time_seconds": 2.0256494589848444 + "deduplicated_train_size": 21121, + "percent_removed": 33.405851935931395, + "build_time_seconds": 1.1878026249978575, + "deduplication_time_seconds": 0.44431537499986007, + "time_seconds": 1.6321179999977176 }, { "dataset": "subj", "original_train_size": 8000, "deduplicated_train_size": 7990, "percent_removed": 0.12499999999999734, - "build_time_seconds": 0.5059439589967951, - "deduplication_time_seconds": 0.12505983305163682, - "time_seconds": 0.6310037920484319 + "build_time_seconds": 0.34191370800181176, + "deduplication_time_seconds": 0.10950004199912655, + "time_seconds": 0.4514137500009383 }, { "dataset": "sst5", "original_train_size": 8544, "deduplicated_train_size": 8526, "percent_removed": 0.2106741573033699, - "build_time_seconds": 0.4805819580797106, - "deduplication_time_seconds": 0.10166720801498741, - "time_seconds": 0.582249166094698 + "build_time_seconds": 0.35040350000053877, + "deduplication_time_seconds": 0.10490237500198418, + "time_seconds": 0.45530587500252295 }, { "dataset": "20_newgroups", "original_train_size": 11314, - "deduplicated_train_size": 10684, - "percent_removed": 5.568322432384654, - "build_time_seconds": 0.610724583035335, - "deduplication_time_seconds": 0.11600329191423953, - "time_seconds": 0.7267278749495745 + "deduplicated_train_size": 10717, + "percent_removed": 5.276648400212125, + "build_time_seconds": 0.49274829199930537, + "deduplication_time_seconds": 0.11272816700147814, + "time_seconds": 0.6054764590007835 }, { "dataset": "hatespeech_offensive", "original_train_size": 22783, - "deduplicated_train_size": 22090, - "percent_removed": 3.0417416494754823, - "build_time_seconds": 0.6471997499465942, - "deduplication_time_seconds": 0.2704670410603285, - "time_seconds": 0.9176667910069227 + "deduplicated_train_size": 22233, + "percent_removed": 2.4140806741868914, + "build_time_seconds": 0.6664763750013663, + "deduplication_time_seconds": 0.2959712080009922, + "time_seconds": 0.9624475830023584 }, { "dataset": "ade", "original_train_size": 17637, - "deduplicated_train_size": 15718, - "percent_removed": 10.880535238419231, - "build_time_seconds": 0.5221591669833288, - "deduplication_time_seconds": 0.20764074998442084, - "time_seconds": 0.7297999169677496 + "deduplicated_train_size": 15723, + "percent_removed": 10.852185745875154, + "build_time_seconds": 0.5355710839976382, + "deduplication_time_seconds": 0.2090475000004517, + "time_seconds": 0.7446185839980899 }, { "dataset": "imdb", "original_train_size": 25000, - "deduplicated_train_size": 24830, - "percent_removed": 0.6800000000000028, - "build_time_seconds": 1.460668999934569, - "deduplication_time_seconds": 0.29758112493436784, - "time_seconds": 1.7582501248689368 + "deduplicated_train_size": 24847, + "percent_removed": 0.6120000000000014, + "build_time_seconds": 1.325603665998642, + "deduplication_time_seconds": 0.31292716700045276, + "time_seconds": 1.6385308329990949 }, { "dataset": "massive_scenario", "original_train_size": 11514, - "deduplicated_train_size": 9366, - "percent_removed": 18.655549765502865, - "build_time_seconds": 0.35503324994351715, - "deduplication_time_seconds": 0.11619104200508446, - "time_seconds": 0.4712242919486016 + "deduplicated_train_size": 9665, + "percent_removed": 16.05871113427132, + "build_time_seconds": 0.3495884580006532, + "deduplication_time_seconds": 0.10384745799819939, + "time_seconds": 0.4534359159988526 }, { "dataset": "student", "original_train_size": 117519, - "deduplicated_train_size": 63856, - "percent_removed": 45.66325445247151, - "build_time_seconds": 2.9044899590080604, - "deduplication_time_seconds": 5.895973875070922, - "time_seconds": 8.800463834078982 + "deduplicated_train_size": 69696, + "percent_removed": 40.69384525055523, + "build_time_seconds": 3.4031454579999263, + "deduplication_time_seconds": 4.874128916999325, + "time_seconds": 8.277274374999251 }, { "dataset": "squad_v2", "original_train_size": 130319, - "deduplicated_train_size": 109698, - "percent_removed": 15.823479308466148, - "build_time_seconds": 6.078755749971606, - "deduplication_time_seconds": 2.7270843340083957, - "time_seconds": 8.805840083980002 + "deduplicated_train_size": 110480, + "percent_removed": 15.22341331655399, + "build_time_seconds": 5.937046999999438, + "deduplication_time_seconds": 3.722715916999732, + "time_seconds": 9.65976291699917 }, { "dataset": "wikitext", "original_train_size": 1801350, - "deduplicated_train_size": 884645, - "percent_removed": 50.88988813945097, - "build_time_seconds": 39.38258587510791, - "deduplication_time_seconds": 44.1503732081037, - "time_seconds": 83.53295908321161 + "deduplicated_train_size": 900554, + "percent_removed": 50.00671718433397, + "build_time_seconds": 44.6954585420026, + "deduplication_time_seconds": 29.619445874999656, + "time_seconds": 74.31490441700225 } ] diff --git a/benchmarks/results/train_test_benchmark_results.json b/benchmarks/results/train_test_benchmark_results.json index 0508700..80d1d58 100644 --- a/benchmarks/results/train_test_benchmark_results.json +++ b/benchmarks/results/train_test_benchmark_results.json @@ -3,11 +3,11 @@ "dataset": "bbc", "train_size": 1225, "test_size": 1000, - "deduplicated_test_size": 870, - "percent_removed": 13.0, - "build_time_seconds": 0.5598082079086453, - "deduplication_time_seconds": 0.1528247919632122, - "time_seconds": 0.7126329998718575 + "deduplicated_test_size": 874, + "percent_removed": 12.6, + "build_time_seconds": 0.1762163750026957, + "deduplication_time_seconds": 0.10950220800077659, + "time_seconds": 0.2857185830034723 }, { "dataset": "senteval_cr", @@ -15,19 +15,19 @@ "test_size": 753, "deduplicated_test_size": 750, "percent_removed": 0.3984063745019917, - "build_time_seconds": 0.10847400000784546, - "deduplication_time_seconds": 0.019297500024549663, - "time_seconds": 0.12777150003239512 + "build_time_seconds": 0.12137100000109058, + "deduplication_time_seconds": 0.027684791002684506, + "time_seconds": 0.14905579100377508 }, { "dataset": "tweet_sentiment_extraction", "train_size": 27481, "test_size": 3534, - "deduplicated_test_size": 3412, - "percent_removed": 3.452178834182229, - "build_time_seconds": 1.3568968329345807, - "deduplication_time_seconds": 0.17268049996346235, - "time_seconds": 1.529577332898043 + "deduplicated_test_size": 3411, + "percent_removed": 3.4804753820034007, + "build_time_seconds": 1.3021473750013683, + "deduplication_time_seconds": 0.1789753750017553, + "time_seconds": 1.4811227500031237 }, { "dataset": "emotion", @@ -35,9 +35,9 @@ "test_size": 2000, "deduplicated_test_size": 1926, "percent_removed": 3.7000000000000033, - "build_time_seconds": 0.5511152499821037, - "deduplication_time_seconds": 0.10135454102419317, - "time_seconds": 0.6524697910062969 + "build_time_seconds": 0.4828177079980378, + "deduplication_time_seconds": 0.09223724999901606, + "time_seconds": 0.5750549579970539 }, { "dataset": "amazon_counterfactual", @@ -45,29 +45,29 @@ "test_size": 5000, "deduplicated_test_size": 4990, "percent_removed": 0.20000000000000018, - "build_time_seconds": 0.2848535830853507, - "deduplication_time_seconds": 0.22846354101784527, - "time_seconds": 0.513317124103196 + "build_time_seconds": 0.22153624999918975, + "deduplication_time_seconds": 0.21453241600102047, + "time_seconds": 0.4360686660002102 }, { "dataset": "ag_news", "train_size": 120000, "test_size": 7600, - "deduplicated_test_size": 6198, - "percent_removed": 18.447368421052634, - "build_time_seconds": 3.0319770000642166, - "deduplication_time_seconds": 0.7034984159981832, - "time_seconds": 3.7354754160623997 + "deduplicated_test_size": 6201, + "percent_removed": 18.407894736842103, + "build_time_seconds": 3.275759290998394, + "deduplication_time_seconds": 0.6717791250011942, + "time_seconds": 3.9475384159995883 }, { "dataset": "enron_spam", "train_size": 31716, "test_size": 2000, - "deduplicated_test_size": 1060, - "percent_removed": 47.0, - "build_time_seconds": 1.3818323339801282, - "deduplication_time_seconds": 0.553584959008731, - "time_seconds": 1.9354172929888591 + "deduplicated_test_size": 1064, + "percent_removed": 46.8, + "build_time_seconds": 1.1878026249978575, + "deduplication_time_seconds": 0.37047674999848823, + "time_seconds": 1.5582793749963457 }, { "dataset": "subj", @@ -75,9 +75,9 @@ "test_size": 2000, "deduplicated_test_size": 1999, "percent_removed": 0.04999999999999449, - "build_time_seconds": 0.5059439589967951, - "deduplication_time_seconds": 0.11624520795885473, - "time_seconds": 0.6221891669556499 + "build_time_seconds": 0.34191370800181176, + "deduplication_time_seconds": 0.09621087499908754, + "time_seconds": 0.4381245830008993 }, { "dataset": "sst5", @@ -85,9 +85,9 @@ "test_size": 2210, "deduplicated_test_size": 2205, "percent_removed": 0.2262443438914019, - "build_time_seconds": 0.4805819580797106, - "deduplication_time_seconds": 0.11375170899555087, - "time_seconds": 0.5943336670752615 + "build_time_seconds": 0.35040350000053877, + "deduplication_time_seconds": 0.09645383300085086, + "time_seconds": 0.44685733300138963 }, { "dataset": "20_newgroups", @@ -95,9 +95,9 @@ "test_size": 7532, "deduplicated_test_size": 7098, "percent_removed": 5.762081784386619, - "build_time_seconds": 0.610724583035335, - "deduplication_time_seconds": 1.6346445409581065, - "time_seconds": 2.2453691239934415 + "build_time_seconds": 0.49274829199930537, + "deduplication_time_seconds": 1.0146297089995642, + "time_seconds": 1.5073780009988695 }, { "dataset": "hatespeech_offensive", @@ -105,49 +105,49 @@ "test_size": 2000, "deduplicated_test_size": 1925, "percent_removed": 3.749999999999998, - "build_time_seconds": 0.6471997499465942, - "deduplication_time_seconds": 0.12372829194646329, - "time_seconds": 0.7709280418930575 + "build_time_seconds": 0.6664763750013663, + "deduplication_time_seconds": 0.11403075000271201, + "time_seconds": 0.7805071250040783 }, { "dataset": "ade", "train_size": 17637, "test_size": 5879, - "deduplicated_test_size": 4952, - "percent_removed": 15.76798775301922, - "build_time_seconds": 0.5221591669833288, - "deduplication_time_seconds": 0.28758599993307143, - "time_seconds": 0.8097451669164002 + "deduplicated_test_size": 4953, + "percent_removed": 15.75097805749277, + "build_time_seconds": 0.5355710839976382, + "deduplication_time_seconds": 0.2526170000019192, + "time_seconds": 0.7881880839995574 }, { "dataset": "imdb", "train_size": 25000, "test_size": 25000, - "deduplicated_test_size": 24795, - "percent_removed": 0.8199999999999985, - "build_time_seconds": 1.460668999934569, - "deduplication_time_seconds": 1.3489695829339325, - "time_seconds": 2.8096385828685015 + "deduplicated_test_size": 24797, + "percent_removed": 0.8120000000000016, + "build_time_seconds": 1.325603665998642, + "deduplication_time_seconds": 1.2251233749993844, + "time_seconds": 2.5507270409980265 }, { "dataset": "massive_scenario", "train_size": 11514, "test_size": 2974, - "deduplicated_test_size": 2190, - "percent_removed": 26.36180228648285, - "build_time_seconds": 0.35503324994351715, - "deduplication_time_seconds": 0.10878237499855459, - "time_seconds": 0.46381562494207174 + "deduplicated_test_size": 2188, + "percent_removed": 26.429051782111635, + "build_time_seconds": 0.3495884580006532, + "deduplication_time_seconds": 0.11969712500285823, + "time_seconds": 0.4692855830035114 }, { "dataset": "student", "train_size": 117519, "test_size": 5000, - "deduplicated_test_size": 2393, - "percent_removed": 52.14, - "build_time_seconds": 2.9044899590080604, - "deduplication_time_seconds": 0.8721794589655474, - "time_seconds": 3.776669417973608 + "deduplicated_test_size": 2400, + "percent_removed": 52.0, + "build_time_seconds": 3.4031454579999263, + "deduplication_time_seconds": 0.9607139999970968, + "time_seconds": 4.363859457997023 }, { "dataset": "squad_v2", @@ -155,18 +155,18 @@ "test_size": 11873, "deduplicated_test_size": 11863, "percent_removed": 0.08422471153036737, - "build_time_seconds": 6.078755749971606, - "deduplication_time_seconds": 1.0497459589969367, - "time_seconds": 7.1285017089685425 + "build_time_seconds": 5.937046999999438, + "deduplication_time_seconds": 1.1297352090005006, + "time_seconds": 7.066782208999939 }, { "dataset": "wikitext", "train_size": 1801350, "test_size": 4358, - "deduplicated_test_size": 2139, - "percent_removed": 50.91785222579165, - "build_time_seconds": 39.38258587510791, - "deduplication_time_seconds": 0.9325925830053166, - "time_seconds": 40.31517845811322 + "deduplicated_test_size": 2134, + "percent_removed": 51.032583754015604, + "build_time_seconds": 44.6954585420026, + "deduplication_time_seconds": 1.594421958001476, + "time_seconds": 46.289880500004074 } ] diff --git a/semhash/semhash.py b/semhash/semhash.py index b72225f..56c2a66 100644 --- a/semhash/semhash.py +++ b/semhash/semhash.py @@ -9,6 +9,7 @@ from model2vec import StaticModel from pyversity import Strategy, diversify from vicinity import Backend +from vicinity.datatypes import SingleQueryResult from semhash.datamodels import DeduplicationResult, DuplicateRecord, FilterResult from semhash.index import Index @@ -44,7 +45,7 @@ def __init__(self, index: Index, model: Encoder, columns: Sequence[str], was_str self.model = model self.columns = columns self._was_string = was_string - self._ranking_cache: FilterResult | None = None + self._ranking_cache: tuple[FilterResult, np.ndarray] | None = None @classmethod def from_records( @@ -284,10 +285,12 @@ def find_representative( :param strategy: Diversification strategy (MMR, MSD, DPP, COVER, SSD). Default is MMR. :return: A FilterResult with the diversified candidates. """ - ranking = self._rank_by_average_similarity(records) + ranking, embeddings = self._rank_by_average_similarity(records) if candidate_limit == "auto": candidate_limit = compute_candidate_limit(total=len(ranking.selected), selection_size=selection_size) - return self._to_output(self._diversify(ranking, candidate_limit, selection_size, diversity, strategy)) + return self._to_output( + self._diversify(ranking, embeddings, candidate_limit, selection_size, diversity, strategy) + ) def self_find_representative( self, @@ -310,10 +313,12 @@ def self_find_representative( :param strategy: Diversification strategy (MMR, MSD, DPP, COVER, SSD). Default is MMR. :return: A FilterResult with the diversified representatives. """ - ranking = self._self_rank_by_average_similarity() + ranking, embeddings = self._self_rank_by_average_similarity() if candidate_limit == "auto": candidate_limit = compute_candidate_limit(total=len(ranking.selected), selection_size=selection_size) - return self._to_output(self._diversify(ranking, candidate_limit, selection_size, diversity, strategy)) + return self._to_output( + self._diversify(ranking, embeddings, candidate_limit, selection_size, diversity, strategy) + ) def filter_outliers( self, @@ -333,7 +338,8 @@ def filter_outliers( """ if outlier_percentage < 0.0 or outlier_percentage > 1.0: raise ValueError("outlier_percentage must be between 0 and 1") - return self._split_into_outliers(self._rank_by_average_similarity(records), outlier_percentage) + ranking, _ = self._rank_by_average_similarity(records) + return self._split_into_outliers(ranking, outlier_percentage) def self_filter_outliers( self, @@ -351,7 +357,7 @@ def self_filter_outliers( """ if outlier_percentage < 0.0 or outlier_percentage > 1.0: raise ValueError("outlier_percentage must be between 0 and 1") - ranking = self._self_rank_by_average_similarity() + ranking, _ = self._self_rank_by_average_similarity() # Exact copies share the score of their group, so every fitted record is returned. groups = {id(group[0]): group for group in self.index.items} ranked = [ @@ -400,66 +406,61 @@ def _to_output(self, result: FilterResult) -> FilterResult: def _rank_by_average_similarity( self, records: Sequence[Record], - ) -> FilterResult: + ) -> tuple[FilterResult, np.ndarray]: """ Rank a given set of records based on the average cosine similarity of the neighbors in the fitted index. :param records: A sequence of records. - :return: A FilterResult containing the ranking (records sorted and their average similarity scores). + :return: A FilterResult containing the ranking, and the embeddings of the records in ranked order. """ dict_records = self._validate_if_strings(records) embeddings = featurize(records=dict_records, columns=self.columns, model=self.model) results = self.index.query_top_k(embeddings, k=100, vectors_are_in_index=False) - - # Compute the average similarity for each record. - sorted_scores = sorted( - ((record, np.mean(sims)) for record, (_, sims) in zip(dict_records, results)), - key=lambda x: x[1], - reverse=True, - ) - selected, scores_selected = zip(*sorted_scores) - - return FilterResult( - selected=list(selected), - filtered=[], - scores_selected=list(scores_selected), - scores_filtered=[], - ) + return self._rank_from_neighbors(dict_records, embeddings, results) def _self_rank_by_average_similarity( self, - ) -> FilterResult: + ) -> tuple[FilterResult, np.ndarray]: """ Rank the records stored in the fitted index based on the average cosine similarity of the neighbors. - :return: A FilterResult containing the ranking. + :return: A FilterResult containing the ranking, and the embeddings of the records in ranked order. """ if self._ranking_cache is not None: return self._ranking_cache dict_records = [record[0] for record in self.index.items] results = self.index.query_top_k(self.index.vectors, k=100, vectors_are_in_index=True) + self._ranking_cache = self._rank_from_neighbors(dict_records, self.index.vectors, results) + return self._ranking_cache - # Compute the average similarity for each record. - sorted_scores = sorted( - ((record, np.mean(sims)) for record, (_, sims) in zip(dict_records, results)), - key=lambda x: x[1], - reverse=True, - ) - selected, scores_selected = zip(*sorted_scores) + @staticmethod + def _rank_from_neighbors( + records: list[dict[str, Any]], embeddings: np.ndarray, neighbors: list[SingleQueryResult] + ) -> tuple[FilterResult, np.ndarray]: + """ + Rank records by the average similarity of their already computed neighbors, keeping their embeddings aligned. + :param records: The records to rank. + :param embeddings: The embeddings of the records. + :param neighbors: The nearest neighbors of each record. + :return: A FilterResult containing the ranking, and the embeddings of the records in ranked order. + """ + scores = [np.mean(sims) for _, sims in neighbors] + # Stable descending sort, so ties keep their input order. + order = sorted(range(len(records)), key=lambda i: scores[i], reverse=True) ranking = FilterResult( - selected=list(selected), + selected=[records[i] for i in order], filtered=[], - scores_selected=list(scores_selected), + scores_selected=[scores[i] for i in order], scores_filtered=[], ) - self._ranking_cache = ranking - return ranking + return ranking, embeddings[order] def _diversify( self, ranked_results: FilterResult, + ranked_embeddings: np.ndarray, candidate_limit: int, selection_size: int, diversity: float, @@ -472,9 +473,8 @@ def _diversify( if not candidates: return FilterResult(selected=[], filtered=[], scores_selected=[], scores_filtered=[]) - embeddings = featurize(records=candidates, columns=self.columns, model=self.model) result = diversify( - embeddings=embeddings, + embeddings=ranked_embeddings[:candidate_limit], scores=np.array(relevance), k=selection_size, strategy=strategy, diff --git a/semhash/version.py b/semhash/version.py index 2094209..2647262 100644 --- a/semhash/version.py +++ b/semhash/version.py @@ -1,2 +1,2 @@ -__version_triple__ = (0, 4, 2) # pragma: no cover +__version_triple__ = (0, 5, 0) # pragma: no cover __version__ = ".".join(map(str, __version_triple__)) # pragma: no cover diff --git a/tests/test_semhash.py b/tests/test_semhash.py index 799baef..df4b6ef 100644 --- a/tests/test_semhash.py +++ b/tests/test_semhash.py @@ -239,30 +239,26 @@ def test_self_filter_outliers(model: Encoder, train_texts: list[str]) -> None: semhash.self_filter_outliers(outlier_percentage=1.5) -def test__diversify(monkeypatch: pytest.MonkeyPatch) -> None: +def test__diversify() -> None: """Test the _diversify method.""" - from semhash import semhash - semhash_instance = SemHash(index=None, model=None, columns=["text"], was_string=True) # Prepare a fake ranking with three records records = ["a", "b", "c"] scores = [3.0, 2.0, 1.0] ranking = FilterResult(selected=records, filtered=[], scores_selected=scores, scores_filtered=[]) - # Create dummy embeddings for the records + # Dummy embeddings for the records, in ranked order embeddings = np.array([[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]]) - # Monkeypatch featurize to return the dummy embeddings - monkeypatch.setattr(semhash, "featurize", lambda records, columns, model: embeddings) # Test diversity=0.0: pure relevance, should pick top 2 by score - result_rel = semhash_instance._diversify(ranking, candidate_limit=3, selection_size=2, diversity=0.0) + result_rel = semhash_instance._diversify(ranking, embeddings, candidate_limit=3, selection_size=2, diversity=0.0) assert result_rel.selected == ["a", "b"] # Test diversity=1.0: pure diversity, should first pick 'a', then pick most dissimilar: 'c' - result_div = semhash_instance._diversify(ranking, candidate_limit=3, selection_size=2, diversity=1.0) + result_div = semhash_instance._diversify(ranking, embeddings, candidate_limit=3, selection_size=2, diversity=1.0) assert result_div.selected == ["a", "c"] # Test empty candidates (candidate_limit=0) - result_empty = semhash_instance._diversify(ranking, candidate_limit=0, selection_size=2, diversity=0.5) + result_empty = semhash_instance._diversify(ranking, embeddings, candidate_limit=0, selection_size=2, diversity=0.5) assert result_empty.selected == [] assert result_empty.filtered == [] assert result_empty.scores_selected == [] @@ -389,6 +385,24 @@ def test_deduplicate_edge_cases(model: Encoder) -> None: semhash_dict.deduplicate([{"col": "a"}, "b"], threshold=0.95) +def test_representatives_reuse_embeddings(model: Encoder, train_texts: list[str], test_texts: list[str]) -> None: + """Representative selection reuses the ranking embeddings instead of encoding candidates again.""" + + class CountingEncoder: + calls = 0 + + def encode(self, inputs: Sequence[Any], **kwargs: Any) -> np.ndarray: + CountingEncoder.calls += 1 + return model.encode(inputs, **kwargs) + + semhash = SemHash.from_records(train_texts, model=CountingEncoder()) + CountingEncoder.calls = 0 + semhash.self_find_representative(selection_size=3) + assert CountingEncoder.calls == 0 + semhash.find_representative(test_texts, selection_size=3) + assert CountingEncoder.calls == 1 + + def test_self_filter_outliers_keeps_exact_copies(model: Encoder, train_texts: list[str]) -> None: """Exact copies are ranked with their group, so every fitted record is returned.""" records = train_texts + ["car", "car"]