Add financial filings sparse encoder v1

README.md

@@ -0,0 +1,333 @@
+---
+language:
+- en
+library_name: sentence-transformers
+pipeline_tag: feature-extraction
+tags:
+- sentence-transformers
+- sparse-encoder
+- splade
+- sparse-retrieval
+- opensearch
+- elasticsearch
+- financial-filings
+base_model: opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill
+datasets:
+- oneryalcin/financial-filings-sparse-retrieval-training
+---
+
+# Financial Filings Sparse Encoder v1
+
+This is a Sentence Transformers `SparseEncoder` / SPLADE-style model fine-tuned for financial filing retrieval.
+
+The practical recommendation from the experiments below is to index document vectors after **top-128 pruning**. In the current proxy retrieval benchmark, top-128 keeps almost all unpruned quality while reducing each document to about 126 active sparse terms.
+
+## TL;DR
+
+| setting | value |
+|---|---:|
+| Base model | `opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill` |
+| Dataset | `oneryalcin/financial-filings-sparse-retrieval-training`, `combined` config |
+| Training recipe | `SpladeLoss(SparseMultipleNegativesRankingLoss)` |
+| Final steps | 1500 |
+| Recommended serving | document `top_k=128` |
+| Triplet accuracy, top-128 | 78.0% |
+| Retrieval proxy Recall@10, top-128 | 67.2% |
+| Retrieval proxy nDCG@10, top-128 | 0.521 |
+
+This is an experiment report as much as a model card. It describes what was tried, why decisions were made, what worked, and what remains unproven.
+
+## Why this exists
+
+Financial filing search is not generic semantic similarity. Queries often refer to company events, accounting concepts, risk factors, segment details, and filing-specific language. A learned sparse model is attractive because it can improve ranking while preserving an Elasticsearch/OpenSearch-style sparse retrieval path.
+
+The goal here was to train a domain-adapted sparse encoder that can be deployed as weighted sparse terms, then test whether the learned sparse signal is meaningfully better than both the base sparse model and a lexical BM25 baseline.
+
+## Starting point
+
+This experiment started after the Sentence Transformers `v5.5.0` release and its new `train-sentence-transformers` agent skill. The release made it easier to create a complete sparse-encoder training workflow: base model choice, loss selection, sparse regularization, top-k pruning checks, and model-card packaging.
+
+Before training, the important constraints were:
+
+- Use a sparse encoder suitable for Elasticsearch/OpenSearch-style retrieval.
+- Keep vectors sparse enough to index in practice.
+- Use the user's existing Hugging Face dataset and local Apple Silicon machine.
+- Avoid overclaiming from pairwise training accuracy alone.
+
+## Base model decision
+
+Chosen base model:
+
+```text
+opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill
+```
+
+Reason: it already has separate query/document sparse encoding behavior and is aligned with OpenSearch neural sparse retrieval. Starting here means fine-tuning adapts a serving-compatible sparse model rather than building a new retrieval stack from scratch.
+
+## Dataset
+
+Dataset:
+
+```text
+oneryalcin/financial-filings-sparse-retrieval-training
+```
+
+Config:
+
+```text
+combined
+```
+
+Each usable row is treated as:
+
+```text
+query
+positive filing chunk
+first non-empty hard negative filing chunk
+```
+
+Training used up to 20,000 requested rows. After dropping rows without a non-empty negative, there were 18,247 usable training triplets.
+
+Evaluation used the first 1,000 usable rows from the held-out `test` split.
+
+Known dataset fields include `query`, `positive`, `negatives`, `query_type`, `company`, and `doc_type`.
+
+## Training recipe
+
+Model type:
+
+```text
+sentence_transformers.sparse_encoder.SparseEncoder
+```
+
+Loss:
+
+```text
+SpladeLoss(
+  SparseMultipleNegativesRankingLoss,
+  query_regularizer_weight=1e-4,
+  document_regularizer_weight=1e-2,
+)
+```
+
+Important settings:
+
+| setting | value |
+|---|---:|
+| max steps | 1500 |
+| batch size | 8 |
+| train rows requested | 20,000 |
+| usable train rows | 18,247 |
+| query regularization | `1e-4` |
+| document regularization | `1e-2` |
+| sampler | no-duplicates batch sampler |
+| platform | local Apple Silicon, MPS where supported |
+| package cutoff | `uv run --exclude-newer=2026-05-13T00:00:00Z` |
+
+Why this recipe:
+
+- `SparseMultipleNegativesRankingLoss` teaches the model to rank the paired positive above in-batch negatives.
+- `SpladeLoss` adds FLOPS-style sparse regularization so the model does not emit unindexably dense document vectors.
+- A no-duplicates sampler avoids false negatives inside contrastive batches.
+- Document regularization matters more than query regularization for index footprint.
+
+## Decision log
+
+| decision point | choice | why |
+|---|---|---|
+| Retrieval family | `SparseEncoder` / SPLADE-style sparse retrieval | Best fit for Elasticsearch/OpenSearch sparse retrieval. |
+| Base model | OpenSearch sparse doc v2 distill | Already query/document routed and serving-aligned. |
+| Vocab | Keep base tokenizer/vocab | Changing vocabulary would be a larger pretraining-style project, not a fine-tune. |
+| Loss | `SpladeLoss + SparseMultipleNegativesRankingLoss` | Combines ranking pressure with sparse regularization. |
+| Sampler | no-duplicates | Reduces accidental false negatives in batch. |
+| Doc regularization | increased to `1e-2` | Lower regularization gave high scores but thousands of active doc terms. |
+| Serving vector size | `top_k=128` | Preserved almost all quality with much smaller document vectors. |
+| Final run length | 1500 steps | 500-step runs were promising; longer training improved both triplet and retrieval proxy scores. |
+| Evaluation strategy | triplet + retrieval proxy + base sparse + BM25 | Avoids relying on a single easy metric. |
+
+## Regularization sweep
+
+The first 500-step sweep showed the main tradeoff. Low document regularization produced good triplet accuracy but extremely dense document vectors.
+
+| doc reg | triplet accuracy | positive doc dims | interpretation |
+|---:|---:|---:|---|
+| `8e-5` | 73.8% | 4204.9 | Too dense. |
+| `1.5e-4` | 74.8% | 3214.8 | Best raw 500-step accuracy, still too dense. |
+| `3e-4` | 73.4% | 2101.4 | Still too dense. |
+| `1e-3` | 72.7% | 1340.5 | Better, still high. |
+| `3e-3` | 73.4% | 657.1 | Practical direction. |
+| `1e-2` | 73.2% | 296.7 | Chosen for longer run. |
+| `3e-2` | 72.3% | 160.9 | Compact, lower quality. |
+
+The final model uses `document_regularizer_weight=1e-2` and 1500 steps.
+
+## Evaluation 1: triplet accuracy
+
+For each held-out row:
+
+```text
+score(query, positive) > score(query, hard_negative)
+```
+
+Metric:
+
+```text
+accuracy = fraction of rows where positive scores higher than the paired negative
+```
+
+| model | doc pruning | accuracy | mean margin | query dims | positive doc dims |
+|---|---|---:|---:|---:|---:|
+| base sparse | unpruned | 49.6% | 0.451 | 15.2 | 371.1 |
+| fine-tuned sparse | unpruned | 78.1% | 3.004 | 15.2 | 340.9 |
+| fine-tuned sparse | top-128 | 78.0% | 2.992 | 15.2 | 126.7 |
+| fine-tuned sparse | top-64 | 75.9% | 2.988 | 15.2 | 64.0 |
+
+Interpretation: the fine-tuned model learned the domain signal strongly. Top-128 preserved almost all pairwise quality.
+
+## Evaluation 2: in-memory retrieval proxy
+
+A retrieval candidate pool was built from the held-out test split:
+
+- 1,000 held-out queries
+- all unique positives from those rows
+- all unique first hard negatives from those rows
+- 1,912 unique candidate chunks total
+
+For each query, the known positive chunk is the only labeled relevant document. The model ranks all 1,912 candidate chunks by sparse dot product. BM25 ranks the same candidate corpus with a local lexical implementation.
+
+Metrics:
+
+- `Recall@1`
+- `Recall@5`
+- `Recall@10`
+- `Recall@20`
+- `MRR@10`
+- `nDCG@10`
+- median rank
+
+| model | pruning | Recall@1 | Recall@5 | Recall@10 | Recall@20 | MRR@10 | nDCG@10 | median rank |
+|---|---|---:|---:|---:|---:|---:|---:|---:|
+| fine-tuned sparse | unpruned | 39.0% | 58.9% | 67.5% | 75.0% | 0.479 | 0.526 | 3 |
+| fine-tuned sparse | top-128 | 38.6% | 57.8% | 67.2% | 73.9% | 0.473 | 0.521 | 3 |
+| fine-tuned sparse | top-64 | 35.0% | 55.7% | 64.8% | 72.4% | 0.442 | 0.491 | 4 |
+| base sparse | unpruned | 32.1% | 56.5% | 63.7% | 69.9% | 0.431 | 0.481 | 3 |
+| base sparse | top-128 | 31.4% | 54.4% | 62.7% | 69.7% | 0.422 | 0.472 | 3 |
+| base sparse | top-64 | 29.5% | 52.6% | 59.5% | 66.5% | 0.396 | 0.444 | 4 |
+| BM25 | lexical | 24.0% | 58.2% | 64.1% | 68.6% | 0.397 | 0.457 | 3 |
+
+Interpretation:
+
+- Top-128 fine-tuned sparse is the best current deployment candidate.
+- BM25 is competitive at Recall@10, but worse at early precision and ranking quality.
+- Fine-tuning mostly improved early ranking: Recall@1 rose from 32.1% for the base sparse model to 39.0% unpruned, and 38.6% at top-128.
+- Top-64 is usable only if index size or latency dominates quality.
+
+## Recommended serving configuration
+
+Use document top-k pruning:
+
+```text
+document_top_k = 128
+```
+
+Rationale:
+
+| setting | Recall@10 | nDCG@10 | doc active dims |
+|---|---:|---:|---:|
+| unpruned | 67.5% | 0.526 | 319.8 |
+| top-128 | 67.2% | 0.521 | 126.5 |
+| top-64 | 64.8% | 0.491 | 64.0 |
+
+Top-128 gives almost the same retrieval quality as unpruned with a much smaller sparse index footprint.
+
+## Usage
+
+```python
+from sentence_transformers.sparse_encoder import SparseEncoder
+
+model = SparseEncoder("oneryalcin/financial-filings-sparse-encoder-v1")
+
+query_vectors = model.encode_query([
+    "What does the company say about liquidity risk?"
+])
+
+document_vectors = model.encode_document([
+    "The company discusses liquidity and capital resources in the MD&A section..."
+])
+```
+
+For production-style sparse indexing, keep the highest-weighted 128 dimensions per document vector before indexing.
+
+## Reproduction
+
+The local runs used `uv` with a package-date cutoff:
+
+```bash
+uv run --exclude-newer=2026-05-13T00:00:00Z scripts/train_fin_sparse_encoder_v2.py \
+  --train-size 20000 \
+  --eval-size 1000 \
+  --max-steps 1500 \
+  --batch-size 8 \
+  --eval-batch-size 16 \
+  --query-reg 1e-4 \
+  --doc-reg 1e-2 \
+  --run-name fin-sparse-encoder-doc-v2-reg1e-2-1500
+```
+
+Triplet top-k evaluation:
+
+```bash
+uv run --exclude-newer=2026-05-13T00:00:00Z scripts/eval_fin_sparse_topk.py \
+  --eval-size 1000 \
+  --batch-size 16 \
+  --topks none,128,64 \
+  --models ./
+```
+
+Retrieval proxy evaluation:
+
+```bash
+uv run --exclude-newer=2026-05-13T00:00:00Z scripts/eval_fin_sparse_retrieval_proxy.py \
+  --model ./ \
+  --eval-size 1000 \
+  --batch-size 16 \
+  --topks none,128,64
+```
+
+BM25 baseline:
+
+```bash
+uv run --exclude-newer=2026-05-13T00:00:00Z scripts/eval_fin_bm25_retrieval_proxy.py \
+  --eval-size 1000
+```
+
+## Files in this repo
+
+| path | purpose |
+|---|---|
+| `scripts/train_fin_sparse_encoder_v2.py` | Training script used for the final run. |
+| `scripts/eval_fin_sparse_topk.py` | Pairwise triplet accuracy with document top-k pruning. |
+| `scripts/eval_fin_sparse_retrieval_proxy.py` | In-memory retrieval proxy over held-out positives and negatives. |
+| `scripts/eval_fin_bm25_retrieval_proxy.py` | Local BM25 baseline over the same held-out candidate pool. |
+| `benchmarks/regularization_sweep_500_steps.json` | 500-step document-regularization sweep. |
+| `benchmarks/triplet_topk_1000.json` | Final triplet/top-k results. |
+| `benchmarks/retrieval_proxy_1000.json` | Final retrieval proxy and baseline results. |
+
+## Limitations
+
+These results are promising but not final production proof.
+
+- The retrieval benchmark is a proxy, not a full production OpenSearch/Elasticsearch benchmark.
+- Each query has only one labeled positive, so other genuinely relevant chunks may be counted as false competitors.
+- The candidate pool has 1,912 chunks, not millions.
+- Evaluation used the first 1,000 usable held-out examples, not the full test split.
+- BM25 is a local baseline, not a tuned OpenSearch BM25 setup.
+- End-to-end index size, latency, shard behavior, and hybrid retrieval have not been measured yet.
+
+## Next steps
+
+- Run the retrieval proxy on the full held-out test split.
+- Run a real OpenSearch/Elasticsearch index benchmark with top-128 sparse vectors.
+- Test hybrid BM25 + learned sparse retrieval.
+- Add query-type, company, and filing-type slice analysis.
+- Add human-labeled financial filing retrieval judgments.

benchmarks/regularization_sweep_500_steps.json

@@ -0,0 +1,25 @@
+{
+  "dataset": "oneryalcin/financial-filings-sparse-retrieval-training",
+  "config": "combined",
+  "split": "test",
+  "rows": 1000,
+  "base_model": "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill",
+  "constant_settings": {
+    "steps": 500,
+    "train_size_requested": 20000,
+    "usable_train_rows": 18247,
+    "batch_size": 8,
+    "query_regularizer_weight": 0.0001,
+    "loss": "SpladeLoss(SparseMultipleNegativesRankingLoss)",
+    "sampler": "BatchSamplers.NO_DUPLICATES"
+  },
+  "results": [
+    {"document_regularizer_weight": 0.00008, "accuracy": 0.738, "positive_doc_active_dims": 4204.911, "negative_doc_active_dims": 3784.729, "note": "High score but too dense for practical sparse indexing."},
+    {"document_regularizer_weight": 0.00015, "accuracy": 0.748, "positive_doc_active_dims": 3214.797, "negative_doc_active_dims": 2956.645, "note": "Best raw 500-step triplet accuracy but still too dense."},
+    {"document_regularizer_weight": 0.0003, "accuracy": 0.734, "positive_doc_active_dims": 2101.448, "negative_doc_active_dims": 1875.179},
+    {"document_regularizer_weight": 0.001, "accuracy": 0.727, "positive_doc_active_dims": 1340.540, "negative_doc_active_dims": 1136.954},
+    {"document_regularizer_weight": 0.003, "accuracy": 0.734, "positive_doc_active_dims": 657.067, "negative_doc_active_dims": 563.964},
+    {"document_regularizer_weight": 0.01, "accuracy": 0.732, "positive_doc_active_dims": 296.654, "negative_doc_active_dims": 261.160, "note": "Chosen for longer run because quality stayed strong while native vectors became much more indexable."},
+    {"document_regularizer_weight": 0.03, "accuracy": 0.723, "positive_doc_active_dims": 160.887, "negative_doc_active_dims": 143.986, "note": "Compact but lower quality than 0.01."}
+  ]
+}

benchmarks/retrieval_proxy_1000.json

@@ -0,0 +1,109 @@
+{
+  "dataset": "oneryalcin/financial-filings-sparse-retrieval-training",
+  "config": "combined",
+  "split": "test",
+  "rows": 1000,
+  "corpus_size": 1912,
+  "method": "Build an in-memory candidate corpus from all unique positives and first hard negatives in the 1000 held-out rows. For each query, rank all 1912 candidate chunks by sparse dot product or BM25 score. Only the known positive is labeled relevant.",
+  "results": [
+    {
+      "model": "oneryalcin/financial-filings-sparse-encoder-v1",
+      "doc_topk": null,
+      "doc_active_dims": 319.79864501953125,
+      "query_active_dims": 15.227999687194824,
+      "recall_at_1": 0.38999998569488525,
+      "recall_at_5": 0.5889999866485596,
+      "recall_at_10": 0.675000011920929,
+      "recall_at_20": 0.75,
+      "mrr_at_10": 0.479160338640213,
+      "ndcg_at_10": 0.5258780717849731,
+      "mean_rank": 51.290000915527344,
+      "median_rank": 3.0
+    },
+    {
+      "model": "oneryalcin/financial-filings-sparse-encoder-v1",
+      "doc_topk": 128,
+      "doc_active_dims": 126.48535919189453,
+      "query_active_dims": 15.227999687194824,
+      "recall_at_1": 0.38600000739097595,
+      "recall_at_5": 0.578000009059906,
+      "recall_at_10": 0.671999990940094,
+      "recall_at_20": 0.7390000224113464,
+      "mrr_at_10": 0.4734043478965759,
+      "ndcg_at_10": 0.5206189155578613,
+      "mean_rank": 52.779998779296875,
+      "median_rank": 3.0
+    },
+    {
+      "model": "oneryalcin/financial-filings-sparse-encoder-v1",
+      "doc_topk": 64,
+      "doc_active_dims": 63.98692321777344,
+      "query_active_dims": 15.227999687194824,
+      "recall_at_1": 0.3499999940395355,
+      "recall_at_5": 0.5569999814033508,
+      "recall_at_10": 0.6480000019073486,
+      "recall_at_20": 0.7239999771118164,
+      "mrr_at_10": 0.4424420893192291,
+      "ndcg_at_10": 0.4913859963417053,
+      "mean_rank": 56.25199890136719,
+      "median_rank": 4.0
+    },
+    {
+      "model": "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill",
+      "doc_topk": null,
+      "doc_active_dims": 377.7003173828125,
+      "query_active_dims": 15.227999687194824,
+      "recall_at_1": 0.32100000977516174,
+      "recall_at_5": 0.5649999976158142,
+      "recall_at_10": 0.6370000243186951,
+      "recall_at_20": 0.6990000009536743,
+      "mrr_at_10": 0.4311789870262146,
+      "ndcg_at_10": 0.4811413586139679,
+      "mean_rank": 98.5719985961914,
+      "median_rank": 3.0
+    },
+    {
+      "model": "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill",
+      "doc_topk": 128,
+      "doc_active_dims": 127.96809387207031,
+      "query_active_dims": 15.227999687194824,
+      "recall_at_1": 0.3140000104904175,
+      "recall_at_5": 0.5440000295639038,
+      "recall_at_10": 0.6269999742507935,
+      "recall_at_20": 0.6970000267028809,
+      "mrr_at_10": 0.4224868714809418,
+      "ndcg_at_10": 0.4720706045627594,
+      "mean_rank": 95.6719970703125,
+      "median_rank": 3.0
+    },
+    {
+      "model": "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill",
+      "doc_topk": 64,
+      "doc_active_dims": 64.0,
+      "query_active_dims": 15.227999687194824,
+      "recall_at_1": 0.29499998688697815,
+      "recall_at_5": 0.5260000228881836,
+      "recall_at_10": 0.5950000286102295,
+      "recall_at_20": 0.6650000214576721,
+      "mrr_at_10": 0.3963717818260193,
+      "ndcg_at_10": 0.44444921612739563,
+      "mean_rank": 96.0739974975586,
+      "median_rank": 4.0
+    },
+    {
+      "model": "bm25",
+      "recall_at_1": 0.24,
+      "recall_at_5": 0.582,
+      "recall_at_10": 0.641,
+      "recall_at_20": 0.686,
+      "mrr_at_10": 0.3968968253968254,
+      "ndcg_at_10": 0.45737284741632095,
+      "mean_rank": 158.692,
+      "median_rank": 3,
+      "k1": 1.2,
+      "b": 0.75,
+      "avg_doc_len": 213.8655857740586,
+      "median_doc_len": 222
+    }
+  ]
+}

benchmarks/triplet_topk_1000.json

@@ -0,0 +1,48 @@
+{
+  "dataset": "oneryalcin/financial-filings-sparse-retrieval-training",
+  "config": "combined",
+  "split": "test",
+  "rows": 1000,
+  "method": "For each held-out row, score(query, positive) and score(query, first hard negative) with sparse dot product. Accuracy is the fraction where positive scores higher.",
+  "results": [
+    {
+      "model": "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill",
+      "doc_topk": null,
+      "accuracy": 0.4960,
+      "mean_margin": 0.4508,
+      "query_active_dims": 15.228,
+      "positive_doc_active_dims": 371.084,
+      "negative_doc_active_dims": 354.758
+    },
+    {
+      "model": "oneryalcin/financial-filings-sparse-encoder-v1",
+      "doc_topk": null,
+      "accuracy": 0.781000018119812,
+      "mean_margin": 3.0043423175811768,
+      "median_margin": 2.4458980560302734,
+      "query_active_dims": 15.227999687194824,
+      "positive_doc_active_dims": 340.94500732421875,
+      "negative_doc_active_dims": 298.7919921875
+    },
+    {
+      "model": "oneryalcin/financial-filings-sparse-encoder-v1",
+      "doc_topk": 128,
+      "accuracy": 0.7799999713897705,
+      "mean_margin": 2.9922006130218506,
+      "median_margin": 2.4021129608154297,
+      "query_active_dims": 15.227999687194824,
+      "positive_doc_active_dims": 126.6719970703125,
+      "negative_doc_active_dims": 126.36499786376953
+    },
+    {
+      "model": "oneryalcin/financial-filings-sparse-encoder-v1",
+      "doc_topk": 64,
+      "accuracy": 0.7590000033378601,
+      "mean_margin": 2.988410234451294,
+      "median_margin": 2.3834714889526367,
+      "query_active_dims": 15.227999687194824,
+      "positive_doc_active_dims": 63.986000061035156,
+      "negative_doc_active_dims": 63.98899841308594
+    }
+  ]
+}

config_sentence_transformers.json

@@ -0,0 +1,14 @@
+{
+  "__version__": {
+    "pytorch": "2.11.0",
+    "sentence_transformers": "5.5.0",
+    "transformers": "5.8.0"
+  },
+  "default_prompt_name": null,
+  "model_type": "SparseEncoder",
+  "prompts": {
+    "document": "",
+    "query": ""
+  },
+  "similarity_fn_name": "dot"
+}

document_0_MLMTransformer/config.json

@@ -0,0 +1,27 @@
+{
+  "activation": "gelu",
+  "architectures": [
+    "DistilBertForMaskedLM"
+  ],
+  "attention_dropout": 0.1,
+  "bos_token_id": null,
+  "dim": 768,
+  "dropout": 0.1,
+  "dtype": "float32",
+  "eos_token_id": null,
+  "hidden_dim": 3072,
+  "initializer_range": 0.02,
+  "max_position_embeddings": 512,
+  "model_type": "distilbert",
+  "n_heads": 12,
+  "n_layers": 6,
+  "pad_token_id": 0,
+  "qa_dropout": 0.1,
+  "seq_classif_dropout": 0.2,
+  "sinusoidal_pos_embds": false,
+  "tie_weights_": true,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.8.0",
+  "use_cache": false,
+  "vocab_size": 30522
+}

document_0_MLMTransformer/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:74310f44e4ea4067d85d4ab39161a1eed74b96ead822dea488bdbda9d034a304
+size 267954768

document_0_MLMTransformer/sentence_bert_config.json

@@ -0,0 +1,10 @@
+{
+    "transformer_task": "fill-mask",
+    "modality_config": {
+        "text": {
+            "method": "forward",
+            "method_output_name": "logits"
+        }
+    },
+    "module_output_name": "token_embeddings"
+}

document_0_MLMTransformer/tokenizer_config.json

@@ -0,0 +1,16 @@
+{
+  "backend": "tokenizers",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "do_lower_case": true,
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "[MASK]",
+  "model_max_length": 384,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "DistilBertTokenizer",
+  "unk_token": "[UNK]"
+}

document_1_SpladePooling/config.json

@@ -0,0 +1,5 @@
+{
+    "pooling_strategy": "max",
+    "activation_function": "relu",
+    "embedding_dimension": 30522
+}

modules.json

@@ -0,0 +1,8 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.base.modules.router.Router"
+  }
+]

query_0_SparseStaticEmbedding/config.json

@@ -0,0 +1,3 @@
+{
+    "frozen": true
+}

query_0_SparseStaticEmbedding/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:711ec64837a7962d2ae106996079782b7ee87860089a0b2348bf7cb840f252d3
+size 122168

query_0_SparseStaticEmbedding/tokenizer_config.json

@@ -0,0 +1,16 @@
+{
+  "backend": "tokenizers",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "do_lower_case": true,
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "[MASK]",
+  "model_max_length": 512,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "DistilBertTokenizer",
+  "unk_token": "[UNK]"
+}

router_config.json

@@ -0,0 +1,21 @@
+{
+    "types": {
+        "query_0_SparseStaticEmbedding": "sentence_transformers.sparse_encoder.modules.sparse_static_embedding.SparseStaticEmbedding",
+        "document_0_MLMTransformer": "sentence_transformers.sparse_encoder.modules.mlm_transformer.MLMTransformer",
+        "document_1_SpladePooling": "sentence_transformers.sparse_encoder.modules.splade_pooling.SpladePooling"
+    },
+    "structure": {
+        "query": [
+            "query_0_SparseStaticEmbedding"
+        ],
+        "document": [
+            "document_0_MLMTransformer",
+            "document_1_SpladePooling"
+        ]
+    },
+    "parameters": {
+        "default_route": "document",
+        "allow_empty_key": true,
+        "route_mappings": {}
+    }
+}

scripts/eval_fin_bm25_retrieval_proxy.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#   "datasets",
+# ]
+# ///
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import math
+import re
+from collections import Counter, OrderedDict, defaultdict
+
+from datasets import load_dataset
+
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(message)s")
+TOKEN_RE = re.compile(r"[a-z0-9]+")
+
+
+def tokenize(text: str) -> list[str]:
+    return TOKEN_RE.findall(text.lower())
+
+
+def first_negative(row: dict) -> str | None:
+    negatives = row.get("negatives")
+    if isinstance(negatives, list):
+        for negative in negatives:
+            if isinstance(negative, str) and negative.strip():
+                return negative
+    if isinstance(negatives, str) and negatives.strip():
+        return negatives
+    return None
+
+
+def load_rows(limit: int) -> list[dict]:
+    dataset = load_dataset(
+        "oneryalcin/financial-filings-sparse-retrieval-training",
+        "combined",
+        split="test",
+    )
+    rows: list[dict] = []
+    for row in dataset:
+        query = row.get("query")
+        positive = row.get("positive")
+        negative = first_negative(row)
+        if query and positive and negative:
+            rows.append({"query": query, "positive": positive, "negative": negative})
+        if len(rows) >= limit:
+            break
+    return rows
+
+
+def dedupe_texts(rows: list[dict]) -> tuple[list[str], list[int]]:
+    corpus = OrderedDict()
+    positive_ids: list[int] = []
+    for row in rows:
+        for key in ("positive", "negative"):
+            text = row[key]
+            if text not in corpus:
+                corpus[text] = len(corpus)
+        positive_ids.append(corpus[row["positive"]])
+    return list(corpus.keys()), positive_ids
+
+
+def build_bm25(corpus: list[str], k1: float, b: float) -> tuple[dict[str, list[tuple[int, float]]], list[int], float]:
+    doc_lens: list[int] = []
+    doc_counts: list[Counter[str]] = []
+    document_frequency: Counter[str] = Counter()
+
+    for text in corpus:
+        counts = Counter(tokenize(text))
+        doc_counts.append(counts)
+        doc_lens.append(sum(counts.values()))
+        document_frequency.update(counts.keys())
+
+    corpus_size = len(corpus)
+    avg_doc_len = sum(doc_lens) / max(corpus_size, 1)
+    inverted: dict[str, list[tuple[int, float]]] = defaultdict(list)
+
+    for doc_id, counts in enumerate(doc_counts):
+        norm = k1 * (1 - b + b * doc_lens[doc_id] / avg_doc_len)
+        for term, tf in counts.items():
+            df = document_frequency[term]
+            idf = math.log(1 + (corpus_size - df + 0.5) / (df + 0.5))
+            score = idf * (tf * (k1 + 1)) / (tf + norm)
+            inverted[term].append((doc_id, score))
+
+    return inverted, doc_lens, avg_doc_len
+
+
+def rank_query(query: str, inverted: dict[str, list[tuple[int, float]]], corpus_size: int) -> list[float]:
+    scores = [0.0] * corpus_size
+    for term, qtf in Counter(tokenize(query)).items():
+        for doc_id, score in inverted.get(term, ()):
+            scores[doc_id] += score * qtf
+    return scores
+
+
+def retrieval_metrics(all_scores: list[list[float]], positive_ids: list[int]) -> dict:
+    ranks: list[int] = []
+    for scores, positive_id in zip(all_scores, positive_ids):
+        positive_score = scores[positive_id]
+        rank = 1 + sum(score > positive_score for score in scores)
+        ranks.append(rank)
+
+    ranks_sorted = sorted(ranks)
+    count = len(ranks)
+    mrr10 = sum((1.0 / rank) if rank <= 10 else 0.0 for rank in ranks) / count
+    ndcg10 = sum((1.0 / math.log2(rank + 1)) if rank <= 10 else 0.0 for rank in ranks) / count
+    return {
+        "recall_at_1": sum(rank <= 1 for rank in ranks) / count,
+        "recall_at_5": sum(rank <= 5 for rank in ranks) / count,
+        "recall_at_10": sum(rank <= 10 for rank in ranks) / count,
+        "recall_at_20": sum(rank <= 20 for rank in ranks) / count,
+        "mrr_at_10": mrr10,
+        "ndcg_at_10": ndcg10,
+        "mean_rank": sum(ranks) / count,
+        "median_rank": ranks_sorted[count // 2],
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--eval-size", type=int, default=1000)
+    parser.add_argument("--k1", type=float, default=1.2)
+    parser.add_argument("--b", type=float, default=0.75)
+    args = parser.parse_args()
+
+    rows = load_rows(args.eval_size)
+    queries = [row["query"] for row in rows]
+    corpus, positive_ids = dedupe_texts(rows)
+    logging.info("rows=%d corpus=%d k1=%s b=%s", len(rows), len(corpus), args.k1, args.b)
+
+    inverted, doc_lens, avg_doc_len = build_bm25(corpus, args.k1, args.b)
+    scores = [rank_query(query, inverted, len(corpus)) for query in queries]
+    result = {
+        "model": "bm25",
+        "rows": len(rows),
+        "corpus_size": len(corpus),
+        "avg_doc_len": avg_doc_len,
+        "median_doc_len": sorted(doc_lens)[len(doc_lens) // 2],
+        "k1": args.k1,
+        "b": args.b,
+        **retrieval_metrics(scores, positive_ids),
+    }
+    logging.info("BM25_RESULT %s", json.dumps(result, sort_keys=True))
+    print(json.dumps(result, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()

scripts/eval_fin_sparse_retrieval_proxy.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#   "datasets",
+#   "sentence-transformers==5.5.0",
+#   "torch",
+# ]
+# ///
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+from collections import OrderedDict
+
+import torch
+from datasets import load_dataset
+from sentence_transformers.sparse_encoder import SparseEncoder
+
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(message)s")
+
+
+def parse_topks(value: str) -> list[int | None]:
+    out: list[int | None] = []
+    for item in value.split(","):
+        item = item.strip().lower()
+        if item in {"none", "all", "null"}:
+            out.append(None)
+        else:
+            out.append(int(item))
+    return out
+
+
+def first_negative(row: dict) -> str | None:
+    negatives = row.get("negatives")
+    if isinstance(negatives, list):
+        for negative in negatives:
+            if isinstance(negative, str) and negative.strip():
+                return negative
+    if isinstance(negatives, str) and negatives.strip():
+        return negatives
+    return None
+
+
+def load_rows(limit: int) -> list[dict]:
+    dataset = load_dataset(
+        "oneryalcin/financial-filings-sparse-retrieval-training",
+        "combined",
+        split="test",
+    )
+    rows: list[dict] = []
+    for row in dataset:
+        query = row.get("query")
+        positive = row.get("positive")
+        negative = first_negative(row)
+        if query and positive and negative:
+            rows.append({"query": query, "positive": positive, "negative": negative})
+        if len(rows) >= limit:
+            break
+    return rows
+
+
+def dedupe_texts(rows: list[dict]) -> tuple[list[str], list[int], list[int]]:
+    corpus = OrderedDict()
+    positive_ids: list[int] = []
+    negative_ids: list[int] = []
+    for row in rows:
+        for key in ("positive", "negative"):
+            text = row[key]
+            if text not in corpus:
+                corpus[text] = len(corpus)
+        positive_ids.append(corpus[row["positive"]])
+        negative_ids.append(corpus[row["negative"]])
+    return list(corpus.keys()), positive_ids, negative_ids
+
+
+def densify(encoded) -> torch.Tensor:
+    if isinstance(encoded, torch.Tensor):
+        tensor = encoded
+        if tensor.layout != torch.strided:
+            tensor = tensor.to_dense()
+    elif hasattr(encoded, "to_dense"):
+        tensor = encoded.to_dense()
+    else:
+        tensor = torch.as_tensor(encoded)
+    return tensor.float().cpu()
+
+
+def active_dims(tensor: torch.Tensor) -> float:
+    return (tensor > 0).sum(dim=1).float().mean().item()
+
+
+def apply_doc_topk(docs: torch.Tensor, topk: int | None) -> torch.Tensor:
+    if topk is None or topk >= docs.shape[1]:
+        return docs
+    values, indices = torch.topk(docs, k=topk, dim=1)
+    pruned = torch.zeros_like(docs)
+    pruned.scatter_(1, indices, values)
+    return pruned
+
+
+def retrieval_metrics(scores: torch.Tensor, positive_ids: list[int]) -> dict:
+    positive = torch.tensor(positive_ids, dtype=torch.long)
+    positive_scores = scores[torch.arange(scores.shape[0]), positive]
+    ranks = (scores > positive_scores.unsqueeze(1)).sum(dim=1) + 1
+    ranks_f = ranks.float()
+
+    mrr10 = torch.where(ranks <= 10, 1.0 / ranks_f, torch.zeros_like(ranks_f)).mean().item()
+    ndcg10 = torch.where(
+        ranks <= 10,
+        1.0 / torch.log2(ranks_f + 1.0),
+        torch.zeros_like(ranks_f),
+    ).mean().item()
+    return {
+        "recall_at_1": (ranks <= 1).float().mean().item(),
+        "recall_at_5": (ranks <= 5).float().mean().item(),
+        "recall_at_10": (ranks <= 10).float().mean().item(),
+        "recall_at_20": (ranks <= 20).float().mean().item(),
+        "mrr_at_10": mrr10,
+        "ndcg_at_10": ndcg10,
+        "mean_rank": ranks_f.mean().item(),
+        "median_rank": ranks_f.median().item(),
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", required=True)
+    parser.add_argument("--eval-size", type=int, default=1000)
+    parser.add_argument("--batch-size", type=int, default=16)
+    parser.add_argument("--topks", default="none,128,64")
+    args = parser.parse_args()
+
+    rows = load_rows(args.eval_size)
+    queries = [row["query"] for row in rows]
+    corpus, positive_ids, negative_ids = dedupe_texts(rows)
+    logging.info("rows=%d corpus=%d topks=%s", len(rows), len(corpus), args.topks)
+
+    model = SparseEncoder(args.model, device="cpu")
+    query_vectors = densify(
+        model.encode_query(queries, batch_size=args.batch_size, convert_to_tensor=True)
+    )
+    doc_vectors = densify(
+        model.encode_document(corpus, batch_size=args.batch_size, convert_to_tensor=True)
+    )
+
+    query_active = active_dims(query_vectors)
+    doc_active = active_dims(doc_vectors)
+
+    for topk in parse_topks(args.topks):
+        pruned_docs = apply_doc_topk(doc_vectors, topk)
+        scores = query_vectors @ pruned_docs.T
+        metrics = retrieval_metrics(scores, positive_ids)
+        result = {
+            "model": args.model,
+            "rows": len(rows),
+            "corpus_size": len(corpus),
+            "doc_topk": topk,
+            "query_active_dims": query_active,
+            "doc_active_dims": active_dims(pruned_docs),
+            **metrics,
+        }
+        logging.info("RETRIEVAL_RESULT %s", json.dumps(result, sort_keys=True))
+        print(json.dumps(result, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()

scripts/eval_fin_sparse_topk.py

@@ -0,0 +1,144 @@
+#!/usr/bin/env python3
+# /// script
+# dependencies = [
+#   "sentence-transformers[train]==5.5.0",
+#   "datasets",
+#   "torch",
+# ]
+# ///
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+from dataclasses import dataclass
+
+import torch
+from datasets import load_dataset
+from sentence_transformers.sparse_encoder import SparseEncoder
+
+
+DATASET = "oneryalcin/financial-filings-sparse-retrieval-training"
+CONFIG = "combined"
+
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(message)s")
+for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
+    logging.getLogger(noisy).setLevel(logging.WARNING)
+
+
+@dataclass(frozen=True)
+class EvalBatch:
+    queries: list[str]
+    positives: list[str]
+    negatives: list[str]
+
+
+def dense(tensor: torch.Tensor) -> torch.Tensor:
+    if tensor.is_sparse:
+        return tensor.to_dense()
+    return tensor
+
+
+def doc_topk(tensor: torch.Tensor, k: int | None) -> torch.Tensor:
+    if k is None or k <= 0 or k >= tensor.shape[1]:
+        return tensor
+    values, indices = torch.topk(tensor, k=k, dim=1)
+    out = torch.zeros_like(tensor)
+    return out.scatter(1, indices, values)
+
+
+def active_dims(tensor: torch.Tensor) -> float:
+    return (tensor != 0).sum(dim=1).float().mean().item()
+
+
+def margins(queries: torch.Tensor, positives: torch.Tensor, negatives: torch.Tensor) -> torch.Tensor:
+    positive_scores = (queries * positives).sum(dim=1)
+    negative_scores = (queries * negatives).sum(dim=1)
+    return positive_scores - negative_scores
+
+
+def load_eval_rows(limit: int) -> EvalBatch:
+    rows = load_dataset(DATASET, CONFIG, split="test", streaming=False)
+    queries: list[str] = []
+    positives: list[str] = []
+    negatives: list[str] = []
+
+    for row in rows:
+        negatives_list = row.get("negatives") or []
+        if not negatives_list:
+            continue
+        queries.append(row["query"])
+        positives.append(row["positive"])
+        negatives.append(negatives_list[0])
+        if len(queries) >= limit:
+            break
+
+    return EvalBatch(queries=queries, positives=positives, negatives=negatives)
+
+
+def eval_model(model_name: str, batch: EvalBatch, topks: list[int | None], batch_size: int) -> list[dict[str, float | int | str | None]]:
+    logging.info("Loading %s", model_name)
+    model = SparseEncoder(model_name, trust_remote_code=True, device="cpu")
+
+    logging.info("Encoding %d triplets", len(batch.queries))
+    queries = dense(model.encode_query(batch.queries, batch_size=batch_size, convert_to_tensor=True, show_progress_bar=True)).cpu()
+    positives_raw = dense(model.encode_document(batch.positives, batch_size=batch_size, convert_to_tensor=True, show_progress_bar=True)).cpu()
+    negatives_raw = dense(model.encode_document(batch.negatives, batch_size=batch_size, convert_to_tensor=True, show_progress_bar=True)).cpu()
+
+    results: list[dict[str, float | int | str | None]] = []
+    for k in topks:
+        positives = doc_topk(positives_raw, k)
+        negatives = doc_topk(negatives_raw, k)
+        margin = margins(queries, positives, negatives)
+        results.append(
+            {
+                "model": model_name,
+                "doc_topk": k,
+                "accuracy": (margin > 0).float().mean().item(),
+                "mean_margin": margin.mean().item(),
+                "median_margin": margin.median().item(),
+                "query_active_dims": active_dims(queries),
+                "positive_doc_active_dims": active_dims(positives),
+                "negative_doc_active_dims": active_dims(negatives),
+            }
+        )
+
+    return results
+
+
+def parse_topks(raw: str) -> list[int | None]:
+    out: list[int | None] = []
+    for item in raw.split(","):
+        item = item.strip().lower()
+        if item in {"none", "all", "0"}:
+            out.append(None)
+        else:
+            out.append(int(item))
+    return out
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--models", nargs="+", required=True)
+    parser.add_argument("--eval-size", type=int, default=1000)
+    parser.add_argument("--batch-size", type=int, default=16)
+    parser.add_argument("--topks", default="none,1024,512,256,128,64")
+    args = parser.parse_args()
+
+    eval_batch = load_eval_rows(args.eval_size)
+    topks = parse_topks(args.topks)
+    logging.info("rows=%d topks=%s", len(eval_batch.queries), topks)
+
+    all_results: list[dict[str, float | int | str | None]] = []
+    for model_name in args.models:
+        all_results.extend(eval_model(model_name, eval_batch, topks, args.batch_size))
+
+    for result in all_results:
+        logging.info("TOPK_RESULT %s", json.dumps(result, sort_keys=True))
+
+    print(json.dumps(all_results, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()

scripts/train_fin_sparse_encoder_v2.py

@@ -0,0 +1,238 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "sentence-transformers[train]==5.5.0",
+#     "datasets>=2.19.0",
+#     "accelerate>=0.26.0",
+# ]
+# ///
+"""Financial filings SparseEncoder V2.
+
+Design choices:
+- stable Apache/OpenSearch base: doc-v2-distill, not brittle custom v3-gte;
+- clean dataset shape: query, positive, negative only;
+- explicit Router mapping: query -> query route, positive/negative -> document route;
+- stronger document FLOPS regularization than the old run to avoid dense docs;
+- custom CPU evaluator because sparse materialization is not safe on MPS here.
+"""
+
+from __future__ import annotations
+
+import argparse
+import ast
+import json
+import logging
+import os
+from pathlib import Path
+
+import torch
+from datasets import Dataset, load_dataset
+
+from sentence_transformers import (
+    SparseEncoder,
+    SparseEncoderModelCardData,
+    SparseEncoderTrainer,
+    SparseEncoderTrainingArguments,
+)
+from sentence_transformers.base.sampler import BatchSamplers
+from sentence_transformers.sparse_encoder.data_collator import SparseEncoderDataCollator
+from sentence_transformers.sparse_encoder.losses import SparseMultipleNegativesRankingLoss, SpladeLoss
+
+
+DATASET = "oneryalcin/financial-filings-sparse-retrieval-training"
+CONFIG = "combined"
+BASE_MODEL = "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill"
+RUN_NAME = "fin-sparse-encoder-doc-v2-clean-router"
+OUTPUT_DIR = Path("models") / RUN_NAME
+LOG_DIR = Path("logs")
+
+
+def setup_logging(run_name: str) -> None:
+    LOG_DIR.mkdir(exist_ok=True)
+    logging.basicConfig(
+        format="%(asctime)s - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+        level=logging.INFO,
+        handlers=[logging.StreamHandler(), logging.FileHandler(LOG_DIR / f"{run_name}.log")],
+        force=True,
+    )
+    for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
+        logging.getLogger(noisy).setLevel(logging.WARNING)
+
+
+def first_negative(value) -> str:
+    if isinstance(value, list):
+        return value[0] if value else ""
+    if isinstance(value, str):
+        parsed = ast.literal_eval(value)
+        return parsed[0] if parsed else ""
+    raise TypeError(f"Unsupported negatives value: {type(value)}")
+
+
+def clean_triplets(ds, limit: int | None = None) -> Dataset:
+    if limit is not None:
+        ds = ds.select(range(min(limit, len(ds))))
+    queries: list[str] = []
+    positives: list[str] = []
+    negatives: list[str] = []
+    for query, positive, raw_negative in zip(ds["query"], ds["positive"], ds["negatives"], strict=True):
+        negative = first_negative(raw_negative)
+        if not negative:
+            continue
+        queries.append(query)
+        positives.append(positive)
+        negatives.append(negative)
+    return Dataset.from_dict(
+        {
+            "query": queries,
+            "positive": positives,
+            "negative": negatives,
+        }
+    )
+
+
+def dense(x: torch.Tensor) -> torch.Tensor:
+    return x.to_dense() if x.is_sparse else x
+
+
+def encode_query(model: SparseEncoder, texts: list[str], batch_size: int) -> torch.Tensor:
+    return dense(model.encode_query(texts, batch_size=batch_size, convert_to_tensor=True, show_progress_bar=False))
+
+
+def encode_document(model: SparseEncoder, texts: list[str], batch_size: int) -> torch.Tensor:
+    return dense(model.encode_document(texts, batch_size=batch_size, convert_to_tensor=True, show_progress_bar=False))
+
+
+def domain_triplet_eval(model: SparseEncoder, eval_ds: Dataset, batch_size: int) -> dict[str, float]:
+    original_device = str(model.device)
+    model.to("cpu")
+    queries = list(eval_ds["query"])
+    positives = list(eval_ds["positive"])
+    negatives = list(eval_ds["negative"])
+    with torch.no_grad():
+        query_emb = encode_query(model, queries, batch_size)
+        pos_emb = encode_document(model, positives, batch_size)
+        neg_emb = encode_document(model, negatives, batch_size)
+        margins = (query_emb * pos_emb).sum(dim=1) - (query_emb * neg_emb).sum(dim=1)
+        result = {
+            "accuracy": (margins > 0).float().mean().item(),
+            "mean_margin": margins.float().mean().item(),
+            "median_margin": margins.float().median().item(),
+            "query_active_dims": (query_emb != 0).sum(dim=1).float().mean().item(),
+            "positive_active_dims": (pos_emb != 0).sum(dim=1).float().mean().item(),
+            "negative_active_dims": (neg_emb != 0).sum(dim=1).float().mean().item(),
+        }
+    if original_device != "cpu":
+        model.to(original_device)
+    return result
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--train-size", type=int, default=10_000)
+    parser.add_argument("--eval-size", type=int, default=500)
+    parser.add_argument("--max-steps", type=int, default=100)
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument("--eval-batch-size", type=int, default=16)
+    parser.add_argument("--max-seq-length", type=int, default=384)
+    parser.add_argument("--query-reg", type=float, default=1e-4)
+    parser.add_argument("--doc-reg", type=float, default=8e-5)
+    parser.add_argument("--run-name", default=RUN_NAME)
+    cli = parser.parse_args()
+
+    output_dir = Path("models") / cli.run_name
+    setup_logging(cli.run_name)
+    logging.info("Torch: %s", torch.__version__)
+    logging.info("MPS available: %s", torch.backends.mps.is_available())
+    logging.info("Loading base model: %s", BASE_MODEL)
+    model = SparseEncoder(
+        BASE_MODEL,
+        model_card_data=SparseEncoderModelCardData(
+            language="en",
+            license="apache-2.0",
+            model_name=f"Financial filings sparse encoder V2 ({cli.run_name})",
+        ),
+    )
+    model.max_seq_length = cli.max_seq_length
+    logging.info("Training device: %s", model.device)
+    logging.info("max_seq_length: %s", model.max_seq_length)
+
+    raw_train = load_dataset(DATASET, CONFIG, split="train")
+    raw_test = load_dataset(DATASET, CONFIG, split="test")
+    train_dataset = clean_triplets(raw_train, cli.train_size)
+    eval_dataset = clean_triplets(raw_test, cli.eval_size)
+    logging.info("train rows: %s | eval rows: %s", len(train_dataset), len(eval_dataset))
+
+    logging.info("Baseline domain eval:")
+    baseline = domain_triplet_eval(model, eval_dataset, cli.eval_batch_size)
+    logging.info("BASELINE: %s", json.dumps(baseline, sort_keys=True))
+
+    loss = SpladeLoss(
+        model=model,
+        loss=SparseMultipleNegativesRankingLoss(model=model),
+        query_regularizer_weight=cli.query_reg,
+        document_regularizer_weight=cli.doc_reg,
+    )
+    data_collator = SparseEncoderDataCollator(
+        preprocess_fn=model.preprocess,
+        router_mapping={
+            "query": "query",
+            "positive": "document",
+            "negative": "document",
+        },
+    )
+
+    args = SparseEncoderTrainingArguments(
+        output_dir=str(output_dir),
+        max_steps=cli.max_steps,
+        num_train_epochs=1,
+        per_device_train_batch_size=cli.batch_size,
+        learning_rate=2e-5,
+        weight_decay=0.01,
+        warmup_steps=0.1,
+        lr_scheduler_type="linear",
+        bf16=False,
+        fp16=False,
+        batch_sampler=BatchSamplers.NO_DUPLICATES,
+        eval_strategy="no",
+        save_strategy="no",
+        logging_steps=10,
+        logging_first_step=True,
+        dataloader_pin_memory=False,
+        report_to="none",
+        run_name=cli.run_name,
+        seed=12,
+    )
+    trainer = SparseEncoderTrainer(
+        model=model,
+        args=args,
+        train_dataset=train_dataset,
+        loss=loss,
+        data_collator=data_collator,
+    )
+    trainer.train()
+
+    logging.info("Post-training domain eval:")
+    score = domain_triplet_eval(model, eval_dataset, cli.eval_batch_size)
+    logging.info("SCORE: %s", json.dumps(score, sort_keys=True))
+    delta = score["accuracy"] - baseline["accuracy"]
+    verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
+    logging.info(
+        "VERDICT: %s | score=%.4f | baseline=%.4f | delta=%+.4f | query_active=%.1f doc_active=%.1f",
+        verdict,
+        score["accuracy"],
+        baseline["accuracy"],
+        delta,
+        score["query_active_dims"],
+        score["positive_active_dims"],
+    )
+
+    final_dir = output_dir / "final"
+    model.save_pretrained(str(final_dir))
+    logging.info("Saved final model to %s", final_dir)
+
+
+if __name__ == "__main__":
+    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+    main()