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), and ``idf.bin`` (a 65536-byte
8quantized IDF table). Keeping both directions in one module means a
9format change cannot land on one side only: ``scripts/train.py`` calls
10:func:`write_model_artifacts` and :func:`read_models`, while
11:mod:`chardet.models` calls :func:`parse_models_bin` and
12:func:`parse_rowmax_bin` behind its cached accessors.
14Note: ``from __future__ import annotations`` is intentionally omitted
15because this module is compiled with mypyc, which does not support PEP 563
16string annotations.
17"""
19import hashlib
20import math
21import struct
22import zlib
23from pathlib import Path
25#: models.bin magic: the v2 dense zlib-compressed format.
26MODELS_MAGIC = b"CMD2"
27#: rowmax.bin layout: magic, SHA-256 of the matching models.bin, then one
28#: 256-byte row-maxima table per model in models.bin header order.
29ROWMAX_MAGIC = b"CRM1"
30ROWMAX_HEADER_SIZE = 4 + 32
32_unpack_uint32 = struct.Struct(">I").unpack_from
33_unpack_float64 = struct.Struct(">d").unpack_from
36def _decompress_tables(
37 data: bytes, offset: int, names: list[str], chunk_size: int = 262144
38) -> dict[str, bytes]:
39 """Decompress the model tables from ``data[offset:]``, one per name.
41 Each model is stored as its own bytes object rather than a memoryview
42 slice of one big blob: mypyc compiles bytes indexing in the scoring hot
43 loop to a native C array access, while memoryview indexing goes through
44 a boxed generic call. Decompression is incremental, one 64 KB table at
45 a time — materializing the whole multi-megabyte blob and slicing it
46 would transiently double the allocation and strand the freed pages in
47 process RSS. Trailing compressed bytes are ignored, as with whole-blob
48 ``zlib.decompress``; the decompressed size is validated instead.
50 :raises ValueError: If the decompressed size is not exactly
51 ``len(names) * 65536``.
52 """
53 num_models = len(names)
54 expected_size = num_models * 65536
55 decomp = zlib.decompressobj()
56 models: dict[str, bytes] = {}
57 table = bytearray()
58 produced = 0
59 pos = offset
60 end = len(data)
61 flushed = False
62 while len(models) < num_models:
63 need = 65536 - len(table)
64 if decomp.unconsumed_tail:
65 piece = decomp.decompress(decomp.unconsumed_tail, need)
66 elif pos < end:
67 chunk = data[pos : pos + chunk_size]
68 pos += len(chunk)
69 piece = decomp.decompress(chunk, need)
70 elif not flushed:
71 flushed = True
72 piece = decomp.flush()
73 if not piece:
74 break # stream exhausted early -> size mismatch below
75 else:
76 break # stream exhausted early -> size mismatch below
77 produced += len(piece)
78 table += piece
79 if len(table) == 65536:
80 models[names[len(models)]] = bytes(table)
81 table.clear()
82 # Unreachable with CPython's zlib — decompress() pieces are capped at
83 # ``need``, and a flush() strand (the tail of one back-reference cut
84 # mid-copy, at most 258 bytes) only ever lands in a fresh table —
85 # but a decompressor that flushed more than asked must not corrupt
86 # the table split silently.
87 elif len(table) > 65536: # pragma: no cover
88 break # oversized flush -> size mismatch below
89 if len(models) == num_models:
90 # Drain any leftover decompressed output so extra data is caught,
91 # including surplus tables in compressed input not yet fed to the
92 # decompressor. Bytes after the stream's end marker (decomp.eof)
93 # are ignored, matching whole-blob zlib.decompress behavior.
94 extra = b""
95 if decomp.unconsumed_tail:
96 extra = decomp.decompress(decomp.unconsumed_tail, 1)
97 while not extra and not decomp.eof and pos < end:
98 chunk = data[pos : pos + chunk_size]
99 pos += len(chunk)
100 extra = decomp.decompress(chunk, 1)
101 if not extra and not decomp.eof and not flushed:
102 extra = decomp.flush()
103 produced += len(extra)
104 if produced != expected_size or len(models) != num_models:
105 msg = (
106 f"corrupt models.bin: decompressed size {produced} "
107 f"!= expected {expected_size}"
108 )
109 raise ValueError(msg)
110 return models
113def parse_models_bin(
114 data: bytes,
115) -> tuple[dict[str, bytes], dict[str, float]]:
116 """Parse the v2 dense zlib-compressed models.bin format.
118 :param data: Raw bytes of models.bin (must be non-empty).
119 :returns: A ``(models, norms)`` tuple.
120 :raises ValueError: If the data is corrupt or truncated.
121 """
122 try:
123 if data[:4] != MODELS_MAGIC:
124 msg = "corrupt models.bin: missing CMD2 magic"
125 raise ValueError(msg)
127 offset = 4 # skip magic
128 (num_models,) = _unpack_uint32(data, offset)
129 offset += 4
131 if num_models > 10_000:
132 msg = f"corrupt models.bin: num_models={num_models} exceeds limit"
133 raise ValueError(msg)
135 names: list[str] = []
136 norms: dict[str, float] = {}
137 for _ in range(num_models):
138 (name_len,) = _unpack_uint32(data, offset)
139 offset += 4
140 if name_len > 256:
141 msg = f"corrupt models.bin: name_len={name_len} exceeds 256"
142 raise ValueError(msg)
143 name = data[offset : offset + name_len].decode("utf-8")
144 offset += name_len
145 (norm,) = _unpack_float64(data, offset)
146 offset += 8
147 names.append(name)
148 norms[name] = norm
150 models = _decompress_tables(data, offset, names)
152 except zlib.error as e:
153 msg = f"corrupt models.bin: {e}"
154 raise ValueError(msg) from e
155 except (struct.error, UnicodeDecodeError) as e:
156 msg = f"corrupt models.bin: {e}"
157 raise ValueError(msg) from e
159 return models, norms
162def parse_rowmax_bin(
163 data: bytes, models_digest: bytes, model_keys: list[str]
164) -> dict[str, bytes] | None:
165 """Parse rowmax.bin into per-model row-maxima tables.
167 :param data: Raw bytes of rowmax.bin.
168 :param models_digest: SHA-256 digest of the current models.bin bytes.
169 :param model_keys: Model keys in models.bin header order.
170 :returns: Mapping of model key to its 256-byte row-maxima table, or
171 ``None`` when the magic, digest, or size does not match — a stale
172 or mismatched file would silently under-estimate row maxima and
173 break the upper bound that prescreening depends on.
174 """
175 if (
176 data[:4] == ROWMAX_MAGIC
177 and data[4:ROWMAX_HEADER_SIZE] == models_digest
178 and len(data) == ROWMAX_HEADER_SIZE + len(model_keys) * 256
179 ):
180 return {
181 key: data[ROWMAX_HEADER_SIZE + i * 256 : ROWMAX_HEADER_SIZE + (i + 1) * 256]
182 for i, key in enumerate(model_keys)
183 }
184 return None
187def rowmax_from_table(table: bytes) -> bytes:
188 """Derive the 256-byte row-maxima table of one dense model table.
190 Entry ``b1`` holds the maximum weight in the model's row for lead byte
191 ``b1``; a row with no bigrams yields 0.
192 """
193 return bytes(max(table[start : start + 256]) for start in range(0, 65536, 256))
196def read_models(models_path: Path) -> dict[str, dict[tuple[int, int], int]]:
197 """Read models.bin into the sparse per-model bigram dicts training uses.
199 The inverse of :func:`write_model_artifacts`'s models.bin output: only
200 non-zero weights appear in the dicts. A missing or empty file reads as
201 no models, so a first-ever training run and a full retrain look the
202 same to the caller.
204 :returns: Mapping of model key to ``{(b1, b2): weight}``.
205 :raises ValueError: If the file exists but is corrupt.
206 """
207 if not models_path.is_file():
208 return {}
209 data = models_path.read_bytes()
210 if not data:
211 return {}
212 tables, _norms = parse_models_bin(data)
213 models: dict[str, dict[tuple[int, int], int]] = {}
214 for name, table in tables.items():
215 bigrams: dict[tuple[int, int], int] = {}
216 for idx in range(65536):
217 weight = table[idx]
218 if weight > 0:
219 bigrams[(idx >> 8, idx & 0xFF)] = weight
220 models[name] = bigrams
221 return models
224def _idf_table(models: dict[str, dict[tuple[int, int], int]]) -> bytes:
225 """Compute the 65536-byte quantized IDF table over all models.
227 For each bigram index the byte holds a scaled inverse document
228 frequency: bigrams present in every model score 1 (minimal signal),
229 bigrams in exactly one model score 255 (maximum signal), and bigrams
230 in no model score 1 (unknown, neutral).
231 """
232 num_models = len(models)
233 doc_freq = [0] * 65536
234 for bigrams in models.values():
235 for b1, b2 in bigrams:
236 doc_freq[(b1 << 8) | b2] += 1
237 max_idf = math.log(num_models) if num_models > 1 else 1.0
238 scale = 254.0 / max_idf if max_idf > 0 else 0.0
239 idf_table = bytearray(65536)
240 for idx in range(65536):
241 df = doc_freq[idx]
242 if df > 0:
243 idf_val = math.log(num_models / df)
244 idf_table[idx] = max(1, round(idf_val * scale) + 1)
245 else:
246 idf_table[idx] = 1
247 return bytes(idf_table)
250def write_model_artifacts(
251 models: dict[str, dict[tuple[int, int], int]],
252 models_path: Path,
253) -> dict[str, int]:
254 """Write the model artifacts as one digest-locked set.
256 ``models.bin`` is written at *models_path* (whatever its basename);
257 ``rowmax.bin`` and ``idf.bin`` are written beside it under their fixed
258 names. ``rowmax.bin`` embeds the SHA-256 of the exact ``models.bin``
259 bytes written here and is deliberately written **last**: it is the
260 commit marker for the whole set, so an interrupted run leaves a digest
261 mismatch that :func:`parse_rowmax_bin` rejects at load time instead of
262 a silently stale sibling (``idf.bin`` is validated by size only).
264 :param models: Mapping of model key to sparse ``{(b1, b2): weight}``.
265 :param models_path: Destination path for models.bin.
266 :returns: Sizes in bytes, keyed by canonical artifact name
267 (``models.bin``, ``rowmax.bin``, ``idf.bin``).
268 """
269 models_path.parent.mkdir(parents=True, exist_ok=True)
270 sorted_names = sorted(models)
272 header = bytearray(MODELS_MAGIC)
273 header += struct.pack("!I", len(sorted_names))
274 tables = bytearray()
275 rowmax_rows: list[bytes] = []
276 for name in sorted_names:
277 # Expand sparse dict to dense 65536-byte table and compute L2 norm
278 table = bytearray(65536)
279 sq_sum = 0
280 for (b1, b2), weight in models[name].items():
281 table[(b1 << 8) | b2] = weight
282 sq_sum += weight * weight
283 name_bytes = name.encode("utf-8")
284 header += struct.pack("!I", len(name_bytes)) + name_bytes
285 header += struct.pack("!d", math.sqrt(sq_sum))
286 # Freeze once and derive the row maxima from the exact bytes being
287 # serialized, so the two can never describe different tables.
288 frozen = bytes(table)
289 tables += frozen
290 rowmax_rows.append(rowmax_from_table(frozen))
292 models_blob = bytes(header) + zlib.compress(bytes(tables), 9)
293 models_path.write_bytes(models_blob)
295 idf_table = _idf_table(models)
296 models_path.with_name("idf.bin").write_bytes(idf_table)
298 rowmax_blob = bytearray(ROWMAX_MAGIC)
299 rowmax_blob += hashlib.sha256(models_blob).digest()
300 for row in rowmax_rows:
301 rowmax_blob += row
302 models_path.with_name("rowmax.bin").write_bytes(rowmax_blob)
304 return {
305 "models.bin": len(models_blob),
306 "rowmax.bin": len(rowmax_blob),
307 "idf.bin": len(idf_table),
308 }