0xKitkat commited on
Commit
08002bd
·
verified ·
1 Parent(s): e1b526d

Upload folder using huggingface_hub

Browse files
Files changed (7) hide show
  1. README.md +168 -0
  2. chunk.py +360 -0
  3. config.json +86 -0
  4. model.safetensors +3 -0
  5. onnx/model.onnx +3 -0
  6. tokenizer.json +0 -0
  7. 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
+ }