Upload folder using huggingface_hub
Browse files- README.md +168 -0
- chunk.py +360 -0
- config.json +86 -0
- model.safetensors +3 -0
- onnx/model.onnx +3 -0
- tokenizer.json +0 -0
- tokenizer_config.json +17 -0
README.md
ADDED
|
@@ -0,0 +1,168 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: apache-2.0
|
| 3 |
+
base_model: answerdotai/ModernBERT-base
|
| 4 |
+
pipeline_tag: token-classification
|
| 5 |
+
tags:
|
| 6 |
+
- text-segmentation
|
| 7 |
+
- semantic-chunking
|
| 8 |
+
- rag
|
| 9 |
+
- chunking
|
| 10 |
+
datasets:
|
| 11 |
+
- saeedabc/wiki727k
|
| 12 |
+
language:
|
| 13 |
+
- en
|
| 14 |
+
---
|
| 15 |
+
|
| 16 |
+
# Semantic chunker (ModernBERT-base)
|
| 17 |
+
|
| 18 |
+
Predicts, for each sentence, whether a topic boundary follows it. Intended as a
|
| 19 |
+
drop-in chunker for RAG pipelines.
|
| 20 |
+
|
| 21 |
+
## What is different here
|
| 22 |
+
|
| 23 |
+
- **2048-token training context.** The closest existing model trains at 1024.
|
| 24 |
+
Median Wikipedia article is ~980 tokens, so 2048 keeps ~82% of documents in a
|
| 25 |
+
single window instead of ~53%.
|
| 26 |
+
- **Supervision at sentence-final tokens only.** Labelling every token puts the
|
| 27 |
+
positive rate near 1% and the model collapses toward never predicting a split.
|
| 28 |
+
Scoring only sentence-final positions raises it to ~12%.
|
| 29 |
+
- **Trained on Wikipedia section structure**, not book paragraphing.
|
| 30 |
+
|
| 31 |
+
## Results
|
| 32 |
+
|
| 33 |
+
Both models are ModernBERT-base token taggers sharing a tokenizer, so they are
|
| 34 |
+
scored at identical positions: P(boundary) read at each sentence's final token.
|
| 35 |
+
Threshold swept 0.05–0.95; the best operating point is reported.
|
| 36 |
+
|
| 37 |
+
### wiki727k test — in-distribution for this model
|
| 38 |
+
|
| 39 |
+
| Model | F1 | Precision | Recall | Threshold |
|
| 40 |
+
|---|---|---|---|---|
|
| 41 |
+
| **this model** | 0.8156 | 0.8390 | 0.7934 | 0.70 |
|
| 42 |
+
| `mirth/chonky_modernbert_base_1` | 0.4986 | 0.4543 | 0.5524 | 0.90 |
|
| 43 |
+
| baseline:every-k (best k, oracle-tuned) (k=2) | 0.1984 | 0.1228 | 0.5159 | — |
|
| 44 |
+
| baseline:base-rate random | 0.1177 | 0.1166 | 0.1187 | — |
|
| 45 |
+
|
| 46 |
+
### PubMed-RCT — out-of-distribution for both models
|
| 47 |
+
|
| 48 |
+
Non-Wikipedia scientific prose with human-assigned section labels. Boundary rate
|
| 49 |
+
is 35.2% here versus ~11.7% on Wikipedia, so
|
| 50 |
+
scores are **not** comparable across the two tables — only across models within
|
| 51 |
+
one table.
|
| 52 |
+
|
| 53 |
+
| Model | F1 | Precision | Recall | Threshold |
|
| 54 |
+
|---|---|---|---|---|
|
| 55 |
+
| **this model** | 0.6932 | 0.7293 | 0.6605 | 0.55 |
|
| 56 |
+
| `mirth/chonky_modernbert_base_1` | 0.5024 | 0.5095 | 0.4954 | 0.85 |
|
| 57 |
+
| baseline:every-k (best k, oracle-tuned) (k=2) | 0.4034 | 0.3500 | 0.4759 | — |
|
| 58 |
+
| baseline:base-rate random | 0.3490 | 0.3521 | 0.3460 | — |
|
| 59 |
+
|
| 60 |
+
### Retrieval — does better boundary detection actually help?
|
| 61 |
+
|
| 62 |
+
Macro-average over 3 corpora (gov_report, qmsum, stackoverflow), embedder `BAAI/bge-small-en-v1.5`, identical 512-token cap on every strategy so chunk size cannot confound the comparison.
|
| 63 |
+
|
| 64 |
+
| Strategy | nDCG@10 | R@1 | R@5 | R@10 | mean chunk tokens |
|
| 65 |
+
|---|---|---|---|---|---|
|
| 66 |
+
| `fixed-512-ovl64` | 0.8625 | 0.8254 | 0.9218 | 0.9500 | 388 |
|
| 67 |
+
| `chonky-min256` | 0.8615 | 0.7995 | 0.9353 | 0.9635 | 274 |
|
| 68 |
+
| **ours-min256** (this model) | 0.8583 | 0.7914 | 0.9328 | 0.9659 | 283 |
|
| 69 |
+
| **ours** (this model) | 0.8575 | 0.7854 | 0.9321 | 0.9685 | 215 |
|
| 70 |
+
| `sentence-8` | 0.8536 | 0.7956 | 0.9267 | 0.9536 | 182 |
|
| 71 |
+
| `fixed-512` | 0.8531 | 0.8019 | 0.9279 | 0.9512 | 385 |
|
| 72 |
+
| `recursive-512` | 0.8528 | 0.7903 | 0.9218 | 0.9598 | 308 |
|
| 73 |
+
| `chonky` | 0.8498 | 0.7702 | 0.9321 | 0.9684 | 146 |
|
| 74 |
+
|
| 75 |
+
Per corpus (nDCG@10):
|
| 76 |
+
|
| 77 |
+
| Strategy | gov_report | qmsum | stackoverflow |
|
| 78 |
+
|---|---|---|---|
|
| 79 |
+
| `fixed-512-ovl64` | **0.9423** | 0.6988 | 0.9464 |
|
| 80 |
+
| `chonky-min256` | 0.9285 | 0.7099 | 0.9461 |
|
| 81 |
+
| **ours-min256** | 0.9225 | 0.7056 | 0.9467 |
|
| 82 |
+
| **ours** | 0.9191 | **0.7254** | 0.9280 |
|
| 83 |
+
| `sentence-8` | 0.9272 | 0.6865 | **0.9472** |
|
| 84 |
+
| `fixed-512` | 0.9348 | 0.6782 | 0.9463 |
|
| 85 |
+
| `recursive-512` | 0.9214 | 0.6910 | 0.9459 |
|
| 86 |
+
| `chonky` | 0.9046 | 0.7118 | 0.9329 |
|
| 87 |
+
|
| 88 |
+
### Reading these numbers honestly
|
| 89 |
+
|
| 90 |
+
`baseline:every-k` splits every k sentences with k chosen to maximise its own
|
| 91 |
+
score — a model-free floor with oracle tuning in its favour. The wiki727k
|
| 92 |
+
comparison flatters this model: Wikipedia section boundaries are its training
|
| 93 |
+
distribution and are out-of-distribution for chonky, which trained on BookCorpus
|
| 94 |
+
paragraphs. **The PubMed table is the fair comparison.**
|
| 95 |
+
|
| 96 |
+
**Better boundaries did not produce a uniform retrieval win, and the retrieval
|
| 97 |
+
table above should be read before adopting this.** That result is consistent
|
| 98 |
+
with the published critique of semantic chunking, and it is reported here rather
|
| 99 |
+
than omitted. Boundary F1 and retrieval quality are different things; this model
|
| 100 |
+
is much better at the first and situationally better at the second.
|
| 101 |
+
|
| 102 |
+
## When this helps, and when it does not
|
| 103 |
+
|
| 104 |
+
**Use it** for long documents whose topic genuinely shifts partway through —
|
| 105 |
+
transcripts, reports, articles, manuals. That is where fixed-size splitting cuts
|
| 106 |
+
through the middle of an idea and where boundary detection pays.
|
| 107 |
+
|
| 108 |
+
**Do not use it** for corpora whose documents are already shorter than your
|
| 109 |
+
embedder's window. Splitting a 150-token document into two 75-token fragments
|
| 110 |
+
makes retrieval worse, not better, no matter how correct the boundary is. On
|
| 111 |
+
such corpora fixed-size chunking is competitive, faster, and has no
|
| 112 |
+
dependencies — use that instead.
|
| 113 |
+
|
| 114 |
+
**`min_chunk_tokens` is a corpus-dependent knob, not a default.** Setting it to
|
| 115 |
+
256 recovered most of the short-document deficit (+0.019 nDCG@10 on
|
| 116 |
+
stackoverflow) but cost almost exactly as much on long transcripts (−0.020 on
|
| 117 |
+
qmsum), by merging away the boundaries that made the model useful there. Set it
|
| 118 |
+
high for short documents, leave it at 0 for long ones, and measure on your own
|
| 119 |
+
corpus rather than trusting either default.
|
| 120 |
+
|
| 121 |
+
**Cost.** Chunking with this model is roughly one to two orders of magnitude
|
| 122 |
+
slower than a fixed-size splitter (per-strategy wall time is recorded in
|
| 123 |
+
`retrieval_eval.json`). On corpora where it does not win, that cost buys
|
| 124 |
+
nothing.
|
| 125 |
+
|
| 126 |
+
## Usage
|
| 127 |
+
|
| 128 |
+
`chunk.py` ships in this repo — grab it alongside the weights:
|
| 129 |
+
|
| 130 |
+
```bash
|
| 131 |
+
hf download 0xKitkat/semantic-chunker-modernbert-base chunk.py --local-dir .
|
| 132 |
+
```
|
| 133 |
+
|
| 134 |
+
```python
|
| 135 |
+
from chunk import SemanticChunker
|
| 136 |
+
|
| 137 |
+
chunker = SemanticChunker("0xKitkat/semantic-chunker-modernbert-base", threshold=0.5, max_chunk_tokens=512)
|
| 138 |
+
for c in chunker.split(document_text):
|
| 139 |
+
print(c.n_tokens, c.text[:80])
|
| 140 |
+
```
|
| 141 |
+
|
| 142 |
+
`max_chunk_tokens` force-splits over-long chunks at the model's lowest-confidence
|
| 143 |
+
interior boundary rather than at an arbitrary offset.
|
| 144 |
+
|
| 145 |
+
## Training
|
| 146 |
+
|
| 147 |
+
| | |
|
| 148 |
+
|---|---|
|
| 149 |
+
| Base | `answerdotai/ModernBERT-base` (149M) |
|
| 150 |
+
| Data | wiki727k, 582,160 train docs / 30,580,099 sentences |
|
| 151 |
+
| Boundary rate | 11.7% |
|
| 152 |
+
| Max length | 2048 |
|
| 153 |
+
| Effective batch | 32 |
|
| 154 |
+
| LR / schedule | 3e-5, linear, 6% warmup |
|
| 155 |
+
| Epochs | 1 |
|
| 156 |
+
| Loss | class-weighted CE (positive weight 3.0) |
|
| 157 |
+
| Hardware | 1x RTX 4070 SUPER (12GB) |
|
| 158 |
+
|
| 159 |
+
Section titles are dropped from the training text: with headings present the
|
| 160 |
+
task partly degenerates into "a heading follows", which does not transfer to the
|
| 161 |
+
unformatted prose a chunker sees in production.
|
| 162 |
+
|
| 163 |
+
## Limitations
|
| 164 |
+
|
| 165 |
+
- English only.
|
| 166 |
+
- Trained on encyclopedic prose; conversational transcripts, code, and tabular
|
| 167 |
+
documents are out of distribution.
|
| 168 |
+
- Boundary F1 is a proxy. If your goal is retrieval quality, measure retrieval.
|
chunk.py
ADDED
|
@@ -0,0 +1,360 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Inference-side semantic chunker.
|
| 2 |
+
|
| 3 |
+
The training task is "does a segment boundary follow this sentence"; this module
|
| 4 |
+
turns that into the thing a RAG pipeline actually wants: a list of text chunks,
|
| 5 |
+
subject to hard size limits an embedder can accept.
|
| 6 |
+
|
| 7 |
+
Two things here are not in the training loop and matter in production:
|
| 8 |
+
|
| 9 |
+
1. Windowed inference. Real documents run far past 2048 tokens. We slide a
|
| 10 |
+
window with sentence-level left context so boundary decisions near a window
|
| 11 |
+
edge still see what came before, instead of being scored cold.
|
| 12 |
+
2. Hard size caps. A semantic boundary model will happily emit a 4000-token
|
| 13 |
+
chunk if the topic does not shift. Embedders truncate at their own limit, so
|
| 14 |
+
`max_chunk_tokens` force-splits at the lowest-confidence interior sentence
|
| 15 |
+
rather than at an arbitrary character offset.
|
| 16 |
+
"""
|
| 17 |
+
|
| 18 |
+
from __future__ import annotations
|
| 19 |
+
|
| 20 |
+
import os
|
| 21 |
+
import re
|
| 22 |
+
from dataclasses import dataclass
|
| 23 |
+
from typing import Callable, Iterable
|
| 24 |
+
|
| 25 |
+
import numpy as np
|
| 26 |
+
from transformers import AutoTokenizer
|
| 27 |
+
|
| 28 |
+
# Sentence splitter: deliberately dependency-free. Handles the common
|
| 29 |
+
# abbreviation cases that naive `.` splitting gets wrong. Pass your own via
|
| 30 |
+
# `sentence_splitter=` if you already have spacy/nltk/pysbd in the pipeline.
|
| 31 |
+
_ABBREV = (
|
| 32 |
+
r"(?<!\bMr)(?<!\bMrs)(?<!\bMs)(?<!\bDr)(?<!\bProf)(?<!\bSt)(?<!\bJr)(?<!\bSr)"
|
| 33 |
+
r"(?<!\bInc)(?<!\bLtd)(?<!\bCo)(?<!\bvs)(?<!\betc)(?<!\bi\.e)(?<!\be\.g)"
|
| 34 |
+
r"(?<!\bFig)(?<!\bNo)(?<!\bVol)(?<!\bApprox)(?<!\b[A-Z])"
|
| 35 |
+
)
|
| 36 |
+
_SENT_RE = re.compile(rf"{_ABBREV}(?<=[.!?])[\"')\]]*\s+(?=[A-Z0-9\"'(\[])")
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def default_sentence_splitter(text: str) -> list[str]:
|
| 40 |
+
out = []
|
| 41 |
+
for block in re.split(r"\n\s*\n", text): # paragraph breaks are always splits
|
| 42 |
+
block = block.strip()
|
| 43 |
+
if not block:
|
| 44 |
+
continue
|
| 45 |
+
out.extend(s.strip() for s in _SENT_RE.split(block) if s.strip())
|
| 46 |
+
return out
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def _softmax_last(x: np.ndarray) -> np.ndarray:
|
| 50 |
+
x = x - x.max(axis=-1, keepdims=True)
|
| 51 |
+
e = np.exp(x)
|
| 52 |
+
return e / e.sum(axis=-1, keepdims=True)
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
class TorchBackend:
|
| 56 |
+
"""Default backend. torch is imported here so the ONNX path stays torch-free."""
|
| 57 |
+
|
| 58 |
+
def __init__(self, model_path: str, device: str | None = None):
|
| 59 |
+
import torch
|
| 60 |
+
from transformers import AutoModelForTokenClassification
|
| 61 |
+
|
| 62 |
+
self._torch = torch
|
| 63 |
+
self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
| 64 |
+
self.model = (
|
| 65 |
+
AutoModelForTokenClassification.from_pretrained(model_path)
|
| 66 |
+
.eval()
|
| 67 |
+
.to(self.device)
|
| 68 |
+
)
|
| 69 |
+
self.id2label = dict(self.model.config.id2label)
|
| 70 |
+
|
| 71 |
+
def __call__(self, input_ids: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
|
| 72 |
+
t = self._torch
|
| 73 |
+
with t.no_grad():
|
| 74 |
+
out = self.model(
|
| 75 |
+
input_ids=t.tensor(input_ids, dtype=t.long, device=self.device),
|
| 76 |
+
attention_mask=t.tensor(attention_mask, dtype=t.long, device=self.device),
|
| 77 |
+
).logits
|
| 78 |
+
return out.float().cpu().numpy()
|
| 79 |
+
|
| 80 |
+
|
| 81 |
+
class OnnxBackend:
|
| 82 |
+
"""CPU inference with onnxruntime only -- no torch in the dependency tree."""
|
| 83 |
+
|
| 84 |
+
def __init__(self, model_dir: str, providers: list[str] | None = None):
|
| 85 |
+
import json
|
| 86 |
+
import onnxruntime as ort
|
| 87 |
+
|
| 88 |
+
path = os.path.join(model_dir, "model.onnx")
|
| 89 |
+
if not os.path.exists(path):
|
| 90 |
+
raise FileNotFoundError(f"no model.onnx under {model_dir}")
|
| 91 |
+
self.sess = ort.InferenceSession(
|
| 92 |
+
path, providers=providers or ["CPUExecutionProvider"]
|
| 93 |
+
)
|
| 94 |
+
with open(os.path.join(model_dir, "config.json"), encoding="utf-8") as f:
|
| 95 |
+
cfg = json.load(f)
|
| 96 |
+
self.id2label = cfg.get("id2label", {0: "O", 1: "semantic-shift"})
|
| 97 |
+
|
| 98 |
+
def __call__(self, input_ids: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
|
| 99 |
+
return self.sess.run(
|
| 100 |
+
None,
|
| 101 |
+
{
|
| 102 |
+
"input_ids": input_ids.astype(np.int64),
|
| 103 |
+
"attention_mask": attention_mask.astype(np.int64),
|
| 104 |
+
},
|
| 105 |
+
)[0]
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
@dataclass
|
| 109 |
+
class Chunk:
|
| 110 |
+
text: str
|
| 111 |
+
start_sentence: int
|
| 112 |
+
end_sentence: int
|
| 113 |
+
n_tokens: int
|
| 114 |
+
boundary_score: float # confidence of the boundary that closed this chunk
|
| 115 |
+
|
| 116 |
+
|
| 117 |
+
class SemanticChunker:
|
| 118 |
+
def __init__(
|
| 119 |
+
self,
|
| 120 |
+
model_path: str,
|
| 121 |
+
threshold: float = 0.5,
|
| 122 |
+
max_chunk_tokens: int = 512,
|
| 123 |
+
min_chunk_tokens: int = 0,
|
| 124 |
+
min_chunk_sentences: int = 1,
|
| 125 |
+
window_tokens: int = 2048,
|
| 126 |
+
context_sentences: int = 3,
|
| 127 |
+
device: str | None = None,
|
| 128 |
+
sentence_splitter: Callable[[str], list[str]] | None = None,
|
| 129 |
+
batch_size: int = 8,
|
| 130 |
+
backend: str = "torch",
|
| 131 |
+
tokenizer_path: str | None = None,
|
| 132 |
+
):
|
| 133 |
+
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path or model_path)
|
| 134 |
+
if backend == "onnx":
|
| 135 |
+
self.backend = OnnxBackend(model_path)
|
| 136 |
+
elif backend == "torch":
|
| 137 |
+
self.backend = TorchBackend(model_path, device)
|
| 138 |
+
else:
|
| 139 |
+
raise ValueError(f"unknown backend {backend!r}, expected 'torch' or 'onnx'")
|
| 140 |
+
self.device = getattr(self.backend, "device", "cpu")
|
| 141 |
+
self.threshold = threshold
|
| 142 |
+
self.max_chunk_tokens = max_chunk_tokens
|
| 143 |
+
self.min_chunk_tokens = min_chunk_tokens
|
| 144 |
+
self.min_chunk_sentences = max(1, min_chunk_sentences)
|
| 145 |
+
self.window_tokens = window_tokens
|
| 146 |
+
self.context_sentences = context_sentences
|
| 147 |
+
self.split_sentences = sentence_splitter or default_sentence_splitter
|
| 148 |
+
self.batch_size = batch_size
|
| 149 |
+
|
| 150 |
+
id2label = dict(self.backend.id2label)
|
| 151 |
+
self.pos_idx = next(
|
| 152 |
+
(int(i) for i, n in id2label.items()
|
| 153 |
+
if str(n).lower() in {"semantic-shift", "separator", "shift"}),
|
| 154 |
+
1,
|
| 155 |
+
)
|
| 156 |
+
|
| 157 |
+
# ------------------------------------------------------------------
|
| 158 |
+
def boundary_scores(self, sentences: list[str]) -> list[float]:
|
| 159 |
+
"""P(boundary follows) for each sentence. Last entry is always 1.0."""
|
| 160 |
+
if not sentences:
|
| 161 |
+
return []
|
| 162 |
+
if len(sentences) == 1:
|
| 163 |
+
return [1.0]
|
| 164 |
+
|
| 165 |
+
toks = self.tokenizer(
|
| 166 |
+
[s if i == 0 else " " + s for i, s in enumerate(sentences)],
|
| 167 |
+
add_special_tokens=False,
|
| 168 |
+
)["input_ids"]
|
| 169 |
+
|
| 170 |
+
budget = self.window_tokens - 2
|
| 171 |
+
windows = [] # (ids, [(sentence_idx, token_pos)])
|
| 172 |
+
start = 0
|
| 173 |
+
n = len(sentences)
|
| 174 |
+
while start < n:
|
| 175 |
+
ids, marks = [], []
|
| 176 |
+
ctx_start = max(0, start - self.context_sentences)
|
| 177 |
+
for i in range(ctx_start, start):
|
| 178 |
+
ids.extend(toks[i])
|
| 179 |
+
i = start
|
| 180 |
+
while i < n:
|
| 181 |
+
t = toks[i] or [self.tokenizer.unk_token_id]
|
| 182 |
+
if len(t) > budget:
|
| 183 |
+
t = t[:budget]
|
| 184 |
+
if len(ids) + len(t) > budget:
|
| 185 |
+
break
|
| 186 |
+
ids.extend(t)
|
| 187 |
+
marks.append((i, len(ids) - 1))
|
| 188 |
+
i += 1
|
| 189 |
+
if i == start: # pathological single sentence
|
| 190 |
+
marks.append((start, max(len(ids) - 1, 0)))
|
| 191 |
+
i = start + 1
|
| 192 |
+
windows.append((ids, marks))
|
| 193 |
+
start = i
|
| 194 |
+
|
| 195 |
+
scores = [0.0] * n
|
| 196 |
+
for b in range(0, len(windows), self.batch_size):
|
| 197 |
+
batch = windows[b : b + self.batch_size]
|
| 198 |
+
maxlen = max(len(w[0]) for w in batch) + 2
|
| 199 |
+
input_ids = np.full(
|
| 200 |
+
(len(batch), maxlen), self.tokenizer.pad_token_id, dtype=np.int64
|
| 201 |
+
)
|
| 202 |
+
attn = np.zeros((len(batch), maxlen), dtype=np.int64)
|
| 203 |
+
for r, (ids, _) in enumerate(batch):
|
| 204 |
+
seq = [self.tokenizer.cls_token_id] + ids + [self.tokenizer.sep_token_id]
|
| 205 |
+
input_ids[r, : len(seq)] = seq
|
| 206 |
+
attn[r, : len(seq)] = 1
|
| 207 |
+
logits = self.backend(input_ids, attn)
|
| 208 |
+
probs = _softmax_last(logits.astype(np.float32))[..., self.pos_idx]
|
| 209 |
+
for r, (_, marks) in enumerate(batch):
|
| 210 |
+
for sent_idx, tok_pos in marks:
|
| 211 |
+
scores[sent_idx] = float(probs[r, tok_pos + 1]) # +1 for CLS
|
| 212 |
+
|
| 213 |
+
scores[-1] = 1.0 # document end is always a boundary
|
| 214 |
+
return scores
|
| 215 |
+
|
| 216 |
+
# ------------------------------------------------------------------
|
| 217 |
+
def _token_lens(self, texts: list[str]) -> list[int]:
|
| 218 |
+
"""One batched tokenizer call, not one per sentence."""
|
| 219 |
+
if not texts:
|
| 220 |
+
return []
|
| 221 |
+
enc = self.tokenizer(texts, add_special_tokens=False)["input_ids"]
|
| 222 |
+
return [len(x) for x in enc]
|
| 223 |
+
|
| 224 |
+
def split(self, text: str) -> list[Chunk]:
|
| 225 |
+
sentences = self.split_sentences(text)
|
| 226 |
+
if not sentences:
|
| 227 |
+
return []
|
| 228 |
+
scores = self.boundary_scores(sentences)
|
| 229 |
+
sent_tokens = self._token_lens(sentences)
|
| 230 |
+
|
| 231 |
+
chunks: list[Chunk] = []
|
| 232 |
+
buf_start = 0
|
| 233 |
+
buf_tokens = 0
|
| 234 |
+
|
| 235 |
+
def emit(end_idx: int, score: float):
|
| 236 |
+
nonlocal buf_start, buf_tokens
|
| 237 |
+
body = " ".join(sentences[buf_start : end_idx + 1]).strip()
|
| 238 |
+
if body:
|
| 239 |
+
chunks.append(
|
| 240 |
+
Chunk(
|
| 241 |
+
text=body,
|
| 242 |
+
start_sentence=buf_start,
|
| 243 |
+
end_sentence=end_idx,
|
| 244 |
+
n_tokens=sum(sent_tokens[buf_start : end_idx + 1]),
|
| 245 |
+
boundary_score=score,
|
| 246 |
+
)
|
| 247 |
+
)
|
| 248 |
+
buf_start = end_idx + 1
|
| 249 |
+
buf_tokens = 0
|
| 250 |
+
|
| 251 |
+
for i, sent in enumerate(sentences):
|
| 252 |
+
buf_tokens += sent_tokens[i]
|
| 253 |
+
n_in_buf = i - buf_start + 1
|
| 254 |
+
|
| 255 |
+
# hard cap wins over semantics: split at the weakest interior
|
| 256 |
+
# boundary so we cut where the model is least confident, not
|
| 257 |
+
# wherever the token counter happened to run out.
|
| 258 |
+
if buf_tokens > self.max_chunk_tokens and n_in_buf > 1:
|
| 259 |
+
interior = range(buf_start, i)
|
| 260 |
+
cut = max(interior, key=lambda j: scores[j])
|
| 261 |
+
emit(cut, scores[cut])
|
| 262 |
+
buf_tokens = sum(sent_tokens[buf_start : i + 1])
|
| 263 |
+
n_in_buf = i - buf_start + 1
|
| 264 |
+
|
| 265 |
+
if scores[i] >= self.threshold and n_in_buf >= self.min_chunk_sentences:
|
| 266 |
+
emit(i, scores[i])
|
| 267 |
+
|
| 268 |
+
if buf_start < len(sentences):
|
| 269 |
+
emit(len(sentences) - 1, 1.0)
|
| 270 |
+
|
| 271 |
+
return self._merge_undersized(chunks, sentences, sent_tokens)
|
| 272 |
+
|
| 273 |
+
def _merge_undersized(self, chunks, sentences, sent_tokens):
|
| 274 |
+
"""Merge chunks below min_chunk_tokens into their neighbour.
|
| 275 |
+
|
| 276 |
+
Measured need, not speculation: on short-document corpora the boundary
|
| 277 |
+
model happily splits a 150-token document in two, and two 75-token
|
| 278 |
+
fragments retrieve worse than one whole document. A topic boundary is
|
| 279 |
+
real there, but acting on it is counterproductive -- the embedder needs
|
| 280 |
+
enough text to place the vector well.
|
| 281 |
+
|
| 282 |
+
Merges forward by preference (keeps reading order), backward for a
|
| 283 |
+
trailing runt, and never produces a chunk exceeding max_chunk_tokens.
|
| 284 |
+
"""
|
| 285 |
+
if self.min_chunk_tokens <= 0 or len(chunks) < 2:
|
| 286 |
+
return chunks
|
| 287 |
+
|
| 288 |
+
out = []
|
| 289 |
+
i = 0
|
| 290 |
+
while i < len(chunks):
|
| 291 |
+
cur = chunks[i]
|
| 292 |
+
while (
|
| 293 |
+
cur.n_tokens < self.min_chunk_tokens
|
| 294 |
+
and i + 1 < len(chunks)
|
| 295 |
+
and cur.n_tokens + chunks[i + 1].n_tokens <= self.max_chunk_tokens
|
| 296 |
+
):
|
| 297 |
+
nxt = chunks[i + 1]
|
| 298 |
+
cur = Chunk(
|
| 299 |
+
text=(cur.text + " " + nxt.text).strip(),
|
| 300 |
+
start_sentence=cur.start_sentence,
|
| 301 |
+
end_sentence=nxt.end_sentence,
|
| 302 |
+
n_tokens=cur.n_tokens + nxt.n_tokens,
|
| 303 |
+
boundary_score=nxt.boundary_score,
|
| 304 |
+
)
|
| 305 |
+
i += 1
|
| 306 |
+
out.append(cur)
|
| 307 |
+
i += 1
|
| 308 |
+
|
| 309 |
+
# a trailing runt has no forward neighbour left; fold it back
|
| 310 |
+
if (
|
| 311 |
+
len(out) > 1
|
| 312 |
+
and out[-1].n_tokens < self.min_chunk_tokens
|
| 313 |
+
and out[-2].n_tokens + out[-1].n_tokens <= self.max_chunk_tokens
|
| 314 |
+
):
|
| 315 |
+
last, prev = out.pop(), out.pop()
|
| 316 |
+
out.append(
|
| 317 |
+
Chunk(
|
| 318 |
+
text=(prev.text + " " + last.text).strip(),
|
| 319 |
+
start_sentence=prev.start_sentence,
|
| 320 |
+
end_sentence=last.end_sentence,
|
| 321 |
+
n_tokens=prev.n_tokens + last.n_tokens,
|
| 322 |
+
boundary_score=last.boundary_score,
|
| 323 |
+
)
|
| 324 |
+
)
|
| 325 |
+
return out
|
| 326 |
+
|
| 327 |
+
def split_text(self, text: str) -> list[str]:
|
| 328 |
+
"""Convenience: just the strings."""
|
| 329 |
+
return [c.text for c in self.split(text)]
|
| 330 |
+
|
| 331 |
+
def batch_split(self, texts: Iterable[str]) -> list[list[str]]:
|
| 332 |
+
return [self.split_text(t) for t in texts]
|
| 333 |
+
|
| 334 |
+
|
| 335 |
+
if __name__ == "__main__":
|
| 336 |
+
import argparse
|
| 337 |
+
import sys
|
| 338 |
+
|
| 339 |
+
ap = argparse.ArgumentParser()
|
| 340 |
+
ap.add_argument("--model", default=r"D:\LocalAI\chunker\out\run1\final")
|
| 341 |
+
ap.add_argument("--threshold", type=float, default=0.5)
|
| 342 |
+
ap.add_argument("--max-chunk-tokens", type=int, default=512)
|
| 343 |
+
ap.add_argument("--file", help="read text from file; otherwise stdin")
|
| 344 |
+
ap.add_argument("--backend", default="torch", choices=["torch", "onnx"])
|
| 345 |
+
ap.add_argument("--tokenizer", default=None, help="only needed for a bare onnx dir")
|
| 346 |
+
args = ap.parse_args()
|
| 347 |
+
|
| 348 |
+
text = (
|
| 349 |
+
open(args.file, encoding="utf-8").read() if args.file else sys.stdin.read()
|
| 350 |
+
)
|
| 351 |
+
ch = SemanticChunker(
|
| 352 |
+
args.model,
|
| 353 |
+
threshold=args.threshold,
|
| 354 |
+
max_chunk_tokens=args.max_chunk_tokens,
|
| 355 |
+
backend=args.backend,
|
| 356 |
+
tokenizer_path=args.tokenizer,
|
| 357 |
+
)
|
| 358 |
+
for i, c in enumerate(ch.split(text)):
|
| 359 |
+
print(f"\n--- chunk {i} ({c.n_tokens} tok, score {c.boundary_score:.3f}) ---")
|
| 360 |
+
print(c.text)
|
config.json
ADDED
|
@@ -0,0 +1,86 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"ModernBertForTokenClassification"
|
| 4 |
+
],
|
| 5 |
+
"attention_bias": false,
|
| 6 |
+
"attention_dropout": 0.0,
|
| 7 |
+
"bos_token_id": 50281,
|
| 8 |
+
"classifier_activation": "gelu",
|
| 9 |
+
"classifier_bias": false,
|
| 10 |
+
"classifier_dropout": 0.0,
|
| 11 |
+
"classifier_pooling": "mean",
|
| 12 |
+
"cls_token_id": 50281,
|
| 13 |
+
"decoder_bias": true,
|
| 14 |
+
"deterministic_flash_attn": false,
|
| 15 |
+
"dtype": "float32",
|
| 16 |
+
"embedding_dropout": 0.0,
|
| 17 |
+
"eos_token_id": 50282,
|
| 18 |
+
"global_attn_every_n_layers": 3,
|
| 19 |
+
"gradient_checkpointing": false,
|
| 20 |
+
"hidden_activation": "gelu",
|
| 21 |
+
"hidden_size": 768,
|
| 22 |
+
"id2label": {
|
| 23 |
+
"0": "semantic-continuity",
|
| 24 |
+
"1": "semantic-shift"
|
| 25 |
+
},
|
| 26 |
+
"initializer_cutoff_factor": 2.0,
|
| 27 |
+
"initializer_range": 0.02,
|
| 28 |
+
"intermediate_size": 1152,
|
| 29 |
+
"label2id": {
|
| 30 |
+
"semantic-continuity": 0,
|
| 31 |
+
"semantic-shift": 1
|
| 32 |
+
},
|
| 33 |
+
"layer_norm_eps": 1e-05,
|
| 34 |
+
"layer_types": [
|
| 35 |
+
"full_attention",
|
| 36 |
+
"sliding_attention",
|
| 37 |
+
"sliding_attention",
|
| 38 |
+
"full_attention",
|
| 39 |
+
"sliding_attention",
|
| 40 |
+
"sliding_attention",
|
| 41 |
+
"full_attention",
|
| 42 |
+
"sliding_attention",
|
| 43 |
+
"sliding_attention",
|
| 44 |
+
"full_attention",
|
| 45 |
+
"sliding_attention",
|
| 46 |
+
"sliding_attention",
|
| 47 |
+
"full_attention",
|
| 48 |
+
"sliding_attention",
|
| 49 |
+
"sliding_attention",
|
| 50 |
+
"full_attention",
|
| 51 |
+
"sliding_attention",
|
| 52 |
+
"sliding_attention",
|
| 53 |
+
"full_attention",
|
| 54 |
+
"sliding_attention",
|
| 55 |
+
"sliding_attention",
|
| 56 |
+
"full_attention"
|
| 57 |
+
],
|
| 58 |
+
"local_attention": 128,
|
| 59 |
+
"max_position_embeddings": 8192,
|
| 60 |
+
"mlp_bias": false,
|
| 61 |
+
"mlp_dropout": 0.0,
|
| 62 |
+
"model_type": "modernbert",
|
| 63 |
+
"norm_bias": false,
|
| 64 |
+
"norm_eps": 1e-05,
|
| 65 |
+
"num_attention_heads": 12,
|
| 66 |
+
"num_hidden_layers": 22,
|
| 67 |
+
"pad_token_id": 50283,
|
| 68 |
+
"position_embedding_type": "absolute",
|
| 69 |
+
"rope_parameters": {
|
| 70 |
+
"full_attention": {
|
| 71 |
+
"rope_theta": 160000.0,
|
| 72 |
+
"rope_type": "default"
|
| 73 |
+
},
|
| 74 |
+
"sliding_attention": {
|
| 75 |
+
"rope_theta": 10000.0,
|
| 76 |
+
"rope_type": "default"
|
| 77 |
+
}
|
| 78 |
+
},
|
| 79 |
+
"sep_token_id": 50282,
|
| 80 |
+
"sparse_pred_ignore_index": -100,
|
| 81 |
+
"sparse_prediction": false,
|
| 82 |
+
"tie_word_embeddings": true,
|
| 83 |
+
"transformers_version": "5.14.1",
|
| 84 |
+
"use_cache": false,
|
| 85 |
+
"vocab_size": 50368
|
| 86 |
+
}
|
model.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:18ea0a1f198166cc336403c2e6ef0dc99bfe2cd19e31a5579dd1a33fb0b460af
|
| 3 |
+
size 598439784
|
onnx/model.onnx
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:5772dddc662de3944429e0f6b05037fb17b6a20706f8d160e2f70af2826e8662
|
| 3 |
+
size 599052097
|
tokenizer.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
tokenizer_config.json
ADDED
|
@@ -0,0 +1,17 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"backend": "tokenizers",
|
| 3 |
+
"clean_up_tokenization_spaces": true,
|
| 4 |
+
"cls_token": "[CLS]",
|
| 5 |
+
"is_local": true,
|
| 6 |
+
"local_files_only": false,
|
| 7 |
+
"mask_token": "[MASK]",
|
| 8 |
+
"model_input_names": [
|
| 9 |
+
"input_ids",
|
| 10 |
+
"attention_mask"
|
| 11 |
+
],
|
| 12 |
+
"model_max_length": 8192,
|
| 13 |
+
"pad_token": "[PAD]",
|
| 14 |
+
"sep_token": "[SEP]",
|
| 15 |
+
"tokenizer_class": "TokenizersBackend",
|
| 16 |
+
"unk_token": "[UNK]"
|
| 17 |
+
}
|