1"""Model loading and bigram scoring utilities.
2
3Note: ``from __future__ import annotations`` is intentionally omitted because
4this module is compiled with mypyc, which does not support PEP 563 string
5annotations.
6"""
7
8import array
9import functools
10import hashlib
11import importlib.resources
12import math
13import warnings
14
15from chardet import _kernel
16from chardet._kernel import dot_packed, pack_profile
17from chardet.models._format import parse_models_bin, parse_rowmax_bin, rowmax_from_table
18from chardet.registry import REGISTRY, lookup_encoding
19
20#: Shared empty packed buffers for profiles scored through another path.
21#: Safe to share because nothing ever writes to a profile's packed
22#: buffers --- dot_packed only reads them, and pack_profile always
23#: returns fresh arrays. Rebuilding these per profile allocated four
24#: throwaway arrays on the focused-profile path, which the surrounding
25#: code exists to keep allocation-free.
26_EMPTY_PACKED = (array.array("i"), array.array("i"))
27
28#: Whether _kernel was compiled into an extension for this install.
29#: The packed buffers only pay off when it was: read from C they are two
30#: contiguous int32 arrays, but read from the interpreter every element
31#: access boxes an int, which is slower than the dense list it replaced.
32#: A build with mypyc but no Cython would otherwise get the packed layout
33#: with an interpreted loop -- measured 3.5x slower than either path
34#: alone -- so profiles keep the dense table instead when it is absent.
35#:
36#: This flag makes two scoring paths and two storage shapes live in one
37#: file, and **no single test run exercises both**: a compiled run never
38#: takes the dense branch, an interpreted run never takes the packed one.
39#: Running the suite against both builds is therefore load-bearing rather
40#: than thorough --- it is the only thing standing between these branches
41#: and silently diverging. Change either one and check the other.
42_KERNEL_COMPILED = _kernel.__file__.endswith((".so", ".pyd"))
43
44# 256-entry membership table for whitespace bytes whose runs training
45# collapsed — native byte indexing under mypyc, used in the bigram-profile
46# hot loop. Covers the ASCII whitespace bytes (space, tab, LF, VT, FF, CR)
47# plus NBSP (0xA0): training's run collapse operates on decoded text where
48# regex \s matches all of these, so Latin models carry no run weight for
49# them and an uncollapsed input run would match only unrelated models
50# where the byte happens to be a letter. 0x85 (NEL in the ISO family) is
51# deliberately absent: it decodes to the ellipsis in windows-1252, whose
52# models legitimately carry ellipsis-run weight.
53_ASCII_WHITESPACE_TABLE = bytes(
54 1 if b in (0x20, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0xA0) else 0 for b in range(256)
55)
56
57# Encodings that map to exactly one language, derived from the registry.
58# Keyed by canonical name only — callers always use canonical names.
59_SINGLE_LANG_MAP: dict[str, str] = {}
60for _enc in REGISTRY.values():
61 if len(_enc.languages) == 1:
62 _SINGLE_LANG_MAP[_enc.name] = _enc.languages[0]
63
64
65@functools.cache
66def _load_models_data() -> tuple[dict[str, bytes], dict[str, float]]:
67 """Load and parse models.bin, returning (models, norms).
68
69 Cached: only reads from disk on first call.
70 """
71 ref = importlib.resources.files("chardet.models").joinpath("models.bin")
72 data = ref.read_bytes()
73
74 if not data:
75 warnings.warn(
76 "chardet models.bin is empty — statistical detection disabled; "
77 "reinstall chardet to fix",
78 RuntimeWarning,
79 stacklevel=2,
80 )
81 return {}, {}
82
83 return parse_models_bin(data)
84
85
86def load_models() -> dict[str, bytes]:
87 """Load all bigram models from the bundled models.bin file.
88
89 Each model is a bytes object of length 65536 (256*256).
90 Index: (b1 << 8) | b2 -> weight (0-255).
91
92 :returns: A dict mapping model key strings to 65536-byte lookup tables.
93 """
94 return _load_models_data()[0]
95
96
97def _build_enc_index(
98 models: dict[str, bytes],
99) -> dict[str, list[tuple[str | None, bytes, str]]]:
100 """Build a grouped index from a models dict.
101
102 :param models: Mapping of ``"lang/encoding"`` keys to 65536-byte tables.
103 :returns: Mapping of encoding name to ``[(lang, model, model_key), ...]``.
104 """
105 index: dict[str, list[tuple[str | None, bytes, str]]] = {}
106 for key, model in models.items():
107 lang, enc = key.split("/", 1)
108 index.setdefault(enc, []).append((lang, model, key))
109
110 # Resolve aliases: if a model key uses a non-canonical name,
111 # copy the entry under the canonical name.
112 for enc_name in list(index):
113 canonical = lookup_encoding(enc_name)
114 if canonical is not None and canonical not in index:
115 index[canonical] = index[enc_name]
116
117 return index
118
119
120@functools.cache
121def get_enc_index() -> dict[str, list[tuple[str | None, bytes, str]]]:
122 """Return a pre-grouped index mapping encoding name -> [(lang, model, model_key), ...]."""
123 return _build_enc_index(load_models())
124
125
126def infer_language(encoding: str) -> str | None:
127 """Return the language for a single-language encoding, or None.
128
129 :param encoding: The canonical encoding name.
130 :returns: An ISO 639-1 language code, or ``None`` if the encoding is
131 multi-language.
132 """
133 return _SINGLE_LANG_MAP.get(encoding)
134
135
136def has_model_variants(encoding: str) -> bool:
137 """Return True if the encoding has language variants in the model index.
138
139 :param encoding: The canonical encoding name.
140 :returns: ``True`` if bigram models exist for this encoding.
141 """
142 return encoding in get_enc_index()
143
144
145def _get_model_norms() -> dict[str, float]:
146 """Return cached L2 norms for all models, keyed by model key string."""
147 return _load_models_data()[1]
148
149
150@functools.cache
151def get_rowmax() -> dict[str, bytes]:
152 """Return per-model row-maximum tables for upper-bound prescreening.
153
154 For each model, entry ``b1`` of its 256-byte table holds the maximum
155 weight in the model's row for lead byte ``b1``. Because every bigram
156 weight is bounded by its row maximum, a dot product against the row
157 maxima (256 terms) upper-bounds the dot product against the full table
158 (65536 terms) — statistical scoring uses this to rule out candidate
159 models without scoring them fully.
160
161 Loads the precomputed ``rowmax.bin`` (written by
162 ``chardet.models._format.write_model_artifacts`` in the same model
163 order as ``models.bin``). The file starts with a ``CRM1`` magic and
164 the SHA-256 of the ``models.bin`` it was derived from: a stale or
165 mismatched file would silently under-estimate row maxima and break
166 the upper bound that pruning depends on, so anything that does not match
167 the *current* ``models.bin`` byte-for-byte is rejected and the tables
168 are derived from the models directly (slower, but always correct).
169 """
170 models = load_models()
171 files = importlib.resources.files("chardet.models")
172 try:
173 data = files.joinpath("rowmax.bin").read_bytes()
174 models_digest = hashlib.sha256(
175 files.joinpath("models.bin").read_bytes()
176 ).digest()
177 except (FileNotFoundError, OSError):
178 data = b""
179 models_digest = b""
180 # models preserves models.bin header order, matching rowmax.bin.
181 tables = parse_rowmax_bin(data, models_digest, list(models))
182 if tables is not None:
183 return tables
184 if models:
185 warnings.warn(
186 "chardet rowmax.bin is missing or does not match models.bin; "
187 "deriving row maxima from the models (slower startup)",
188 RuntimeWarning,
189 stacklevel=2,
190 )
191 return {key: rowmax_from_table(table) for key, table in models.items()}
192
193
194@functools.cache
195def get_idf_weights() -> bytes:
196 """Return a 65536-byte IDF weight table for bigram profile construction.
197
198 Loads a precomputed table from ``idf.bin`` (generated at training time).
199 For each bigram index, the weight reflects how discriminative that bigram
200 is across all models:
201
202 - Bigrams in every model (common ASCII) → weight 1 (minimal signal)
203 - Bigrams in one model → weight 255 (maximum signal)
204 - Bigrams not in any model → weight 1 (unknown, treat as neutral)
205 """
206 ref = importlib.resources.files("chardet.models").joinpath("idf.bin")
207 data = ref.read_bytes()
208 if len(data) != 65536:
209 warnings.warn(
210 f"chardet idf.bin has wrong size ({len(data)}), "
211 "falling back to uniform weights",
212 RuntimeWarning,
213 stacklevel=2,
214 )
215 return b"\x01" * 65536
216 return data
217
218
219class BigramProfile:
220 """Pre-computed bigram frequency distribution for a data sample.
221
222 Computing this once and reusing it across all models reduces per-model
223 scoring from O(n) to O(distinct_bigrams).
224
225 Each bigram is weighted by its IDF (inverse document frequency) across all
226 models — bigrams unique to few models get high weight, bigrams common to
227 all models get weight 1. ``nonzero`` lists the indices carrying weight,
228 in first-encounter order.
229
230 The weights are reachable two ways, and which one a profile fills
231 depends on how it was built:
232
233 * ``idx_arr``/``val_arr`` — parallel ``array('i')`` buffers, what
234 :func:`score_with_profile` reads for a streaming profile. Filled by
235 the streaming constructor only.
236 * ``values`` — a plain list parallel to ``nonzero``, filled by
237 :meth:`from_weighted_freq` for the small focused profiles confusion
238 resolution builds, which are scored inline instead.
239 ``row_freq`` aggregates the weights by lead byte (256 entries) and
240 ``nonzero_rows`` lists the lead bytes with non-zero total; together with
241 per-model row maxima (:func:`get_rowmax`) they let statistical scoring
242 compute a cheap upper bound on a model's score.
243
244 **Input limit.** Weights are packed as int32, which holds any value an
245 input of at most 16 MB can produce (``255`` per occurrence against a
246 2.1-billion ceiling). Detection truncates to ``max_bytes`` long before
247 that; construct a profile directly from a larger buffer and packing
248 raises :exc:`OverflowError`.
249 """
250
251 __slots__ = (
252 "freq",
253 "idx_arr",
254 "input_norm",
255 "nonzero",
256 "nonzero_rows",
257 "row_freq",
258 "val_arr",
259 "values",
260 "weight_sum",
261 )
262
263 def __init__(self, data: bytes) -> None:
264 """Compute the bigram frequency distribution for *data*.
265
266 Each bigram is weighted by its IDF (inverse document frequency) across
267 all loaded models. Bigrams unique to few models get high weight;
268 bigrams common to all models get weight 1.
269
270 :param data: The raw byte data to profile.
271 """
272 total_bigrams = len(data) - 1
273 if total_bigrams <= 0:
274 # Empty lists, not [0]*65536: a no-op profile allocates nothing.
275 self.freq: list[int] = []
276 self.nonzero: list[int] = []
277 self.values: list[int] = []
278 self.idx_arr, self.val_arr = _EMPTY_PACKED
279 self.row_freq: list[int] = []
280 self.nonzero_rows: list[int] = []
281 self.weight_sum: int = 0
282 self.input_norm: float = 0.0
283 return
284
285 idf = get_idf_weights()
286 freq: list[int] = [0] * 65536
287 nonzero: list[int] = []
288 w_sum = 0
289 for i in range(total_bigrams):
290 b1 = data[i]
291 b2 = data[i + 1]
292 # Skip repeated-whitespace bigrams (equivalent to collapsing
293 # whitespace runs, which training does before counting): padding
294 # and indentation carry no encoding signal, and models trained
295 # on lossy transcodes can carry spurious weight for them.
296 if b1 == b2 and _ASCII_WHITESPACE_TABLE[b1]:
297 continue
298 idx = (b1 << 8) | b2
299 w = idf[idx]
300 if freq[idx] == 0:
301 nonzero.append(idx)
302 freq[idx] += w
303 w_sum += w
304 # No ``values``: this path already has the dense table, and
305 # materializing a parallel list over every nonzero bigram costs
306 # more than the sparse constructor's allocation saves.
307 self._finish(freq, nonzero, [], w_sum)
308
309 def _finish(
310 self,
311 freq: list[int],
312 nonzero: list[int],
313 values: list[int],
314 weight_sum: int,
315 ) -> None:
316 """Store the frequency data and derive the norm and row aggregates.
317
318 Single finalization path shared by both constructors so pruning
319 fields cannot silently diverge between them.
320
321 Exactly one of *freq* and *values* carries the weights, never both.
322 The streaming constructor passes the dense *freq* table it had to
323 build and leaves *values* empty; the sparse constructor fills
324 *values* parallel to *nonzero* and passes no table, which is what
325 lets it skip a 65536-entry allocation per call.
326
327 Neither is retained. *freq* is read here to derive the norm, the row
328 aggregates and the packed buffers, and is then dropped -- scoring
329 reads ``idx_arr``/``val_arr`` or ``values``, never the dense table,
330 so keeping it alive would pin 512 KB per profile for nothing.
331 """
332 self.nonzero = nonzero
333 self.values = values
334 self.weight_sum = weight_sum
335 norm_sq = 0
336 row_freq: list[int] = [0] * 256
337 if values:
338 for i in range(len(nonzero)):
339 v = values[i]
340 norm_sq += v * v
341 row_freq[nonzero[i] >> 8] += v
342 else:
343 for idx in nonzero:
344 v = freq[idx]
345 norm_sq += v * v
346 row_freq[idx >> 8] += v
347 self.input_norm = math.sqrt(norm_sq)
348 self.row_freq = row_freq
349 self.nonzero_rows = [b1 for b1 in range(256) if row_freq[b1]]
350 # Dense profiles are scored through the packed kernel when it is
351 # compiled; sparse ones (a median of eight bigrams) are scored inline
352 # either way and skip the packing.
353 if values or not _KERNEL_COMPILED:
354 self.idx_arr, self.val_arr = _EMPTY_PACKED
355 else:
356 self.idx_arr, self.val_arr = pack_profile(nonzero, freq)
357 # Retained only when scoring will read it -- see _KERNEL_COMPILED.
358 self.freq = [] if (values or _KERNEL_COMPILED) else freq
359
360 @classmethod
361 def from_weighted_freq(cls, weighted_freq: dict[int, int]) -> "BigramProfile":
362 """Create a BigramProfile from pre-computed weighted frequencies.
363
364 Computes ``weight_sum`` and ``input_norm`` from *weighted_freq* to
365 ensure consistency between the stored fields.
366
367 Deliberately does not build the dense 65536-entry ``freq`` table
368 the streaming constructor uses. Callers here pass a handful of
369 bigrams — confusion resolution's focused profiles hold a median of
370 eight — and allocating a 65536-element list per call cost about
371 24us, which measured as roughly 40% of the whole bigram-rescore
372 stage. Scoring reads ``nonzero``/``values``, so the table is
373 never needed.
374
375 :param weighted_freq: Mapping of bigram index to weighted count.
376 :returns: A new :class:`BigramProfile` instance.
377 """
378 profile = cls(b"")
379 nonzero: list[int] = []
380 values: list[int] = []
381 w_sum = 0
382 for idx, count in weighted_freq.items():
383 if count:
384 nonzero.append(idx)
385 values.append(count)
386 w_sum += count
387 profile._finish([], nonzero, values, w_sum)
388 return profile
389
390
391def score_with_profile(
392 profile: BigramProfile,
393 model: "bytes | bytearray | memoryview",
394 model_key: str = "",
395) -> float:
396 """Score a pre-computed bigram profile against a single model using cosine similarity.
397
398 ``bytearray``/``memoryview`` tables are accepted for compatibility but
399 copied to ``bytes`` first: the narrow type lets mypyc compile the
400 dot-product loop with native byte indexing (and keeps compiled and
401 pure-Python installs accepting the same argument types).
402 """
403 if not isinstance(model, bytes):
404 model = bytes(model)
405 if profile.input_norm == 0.0:
406 return 0.0
407 norms = _get_model_norms()
408 model_norm = norms.get(model_key) if model_key else None
409 if model_norm is None:
410 sq_sum = 0
411 for i in range(65536):
412 v = model[i]
413 if v:
414 sq_sum += v * v
415 model_norm = math.sqrt(sq_sum)
416 if model_norm == 0.0:
417 return 0.0
418 nonzero = profile.nonzero
419 values = profile.values
420 if values:
421 # Focused confusion profiles hold a median of eight bigrams, so this
422 # branch stays inline: a call into _kernel would cost about what the
423 # loop does. It measured 1.3% of compiled runtime.
424 dot = 0
425 for i in range(len(nonzero)):
426 dot += model[nonzero[i]] * values[i]
427 elif _KERNEL_COMPILED:
428 dot = dot_packed(profile.idx_arr, profile.val_arr, model)
429 else:
430 # No compiled kernel: the dense table beats packed buffers here,
431 # because a list returns a cached int where array('i') boxes a new one.
432 dot = 0
433 freq = profile.freq
434 for idx in nonzero:
435 dot += model[idx] * freq[idx]
436 return dot / (model_norm * profile.input_norm)
437
438
439#: ADR-0005's hand-audited rare-language set, shared with the encoding-side
440#: arbitration gate in ``pipeline.postprocess``. One set, two gates: the
441#: deployment evidence that justifies membership is recorded in the ADR and
442#: any new genuine specimen forces a re-audit of both.
443RARE_LANGUAGES: frozenset[str] = frozenset({"gd", "cy", "ga", "br"})
444
445#: The ANSI/ASCII-art pseudo-language. Never a valid demotion target for
446#: the thin-rare band: swapping a rare label for "zxx" would tell the
447#: orchestrator the text has no linguistic content at all.
448ART_LANGUAGE = "zxx"
449
450#: The thin-margin band for ``demote_thin_rare``: inputs shorter than this
451#: are where bigram cosines stop discriminating (apostrophe-rich English
452#: snippets score as Gaelic). The smallest genuine rare-language files in
453#: the test corpus are 253 bytes (legacy path) and 560 bytes (fill path),
454#: 2x and 4.4x above this gate; the callers in ``pipeline.language`` own
455#: the length judgment because only they know the pre-transcoding size.
456_THIN_RARE_MAX_BYTES = 128
457
458#: Maximum lead over the best prevalent-language variant for a rare win on
459#: a short input to count as noise. Measured mislabels win by <= 0.021.
460#: This is NOT a safety floor: genuine gd/cy snippets measure 0.07+ and the
461#: nearest measured genuine escape is Welsh at 0.043, but short Breton and
462#: Irish can sit inside the band and are the accepted casualty class per
463#: ADR-0005's addendum — widening the margin widens that class.
464_THIN_RARE_MARGIN = 0.03
465
466
467def score_best_language(
468 data: bytes,
469 encoding: str,
470 profile: BigramProfile | None = None,
471 *,
472 demote_thin_rare: bool = False,
473) -> tuple[float, str | None]:
474 """Score data against all language variants of an encoding.
475
476 Returns (best_score, best_language). Uses a pre-grouped index for O(L)
477 lookup where L is the number of language variants for the encoding.
478
479 If *profile* is provided, it is reused instead of recomputing the bigram
480 frequency distribution from *data*.
481
482 :param data: The raw byte data to score.
483 :param encoding: The canonical encoding name to match against.
484 :param profile: Optional pre-computed :class:`BigramProfile` to reuse.
485 :param demote_thin_rare: Pass true only when the caller has judged the
486 *original* input thin (under :data:`_THIN_RARE_MAX_BYTES` before
487 any transcoding). A :data:`RARE_LANGUAGES` winner that leads the
488 best prevalent-language variant by less than the measured noise
489 band is then reported under the prevalent language instead. The
490 length judgment deliberately lives with the caller: this function
491 may receive transcoded bytes or a profile without its source data,
492 so ``len(data)`` here is not a reliable proxy for input size.
493 Language-fill callers pass this; encoding-ranking callers must not,
494 so that candidate ordering stays byte-identical.
495 :returns: A ``(score, language)`` tuple. The score is always the best
496 cosine similarity across variants; the language matches it except
497 when ``demote_thin_rare`` fires, in which case the label is the
498 best prevalent-language variant's while the score remains the
499 rare winner's.
500 """
501 if not data and profile is None:
502 return 0.0, None
503
504 index = get_enc_index()
505 variants = index.get(encoding)
506 if variants is None:
507 return 0.0, None
508
509 if profile is None:
510 profile = BigramProfile(data)
511
512 best_score = 0.0
513 best_lang: str | None = None
514 best_prevalent = 0.0
515 best_prevalent_lang: str | None = None
516 for lang, model, model_key in variants:
517 s = score_with_profile(profile, model, model_key)
518 if s > best_score:
519 best_score = s
520 best_lang = lang
521 if (
522 demote_thin_rare
523 and lang not in RARE_LANGUAGES
524 and lang != ART_LANGUAGE
525 and s > best_prevalent
526 ):
527 best_prevalent = s
528 best_prevalent_lang = lang
529
530 if (
531 demote_thin_rare
532 and best_lang is not None
533 and best_lang in RARE_LANGUAGES
534 and best_prevalent_lang is not None
535 and best_score - best_prevalent < _THIN_RARE_MARGIN
536 ):
537 best_lang = best_prevalent_lang
538
539 return best_score, best_lang