Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""Owner of the model-artifacts file format, write side and read side.
3The model artifacts are the digest-locked file set the trainer produces as
4one unit: ``models.bin`` (``CMD2``: per-model names and L2 norms, then
5zlib-compressed dense bigram tables), ``rowmax.bin`` (``CRM1``: per-model
6row maxima for upper-bound prescreening, digest-locked to the
7``models.bin`` they were computed from), ``idf.bin`` (a 65536-byte
8quantized IDF table), and ``langnorm.bin`` (``CLN1``: per-model norms
9under the IDF-weighted metric the language tiers score in, digest-locked
10like ``rowmax.bin``). Keeping both directions in one module means a
11format change cannot land on one side only: ``scripts/train.py`` calls
12:func:`write_model_artifacts` and :func:`read_models`, while
13:mod:`chardet.models` calls :func:`parse_models_bin` and
14:func:`parse_rowmax_bin` behind its cached accessors.
16Note: ``from __future__ import annotations`` is intentionally omitted
17because this module is compiled with mypyc, which does not support PEP 563
18string annotations.
19"""
21import hashlib
22import math
23import struct
24import zlib
25from pathlib import Path
27#: models.bin magic: the v2 dense zlib-compressed format.
28MODELS_MAGIC = b"CMD2"
29#: rowmax.bin layout: magic, SHA-256 of the matching models.bin, then one
30#: 256-byte row-maxima table per model in models.bin header order.
31ROWMAX_MAGIC = b"CRM1"
32ROWMAX_HEADER_SIZE = 4 + 32
33#: langnorm.bin layout: magic, SHA-256 of the matching models.bin, then one
34#: big-endian float64 per model in models.bin header order: the model's
35#: L2 norm after each weight is scaled by the square root of its bigram's
36#: IDF (see :func:`langnorm_from_table`).
37LANGNORM_MAGIC = b"CLN1"
38LANGNORM_HEADER_SIZE = 4 + 32
40_unpack_uint32 = struct.Struct(">I").unpack_from
41_unpack_float64 = struct.Struct(">d").unpack_from
44def _decompress_tables(
45 data: bytes, offset: int, names: list[str], chunk_size: int = 262144
46) -> dict[str, bytes]:
47 """Decompress the model tables from ``data[offset:]``, one per name.
49 Each model is stored as its own bytes object rather than a memoryview
50 slice of one big blob: mypyc compiles bytes indexing in the scoring hot
51 loop to a native C array access, while memoryview indexing goes through
52 a boxed generic call. Decompression is incremental, one 64 KB table at
53 a time — materializing the whole multi-megabyte blob and slicing it
54 would transiently double the allocation and strand the freed pages in
55 process RSS. Trailing compressed bytes are ignored, as with whole-blob
56 ``zlib.decompress``; the decompressed size is validated instead.
58 :raises ValueError: If the decompressed size is not exactly
59 ``len(names) * 65536``.
60 """
61 num_models = len(names)
62 expected_size = num_models * 65536
63 decomp = zlib.decompressobj()
64 models: dict[str, bytes] = {}
65 table = bytearray()
66 produced = 0
67 pos = offset
68 end = len(data)
69 flushed = False
70 while len(models) < num_models:
71 need = 65536 - len(table)
72 if decomp.unconsumed_tail:
73 piece = decomp.decompress(decomp.unconsumed_tail, need)
74 elif pos < end:
75 chunk = data[pos : pos + chunk_size]
76 pos += len(chunk)
77 piece = decomp.decompress(chunk, need)
78 elif not flushed:
79 flushed = True
80 piece = decomp.flush()
81 if not piece:
82 break # stream exhausted early -> size mismatch below
83 else:
84 break # stream exhausted early -> size mismatch below
85 produced += len(piece)
86 table += piece
87 if len(table) == 65536:
88 models[names[len(models)]] = bytes(table)
89 table.clear()
90 # Unreachable with CPython's zlib — decompress() pieces are capped at
91 # ``need``, and a flush() strand (the tail of one back-reference cut
92 # mid-copy, at most 258 bytes) only ever lands in a fresh table —
93 # but a decompressor that flushed more than asked must not corrupt
94 # the table split silently.
95 elif len(table) > 65536: # pragma: no cover
96 break # oversized flush -> size mismatch below
97 if len(models) == num_models:
98 # Drain any leftover decompressed output so extra data is caught,
99 # including surplus tables in compressed input not yet fed to the
100 # decompressor. Bytes after the stream's end marker (decomp.eof)
101 # are ignored, matching whole-blob zlib.decompress behavior.
102 extra = b""
103 if decomp.unconsumed_tail:
104 extra = decomp.decompress(decomp.unconsumed_tail, 1)
105 while not extra and not decomp.eof and pos < end:
106 chunk = data[pos : pos + chunk_size]
107 pos += len(chunk)
108 extra = decomp.decompress(chunk, 1)
109 if not extra and not decomp.eof and not flushed:
110 extra = decomp.flush()
111 produced += len(extra)
112 if produced != expected_size or len(models) != num_models:
113 msg = (
114 f"corrupt models.bin: decompressed size {produced} "
115 f"!= expected {expected_size}"
116 )
117 raise ValueError(msg)
118 return models
121def parse_models_bin(
122 data: bytes,
123) -> tuple[dict[str, bytes], dict[str, float]]:
124 """Parse the v2 dense zlib-compressed models.bin format.
126 :param data: Raw bytes of models.bin (must be non-empty).
127 :returns: A ``(models, norms)`` tuple.
128 :raises ValueError: If the data is corrupt or truncated.
129 """
130 try:
131 if data[:4] != MODELS_MAGIC:
132 msg = "corrupt models.bin: missing CMD2 magic"
133 raise ValueError(msg)
135 offset = 4 # skip magic
136 (num_models,) = _unpack_uint32(data, offset)
137 offset += 4
139 if num_models > 10_000:
140 msg = f"corrupt models.bin: num_models={num_models} exceeds limit"
141 raise ValueError(msg)
143 names: list[str] = []
144 norms: dict[str, float] = {}
145 for _ in range(num_models):
146 (name_len,) = _unpack_uint32(data, offset)
147 offset += 4
148 if name_len > 256:
149 msg = f"corrupt models.bin: name_len={name_len} exceeds 256"
150 raise ValueError(msg)
151 name = data[offset : offset + name_len].decode("utf-8")
152 offset += name_len
153 (norm,) = _unpack_float64(data, offset)
154 offset += 8
155 names.append(name)
156 norms[name] = norm
158 models = _decompress_tables(data, offset, names)
160 except zlib.error as e:
161 msg = f"corrupt models.bin: {e}"
162 raise ValueError(msg) from e
163 except (struct.error, UnicodeDecodeError) as e:
164 msg = f"corrupt models.bin: {e}"
165 raise ValueError(msg) from e
167 return models, norms
170def parse_rowmax_bin(
171 data: bytes, models_digest: bytes, model_keys: list[str]
172) -> dict[str, bytes] | None:
173 """Parse rowmax.bin into per-model row-maxima tables.
175 :param data: Raw bytes of rowmax.bin.
176 :param models_digest: SHA-256 digest of the current models.bin bytes.
177 :param model_keys: Model keys in models.bin header order.
178 :returns: Mapping of model key to its 256-byte row-maxima table, or
179 ``None`` when the magic, digest, or size does not match — a stale
180 or mismatched file would silently under-estimate row maxima and
181 break the upper bound that prescreening depends on.
182 """
183 if (
184 data[:4] == ROWMAX_MAGIC
185 and data[4:ROWMAX_HEADER_SIZE] == models_digest
186 and len(data) == ROWMAX_HEADER_SIZE + len(model_keys) * 256
187 ):
188 return {
189 key: data[ROWMAX_HEADER_SIZE + i * 256 : ROWMAX_HEADER_SIZE + (i + 1) * 256]
190 for i, key in enumerate(model_keys)
191 }
192 return None
195def parse_langnorm_bin(
196 data: bytes, models_digest: bytes, model_keys: list[str]
197) -> dict[str, float] | None:
198 """Parse langnorm.bin into per-model language-scoring norms.
200 :param data: Raw bytes of langnorm.bin.
201 :param models_digest: SHA-256 digest of the current models.bin bytes.
202 :param model_keys: Model keys in models.bin header order.
203 :returns: Mapping of model key to its IDF-weighted norm, or ``None``
204 when the magic, digest, or size does not match -- a stale file
205 would normalize every language score against the wrong length and
206 silently tilt the comparison between languages.
207 """
208 if (
209 data[:4] == LANGNORM_MAGIC
210 and data[4:LANGNORM_HEADER_SIZE] == models_digest
211 and len(data) == LANGNORM_HEADER_SIZE + len(model_keys) * 8
212 ):
213 return {
214 key: _unpack_float64(data, LANGNORM_HEADER_SIZE + i * 8)[0]
215 for i, key in enumerate(model_keys)
216 }
217 return None
220def langnorm_from_table(table: bytes, idf: bytes) -> float:
221 """Derive one model's norm under the language-scoring metric.
223 The language tiers compare an input profile weighted by IDF against
224 models weighted by the square root of IDF, so a model's length in that
225 space is the root of the sum of ``weight**2 * idf`` over its bigrams.
226 Scoring folds the model-side factor into the profile and divides by
227 this norm, so the table itself is never rescaled.
228 """
229 total = 0
230 for i in range(65536):
231 v = table[i]
232 if v:
233 total += v * v * idf[i]
234 return math.sqrt(total)
237def rowmax_from_table(table: bytes) -> bytes:
238 """Derive the 256-byte row-maxima table of one dense model table.
240 Entry ``b1`` holds the maximum weight in the model's row for lead byte
241 ``b1``; a row with no bigrams yields 0.
242 """
243 return bytes(max(table[start : start + 256]) for start in range(0, 65536, 256))
246def read_models(models_path: Path) -> dict[str, dict[tuple[int, int], int]]:
247 """Read models.bin into the sparse per-model bigram dicts training uses.
249 The inverse of :func:`write_model_artifacts`'s models.bin output: only
250 non-zero weights appear in the dicts. A missing or empty file reads as
251 no models, so a first-ever training run and a full retrain look the
252 same to the caller.
254 :returns: Mapping of model key to ``{(b1, b2): weight}``.
255 :raises ValueError: If the file exists but is corrupt.
256 """
257 if not models_path.is_file():
258 return {}
259 data = models_path.read_bytes()
260 if not data:
261 return {}
262 tables, _norms = parse_models_bin(data)
263 models: dict[str, dict[tuple[int, int], int]] = {}
264 for name, table in tables.items():
265 bigrams: dict[tuple[int, int], int] = {}
266 for idx in range(65536):
267 weight = table[idx]
268 if weight > 0:
269 bigrams[(idx >> 8, idx & 0xFF)] = weight
270 models[name] = bigrams
271 return models
274def _idf_table(models: dict[str, dict[tuple[int, int], int]]) -> bytes:
275 """Compute the 65536-byte quantized IDF table over all models.
277 For each bigram index the byte holds a scaled inverse document
278 frequency: bigrams present in every model score 1 (minimal signal),
279 bigrams in exactly one model score 255 (maximum signal), and bigrams
280 in no model score 1 (unknown, neutral).
281 """
282 num_models = len(models)
283 doc_freq = [0] * 65536
284 for bigrams in models.values():
285 for b1, b2 in bigrams:
286 doc_freq[(b1 << 8) | b2] += 1
287 max_idf = math.log(num_models) if num_models > 1 else 1.0
288 scale = 254.0 / max_idf if max_idf > 0 else 0.0
289 idf_table = bytearray(65536)
290 for idx in range(65536):
291 df = doc_freq[idx]
292 if df > 0:
293 idf_val = math.log(num_models / df)
294 idf_table[idx] = max(1, round(idf_val * scale) + 1)
295 else:
296 idf_table[idx] = 1
297 return bytes(idf_table)
300def write_model_artifacts(
301 models: dict[str, dict[tuple[int, int], int]],
302 models_path: Path,
303) -> dict[str, int]:
304 """Write the model artifacts as one digest-locked set.
306 ``models.bin`` is written at *models_path* (whatever its basename);
307 ``rowmax.bin`` and ``idf.bin`` are written beside it under their fixed
308 names. ``rowmax.bin`` embeds the SHA-256 of the exact ``models.bin``
309 bytes written here and is deliberately written **last**: it is the
310 commit marker for the whole set, so an interrupted run leaves a digest
311 mismatch that :func:`parse_rowmax_bin` rejects at load time instead of
312 a silently stale sibling (``idf.bin`` is validated by size only).
314 :param models: Mapping of model key to sparse ``{(b1, b2): weight}``.
315 :param models_path: Destination path for models.bin.
316 :returns: Sizes in bytes, keyed by canonical artifact name
317 (``models.bin``, ``rowmax.bin``, ``idf.bin``, ``langnorm.bin``).
318 """
319 models_path.parent.mkdir(parents=True, exist_ok=True)
320 sorted_names = sorted(models)
321 idf_table = _idf_table(models)
323 header = bytearray(MODELS_MAGIC)
324 header += struct.pack("!I", len(sorted_names))
325 tables = bytearray()
326 rowmax_rows: list[bytes] = []
327 langnorms: list[float] = []
328 for name in sorted_names:
329 # Expand sparse dict to dense 65536-byte table and compute L2 norm
330 table = bytearray(65536)
331 sq_sum = 0
332 for (b1, b2), weight in models[name].items():
333 table[(b1 << 8) | b2] = weight
334 sq_sum += weight * weight
335 name_bytes = name.encode("utf-8")
336 header += struct.pack("!I", len(name_bytes)) + name_bytes
337 header += struct.pack("!d", math.sqrt(sq_sum))
338 # Freeze once and derive the row maxima from the exact bytes being
339 # serialized, so the two can never describe different tables.
340 frozen = bytes(table)
341 tables += frozen
342 rowmax_rows.append(rowmax_from_table(frozen))
343 langnorms.append(langnorm_from_table(frozen, idf_table))
345 models_blob = bytes(header) + zlib.compress(bytes(tables), 9)
346 models_path.write_bytes(models_blob)
347 models_digest = hashlib.sha256(models_blob).digest()
349 models_path.with_name("idf.bin").write_bytes(idf_table)
351 rowmax_blob = bytearray(ROWMAX_MAGIC)
352 rowmax_blob += models_digest
353 for row in rowmax_rows:
354 rowmax_blob += row
355 models_path.with_name("rowmax.bin").write_bytes(rowmax_blob)
357 langnorm_blob = langnorm_blob_for(langnorms, models_digest)
358 models_path.with_name("langnorm.bin").write_bytes(langnorm_blob)
360 return {
361 "models.bin": len(models_blob),
362 "rowmax.bin": len(rowmax_blob),
363 "idf.bin": len(idf_table),
364 "langnorm.bin": len(langnorm_blob),
365 }
368def langnorm_blob_for(langnorms: list[float], models_digest: bytes) -> bytes:
369 """Serialize langnorm.bin for norms in models.bin header order."""
370 blob = bytearray(LANGNORM_MAGIC)
371 blob += models_digest
372 for norm in langnorms:
373 blob += struct.pack("!d", norm)
374 return bytes(blob)