#!/usr/bin/env python3
# Copyright 2026 The Source-1 Authors
# SPDX-License-Identifier: Apache-2.0
"""Source-1: score text as language-model pretraining data.

A standalone loader and scorer for Source-1, an mmBERT-base encoder fine-tuned with 13 scoring heads. It needs
only torch, transformers, safetensors and tokenizers (no remote code, no other project code), and runs on a GPU or a
CPU.

    from source1 import Source1

    model = Source1.from_pretrained("path/to/Source-1")     # a local directory or a Hugging Face repo id
    model = Source1.from_pretrained("path/to/Source-1", precision="fp32")  # the full-precision weights instead
    doc = model.score("Photosynthesis is how plants ...", title="Photosynthesis")
    doc["overall"], doc["keep"], doc["educational_value"]
    docs = model.score_batch(["first document", {"text": "second document", "title": "A title"}])

    python source1.py --input docs.jsonl --output scores.jsonl     # one JSON object per line, "text" field
    python source1.py --input page.txt                             # a .txt file is one document
    python source1.py --input docs.jsonl.gz                        # gzip, bzip2 and xz files are decompressed

Output: one flat dict per document
----------------------------------
format, topic, content_type    labels: the most likely value (10, 15 and 4 values; see source1.json)
educational_value, reasoning_depth, writing_quality, information_density, reliability
                               quality scores 0-5, higher is better
spam_seo, boilerplate, toxicity
                               red flags 0-5, higher is worse
code_quality, math_quality     gated scores 0-5, or None when the document is not code / not math
overall                        composite 0-5: weighted quality, blended with the gated scores that apply, minus
                               red-flag penalties (see ``composite``)
keep                           False when the drop line of calibration.json matches (see ``DropLine``)
drop_reasons                   the drop-line conditions that matched ([] when kept)
parts, tokens, truncated       chunks the document was split into, its length in tokens, whether any chunk was
                               longer than the model's 8,192 tokens and was cut
chunks                         only when parts > 1: the same fields per chunk, with its part number, its character
                               span in the cleaned text (``clean_text``) and its token count
label_dist, ranges             only when parts > 1: each label value's token share; each score's [min, max]

Scores are expected values (the six levels 0-5 weighted by their probabilities), so they are fractional, rounded
to 3 decimals like every other number here. The scores are the model's own; ``apply_offsets=True`` adds the small
calibration offsets of calibration.json to the quality scores (about 0.02 at most; off by default).

How a document is scored (the same steps the model was trained and evaluated with)
------------------------------------------------------------------------------------
1. ``clean_text``: line endings to "\\n", ASCII control characters other than tab and newline dropped, Unicode NFC,
   trailing spaces dropped, at most two blank lines in a row.
2. ``split_text``: a document longer than 7,808 tokens is split into balanced chunks, each ending at the most
   natural boundary near its ideal end (headings, then paragraphs, lines, sentences, spaces; definitions in code).
3. ``build_input``: each chunk gets a one-line header, a blank line, then the chunk text:
       Source: dataset record | Title: <title> | Part 2 of 3 of a longer document
   Every training input had a Source line, almost always "dataset record", so that is the default; code files had
   "<language> source file" (pass ``code_language="Python"``). Title is added when given, Part when the document has
   more than one chunk. ``url`` is accepted but not shown to the model unless ``show_url=True``: no training input had
   one. On 1,261 held-out benchmark chunks (495 graded held-out chunks from the test split and 766 exam chunks),
   dropping the whole header moved overall by 0.03 on average (at most 0.655), dropping a title by 0.05 on the chunks
   that had one; adding a URL moved it by up to 0.7 and did not improve its rank agreement with the independent
   graders of the model card's evaluation.
4. ``collapse_spaces``: runs of spaces and tabs become one space; at most two empty lines in a row.
5. Tokenized with <bos> and <eos>, at most 8,192 tokens (a longer input is cut at its end).
6. The final hidden states are mean-pooled over the tokens; one linear head per field.
7. A document's chunks are combined by token-weighted vote (labels) and token-weighted mean (scores; toxicity takes
   the maximum; gated scores average over the chunks where they apply). ``overall`` and ``keep`` are then computed
   on the combined scores.

Weights and precision
---------------------
Two copies of the backbone weights: ``model.safetensors`` in bfloat16 (the default, ``precision="bf16"``, half the
size) and ``model.fp32.safetensors`` in float32 (``precision="fp32"``, the full-precision copy). ``precision`` picks
the file; ``dtype`` picks what the model computes in. By default (``dtype="auto"``) it computes in bfloat16 on a GPU
with native bfloat16 (NVIDIA Ampere and newer), as in the project's own evaluation, and in float32 on older GPUs and on
a CPU, where bfloat16 weights are upcast to float32 (pass ``dtype="bf16"`` to compute in bfloat16 there too). Computing
in bfloat16, both files give the same scores, because the bfloat16 file holds exactly the float32 weights rounded to
bfloat16. Computing in float32, the bfloat16 weights move scores slightly away from the float32 weights' (on the same
1,261 benchmark chunks: up to 0.03 on overall and 0.08 on a single field, 3 labels and 1 keep decision changed); use
``precision="fp32", dtype="fp32"`` for full float32. float16 is refused: the
mean pooling overflows its range on long inputs and gives NaN scores. In bfloat16 a chunk's scores depend slightly on
which other inputs share its batch (scoring each benchmark chunk alone instead of in the default batches moved overall
by up to 0.04 and a single field by up to 0.10; 3 labels changed, no keep decision), and float32 differs from bfloat16
by a similar amount (up to 0.03 on overall and 0.09 on a single field; 4 labels and 1 keep decision changed); for
scores that do not depend on the batch, compute in float32 (``dtype="fp32"``) or use ``batch_tokens=1`` (one input per
forward pass).
"""

from __future__ import annotations

import argparse
import ast
import bz2
import gzip
import json
import lzma
import math
import operator
import re
import sys
import time
import unicodedata
import zlib
from bisect import bisect_left
from collections.abc import Iterable, Iterator, Mapping
from pathlib import Path
from typing import Any

import torch
import torch.nn as nn

__version__ = "1.0.0"

MAX_LENGTH = 8192  # tokens per model input, <bos> and <eos> included
CHUNK_TOKENS = 7808  # document tokens per chunk: 8,192 minus 384 kept for the header (as in training)
DEFAULT_SOURCE = "dataset record"
DEFAULT_BATCH_TOKENS = 65536  # padded tokens per forward pass on a GPU
DEFAULT_BATCH_TOKENS_CPU = 16384  # on a CPU: a quarter of the memory, and no slower there
LEVELS = (0, 1, 2, 3, 4, 5)
# The backbone weights of each precision: bfloat16 (the default) and the full-precision float32 copy. The fp32 file
# follows transformers' variant naming (model.<variant>.safetensors), so AutoModel loads it with variant="fp32".
WEIGHTS = {"bf16": "model.safetensors", "fp32": "model.fp32.safetensors"}
DEFAULT_PRECISION = "bf16"
FILES = ("config.json", "heads.safetensors", "source1.json", "tokenizer.json")  # needed besides the weights
CALIBRATION = "calibration.json"  # the calibrated drop line and offsets; required unless drop_line is given
# What from_pretrained also downloads for a Hub repo id (with the weights of the chosen precision only): the
# calibration and the license files.
HUB_EXTRA = (CALIBRATION, "tokenizer_config.json", "LICENSE", "NOTICE", "AUTHORS", "CREDITS_BOOKS.tsv")


def hub_files(precision: str = DEFAULT_PRECISION) -> tuple[str, ...]:
    """The files from_pretrained downloads from a Hugging Face repo for ``precision``."""
    return (WEIGHTS[resolve_precision(precision)], *FILES, *HUB_EXTRA)

# --------------------------------------------------------------------------------------------- text


_JUNK = re.compile("[\x00-\x08\x0b\x0e-\x1f\x7f\ud800-\udfff\ufeff\u200b\ufffe\uffff]")
# The lookbehind keeps this linear: without it, a long run of spaces that ends in no newline is scanned again from
# each of its positions, so a page of spaces takes quadratic time. The result is the same.
_TRAILING_WS = re.compile(r"(?<![ \t])[ \t]+\n")
_BLANK_LINES = re.compile(r"\n{4,}")
_SPACE_RUN = re.compile(r"[ \t]{2,}")
_BLANK_RUN = re.compile(r"\n(?:[ \t]*\n){3,}")


def clean_text(text: str) -> str:
    """The document as the scorer sees it before chunking: "\\n" line endings (a form feed counts as a paragraph
    break), ASCII control characters other than tab and newline, lone surrogates, zero-width spaces and byte-order
    marks dropped (C1 control characters such as U+0085 are kept), Unicode NFC, no trailing spaces, at most two
    blank lines in a row, no blank lines at either end."""
    text = text.replace("\r\n", "\n").replace("\r", "\n").replace("\x0c", "\n\n")
    text = _JUNK.sub("", text)
    text = unicodedata.normalize("NFC", text)
    text = _TRAILING_WS.sub("\n", text)
    text = _BLANK_LINES.sub("\n\n\n", text)
    return text.strip("\n").rstrip()


def collapse_spaces(text: str) -> str:
    """The model's input normalization: runs of 2+ spaces/tabs become one space, 3+ blank lines in a row (lines
    holding only spaces or tabs) become two empty lines; newlines are kept."""
    return _BLANK_RUN.sub("\n\n\n", _SPACE_RUN.sub(" ", text))


_SURROGATES = re.compile("[\ud800-\udfff]")


def _from_bytes(x: Any) -> Any:
    """bytes and bytearray decoded as UTF-8 (invalid bytes become U+FFFD); anything else unchanged."""
    return x.decode("utf-8", errors="replace") if isinstance(x, (bytes, bytearray)) else x


def _one_line(s: Any, limit: int) -> str:
    """A header value on one line: lone surrogates (which cannot be tokenized) dropped, every run of whitespace
    (newlines included) made one space, cut to ``limit`` characters."""
    s = re.sub(r"\s+", " ", _SURROGATES.sub("", str(_from_bytes(s)))).strip()
    return s if len(s) <= limit else s[: limit - 1] + "\u2026"


def source_description(source_type: str | None = None, code_language: str | None = None) -> str:
    """The header's Source value: ``source_type`` when given ("" leaves the Source line out), else
    "<code_language> source file" for code, else "dataset record"."""
    if source_type is not None:
        return str(source_type)
    if code_language:
        return f"{code_language} source file"
    return DEFAULT_SOURCE


def build_input(chunk: str, *, source_type: str | None = DEFAULT_SOURCE, title: str | None = None,
                url: str | None = None, part: int = 1, parts: int = 1) -> str:
    """One chunk as the model reads it (before ``collapse_spaces``): a header line, a blank line, the chunk.

    >>> build_input("Text.", title="On  rivers", part=2, parts=3)
    'Source: dataset record | Title: On rivers | Part 2 of 3 of a longer document\\n\\nText.'
    """
    head = []
    if source_type:
        head.append(f"Source: {_one_line(source_type, 200)}")
    if title:
        head.append(f"Title: {_one_line(title, 200)}")
    if url:
        head.append(f"URL: {_one_line(url, 300)}")
    if parts > 1:
        head.append(f"Part {part} of {parts} of a longer document")
    line = " | ".join(head)
    return f"{line}\n\n{chunk}" if line else chunk


# --------------------------------------------------------------------------------------------- chunking

# Cost of splitting at each boundary level; the distance from the ideal point (as a fraction of the target
# chunk size) is added to it.
_LEVEL_COST = (0.0, 0.12, 0.3, 0.5, 0.8)
_HARD_CUT_COST = 2.0
_HEADING = re.compile(
    r"\n(?=#{1,6} |(?:chapter|CHAPTER|Chapter|PART|Part|BOOK|Book|SECTION|Section|ACT|Act"
    r"|Kapitel|KAPITEL|Chapitre|CHAPITRE|Cap[ií]tulo|CAP[IÍ]TULO|Capitolo|CAPITOLO|Глава|ГЛАВА|Rozdział|ROZDZIAŁ)\b[^\n]{0,80}\n"
    r"|第[一二三四五六七八九十百千〇零0-9]+[章节節回卷部篇][^\n]{0,80}\n"
    r"|[=\-*_]{3,}[ \t]*\n|\x0c|\\(?:chapter|section|subsection)\b)"
)
_CODE_DEF = re.compile(
    r"\n(?=(?:def |async def |class |function |func |fn |pub |impl |struct |enum |interface |trait |type |"
    r"module |package |public |private |protected |internal |static |export |const |let |var |@|#include|# ?%%))"
)
_PARAGRAPH = re.compile(r"\n[ \t]*\n+")
_LINE = re.compile(r"\n")
_SENTENCE = re.compile(r"[.!?…।॥۔؟։።။។៕][\"'”’)\]»]*\s+|[。！？][」』”’）\]]*")
_SPACE = re.compile(r"\s+")


def _joins_previous(ch: str) -> bool:
    """Characters that belong to the one before them: combining marks, ZWJ, variation selectors, skin tones."""
    o = ord(ch)
    return (unicodedata.category(ch) in ("Mn", "Mc", "Me") or o == 0x200D or 0xFE00 <= o <= 0xFE0F
            or 0xE0100 <= o <= 0xE01EF or 0x1F3FB <= o <= 0x1F3FF)


def split_text(text: str, offsets: list[int], max_tokens: int = CHUNK_TOKENS,
               is_code: bool = False) -> list[tuple[int, int, int]]:
    """Balanced chunks of ``text``: ``(start_char, end_char, tokens)`` each, ``offsets`` being the start character
    of every token. A 9k-token document becomes two ~4.5k chunks, not 7.8k plus 1.2k; each chunk ends at the
    cheapest boundary near its ideal end (headings or code definitions, paragraphs, lines, sentences, spaces), with
    a hard cut only as a last resort, never inside a character's combining marks or an emoji sequence."""
    n = len(offsets)
    if not text:
        return []
    if n <= max_tokens:
        return [(0, len(text), n)]
    levels = [_CODE_DEF if is_code else _HEADING, _PARAGRAPH, _LINE, None if is_code else _SENTENCE, _SPACE]
    spans: list[tuple[int, int, int]] = []
    s_tok, s_char = 0, 0
    while n - s_tok > max_tokens:
        remaining = n - s_tok
        k = math.ceil(remaining / max_tokens)
        target = remaining / k
        ideal = s_tok + target
        hard = s_tok + max_tokens  # the chunk ends at or before token `hard`
        lo_tok = max(s_tok + max(1, int(target * 0.5)), n - (k - 1) * max_tokens)
        lo_char, hi_char = int(offsets[lo_tok]), int(offsets[hard])
        best: tuple[float, int, int] | None = None
        for level, pattern in enumerate(levels):
            if best is not None and _LEVEL_COST[level] >= best[0]:
                break  # nothing at this level or later can beat the current best
            if pattern is None:
                continue
            for m in pattern.finditer(text, max(lo_char - 1, 0), min(len(text), hi_char + 256)):
                pos = m.end()
                if not lo_char <= pos <= hi_char:
                    continue
                tok = bisect_left(offsets, pos)
                if not s_tok < tok <= hard:
                    continue
                cost = _LEVEL_COST[level] + abs(tok - ideal) / target
                if best is None or cost < best[0]:
                    best = (cost, pos, tok)
        if best is None or best[0] >= _HARD_CUT_COST:
            c, t = hi_char, hard
            while c > s_char + 1 and c < len(text) and (_joins_previous(text[c]) or text[c - 1] == "\u200d"):
                c -= 1
            if c != hi_char:
                t2 = bisect_left(offsets, c)
                if t2 > s_tok:
                    t = t2
                else:
                    c = hi_char
            best = (_HARD_CUT_COST, c, t)
        _, split_char, split_tok = best
        spans.append((s_char, split_char, split_tok - s_tok))
        s_tok, s_char = split_tok, split_char
    spans.append((s_char, len(text), n - s_tok))
    return spans


def select_chunks(n: int, max_chunks: int) -> list[int]:
    """Indices of at most ``max_chunks`` evenly spaced chunks out of ``n`` (0 keeps all)."""
    if max_chunks <= 0 or n <= max_chunks:
        return list(range(n))
    if max_chunks == 1:
        return [n // 2]
    last = n - 1
    return sorted({round(i * last / (max_chunks - 1)) for i in range(max_chunks)})


# --------------------------------------------------------------------------------------------- scores

_CMP = {ast.Eq: operator.eq, ast.NotEq: operator.ne, ast.Lt: operator.lt, ast.LtE: operator.le,
        ast.Gt: operator.gt, ast.GtE: operator.ge}


class DropLineError(ValueError):
    """A drop line that is not valid (see ``DropLine``)."""


class DropLine:
    """A drop line such as ``toxicity >= 4 or spam_seo >= 3.5 or boilerplate >= 4.5``: comparisons of a score (or
    ``overall``, ``tokens``, ``parts``) with a number or another score, and of a label with one of its values
    (``format == 'news'``; ``==`` and ``!=`` only), joined with ``and`` / ``or`` / ``not`` and parentheses; ``True``
    and ``False`` match always and never. A comparison with a missing score (None) is false, ``!=`` included. A
    chunk or document is kept when the line does not match.

    ``names`` are the fields the line may use and ``labels`` maps each label field to its values; the other names
    are numbers. The line is checked when it is created: a bare field name or constant used as a condition, a string
    that is not a value of its label, a score compared with a string or a label with a number raise DropLineError
    (a ValueError)."""

    _NODES = (ast.Expression, ast.BoolOp, ast.And, ast.Or, ast.UnaryOp, ast.Not, ast.USub, ast.Compare, ast.Name,
              ast.Load, ast.Constant, *_CMP)

    def __init__(self, source: str, names: Iterable[str], labels: Mapping[str, Iterable[str]] | None = None):
        self.source = source.strip()
        self.names = set(names)
        self.labels = {name: list(values) for name, values in (labels or {}).items()}
        try:
            tree = ast.parse(self.source, mode="eval")
        except (SyntaxError, ValueError, RecursionError, MemoryError) as e:
            raise DropLineError(f"drop line {source!r} is not a valid expression ({e})") from None
        for node in ast.walk(tree):
            if not isinstance(node, self._NODES):
                raise DropLineError(f"{type(node).__name__} is not allowed in a drop line: {source!r}")
            if isinstance(node, ast.Name) and node.id not in self.names:
                raise DropLineError(f"unknown name {node.id!r} in drop line {source!r}")
        self.tree = tree.body
        try:
            self._condition(self.tree)
        except RecursionError:
            raise DropLineError(f"drop line {source!r} is nested too deeply") from None
        top_or = isinstance(self.tree, ast.BoolOp) and isinstance(self.tree.op, ast.Or)
        self.terms = list(self.tree.values) if top_or else [self.tree]
        # Evaluated once on stand-in scores, and once with every score missing, so that a line that cannot be
        # evaluated fails here and not in the middle of a run.
        try:
            self.reasons({n: self.labels[n][0] if self.labels.get(n) else 0.0 for n in self.names})
            self.reasons({})
        except RecursionError:
            raise DropLineError(f"drop line {source!r} is nested too deeply") from None

    def _fail(self, node: ast.AST, problem: str) -> None:
        raise DropLineError(f"{ast.unparse(node)!r} {problem}, in drop line {self.source!r}")

    def _condition(self, node: ast.AST) -> None:
        """Check that ``node`` is a condition: a comparison, True / False, or and / or / not of conditions."""
        if isinstance(node, ast.BoolOp):
            for value in node.values:
                self._condition(value)
        elif isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Not):
            self._condition(node.operand)
        elif isinstance(node, ast.Compare):
            kinds = [self._operand(x) for x in (node.left, *node.comparators)]
            for op, (a, b) in zip(node.ops, zip(kinds, kinds[1:])):
                self._check_pair(node, op, a, b)
        elif not (isinstance(node, ast.Constant) and isinstance(node.value, bool)):
            self._fail(node, "is not a condition: compare it with something, as in 'toxicity >= 4'")

    def _operand(self, node: ast.AST) -> tuple[str, Any]:
        """What one side of a comparison is: ("label", name), ("str", value), ("num", name) for a numeric field or
        ("num", None) for a number."""
        if isinstance(node, ast.Name):
            return ("label", node.id) if node.id in self.labels else ("num", node.id)
        if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
            kind = self._operand(node.operand)
            if kind[0] != "num":
                self._fail(node, "negates something that is not a number")
            return kind
        if isinstance(node, ast.Constant):
            v = node.value
            if isinstance(v, str):
                return ("str", v)
            if (isinstance(v, int) and not isinstance(v, bool)) or (isinstance(v, float) and math.isfinite(v)):
                return ("num", None)
            self._fail(node, "is not allowed in a comparison: use a finite number or a label value in quotes")
        self._fail(node, "is not allowed in a comparison: compare fields, numbers and label values")
        raise AssertionError  # not reached

    def _check_pair(self, node: ast.Compare, op: ast.cmpop, a: tuple[str, Any], b: tuple[str, Any]) -> None:
        kinds = {a[0], b[0]}
        if all(k == "str" or name is None for k, name in (a, b)):
            self._fail(node, "compares two constants")
        if kinds == {"num"}:
            return
        if kinds == {"label", "str"}:
            (_, label), (_, value) = (a, b) if a[0] == "label" else (b, a)
            if not isinstance(op, (ast.Eq, ast.NotEq)):
                self._fail(node, f"compares the label {label} by order: use == or !=")
            if value not in self.labels[label]:
                self._fail(node, f"uses {value!r}, which is not a value of {label} (one of: "
                                 f"{', '.join(self.labels[label])})")
            return
        if kinds == {"label"}:
            self._fail(node, "compares two labels: compare a label with one of its values, as in format == 'news'")
        self._fail(node, "compares a score with a string, or a label with a number")

    def reasons(self, env: dict) -> list[str]:
        """The terms of the line that match ``env`` (each ``or`` branch on its own); [] means keep."""
        return [ast.unparse(t) for t in self.terms if self._eval(t, env)]

    def _eval(self, node: ast.AST, env: dict) -> Any:
        if isinstance(node, ast.Constant):
            return node.value
        if isinstance(node, ast.Name):
            return env.get(node.id)
        if isinstance(node, ast.BoolOp):
            value: Any = isinstance(node.op, ast.And)
            for v in node.values:
                value = self._eval(v, env)
                if bool(value) != isinstance(node.op, ast.And):
                    return value
            return value
        if isinstance(node, ast.UnaryOp):
            v = self._eval(node.operand, env)
            if isinstance(node.op, ast.Not):
                return not v
            return None if v is None else -v
        if isinstance(node, ast.Compare):
            left = self._eval(node.left, env)
            for op, comp in zip(node.ops, node.comparators):
                right = self._eval(comp, env)
                if left is None or right is None or not _CMP[type(op)](left, right):
                    return False  # a comparison with a missing value is false, != included
                left = right
            return True
        raise DropLineError(f"unsupported drop-line node {type(node).__name__}")


def _r3(x: float | None) -> float | None:
    return None if x is None or (isinstance(x, float) and math.isnan(x)) else round(float(x), 3)


def composite(schema: dict, scores: dict) -> float | None:
    """The overall score (0-5) of a flat score dict, by the schema in source1.json: the weighted mean of the quality
    scores; when gated scores apply, 80% of that plus 20% of their mean; minus, for each red flag, its penalty
    weight times how far it is above its threshold (spam_seo 0.5, boilerplate 0.4, toxicity 0.8, each above 1);
    clipped to 0-5."""
    q_num = q_den = 0.0
    for name, spec in schema["quality"].items():
        v, w = scores.get(name), float(spec.get("weight", 1.0))
        if v is not None and w > 0:
            q_num += w * float(v)
            q_den += w
    if q_den <= 0:
        return None
    quality = q_num / q_den
    gated = [float(scores[n]) for n in schema.get("gated", {}) if scores.get(n) is not None]
    gw = float((schema.get("composite") or {}).get("gated_weight", 0.2))
    if gated and gw > 0:
        quality = (1 - gw) * quality + gw * (sum(gated) / len(gated))
    penalty = 0.0
    for name, spec in (schema.get("red_flags") or {}).items():
        v, pen = scores.get(name), spec.get("penalty") or {}
        if v is not None and float(pen.get("weight", 0.0)) > 0:
            penalty += float(pen["weight"]) * max(0.0, float(v) - float(pen.get("threshold", 0.0)))
    return round(min(5.0, max(0.0, quality - penalty)), 3)


def gate_applies(schema: dict, name: str, labels: dict) -> bool:
    """Whether gated score ``name`` applies, given the labels: any of its ``applies_when`` labels has one of the
    listed values (code_quality: code content or a code file; math_quality: math-heavy content or the math topic)."""
    for label, values in schema["gated"][name]["applies_when"].items():
        values = [values] if isinstance(values, str) else values
        if labels.get(label) in values:
            return True
    return False


def aggregate(schema: dict, items: list[tuple[int, dict]]) -> dict:
    """Document scores from ``(tokens, chunk_scores)`` pairs: labels by token-weighted vote, scores by token-weighted
    mean (or max / min where the schema says so: toxicity uses max), gated scores over the chunks where they apply.
    With several chunks also ``label_dist`` (token share of each label value) and ``ranges`` ([min, max] per score)."""
    out: dict = {}
    dist_out: dict = {}
    ranges: dict = {}
    total_w = sum(max(w, 1) for w, _ in items)
    for name in schema["labels"]:
        weights: dict[str, float] = {}
        for w, sc in items:
            v = sc.get(name)
            if v is not None:
                weights[v] = weights.get(v, 0.0) + max(w, 1)
        if not weights:
            out[name] = None
            continue
        dist = {k: round(v / total_w, 3) for k, v in sorted(weights.items(), key=lambda kv: -kv[1])}
        out[name] = next(iter(dist))
        if len(items) > 1:
            dist_out[name] = dist
    for group in ("quality", "red_flags", "gated"):
        for name, spec in (schema.get(group) or {}).items():
            vals = [(float(v), max(w, 1)) for w, sc in items if (v := sc.get(name)) is not None]
            if not vals:
                out[name] = None
                continue
            how = spec.get("aggregate", "mean")
            if how == "max":
                agg = max(v for v, _ in vals)
            elif how == "min":
                agg = min(v for v, _ in vals)
            else:
                agg = sum(v * w for v, w in vals) / sum(w for _, w in vals)
            out[name] = _r3(agg)
            if len(vals) > 1:
                ranges[name] = [_r3(min(v for v, _ in vals)), _r3(max(v for v, _ in vals))]
    if dist_out:
        out["label_dist"] = dist_out
    if ranges:
        out["ranges"] = ranges
    return out


# --------------------------------------------------------------------------------------------- model

_DTYPES = {"bfloat16": torch.bfloat16, "bf16": torch.bfloat16, "float32": torch.float32, "fp32": torch.float32}
DTYPE_CHOICES = ("auto", "bf16", "bfloat16", "fp32", "float32")
PRECISION_CHOICES = ("bf16", "fp32")
_FP16 = ("float16 is not supported: summing the hidden states for mean pooling overflows float16's range on long "
         "inputs and gives NaN scores. Use dtype='bf16' (GPUs from NVIDIA Ampere on) or dtype='fp32' (any device).")
_FP16_WEIGHTS = ("there are no float16 weights, and float16 is not supported (summing the hidden states for mean "
                 "pooling overflows its range on long inputs and gives NaN scores). Use precision='bf16' (the "
                 "default, model.safetensors) or precision='fp32' (model.fp32.safetensors).")


def resolve_precision(precision: Any = None) -> str:
    """Which weights to load, "bf16" (model.safetensors, the default) or "fp32" (model.fp32.safetensors): None,
    "bf16" / "bfloat16" / torch.bfloat16, or "fp32" / "float32" / torch.float32. float16 raises ValueError."""
    if precision is None:
        return DEFAULT_PRECISION
    if isinstance(precision, str):
        key = precision.lower().removeprefix("torch.")
        if key in ("float16", "fp16", "half"):
            raise ValueError(_FP16_WEIGHTS)
        if key in ("bf16", "bfloat16"):
            return "bf16"
        if key in ("fp32", "float32"):
            return "fp32"
        raise ValueError(f"unknown precision {precision!r}: use 'bf16' (the default) or 'fp32'")
    if precision == torch.float16:
        raise ValueError(_FP16_WEIGHTS)
    if precision == torch.bfloat16:
        return "bf16"
    if precision == torch.float32:
        return "fp32"
    raise ValueError(f"unknown precision {precision!r}: use 'bf16' (the default) or 'fp32'")


def stored_dtype(path: str | Path) -> str | None:
    """The dtype a safetensors file stores its tensors in ("BF16", "F32", ...; "mixed" when several), read from its
    header; None when the header cannot be read."""
    try:
        with open(path, "rb") as f:
            n = int.from_bytes(f.read(8), "little")
            if not 0 < n <= 100_000_000:
                return None
            header = json.loads(f.read(n))
    except (OSError, ValueError):
        return None
    kinds = {v.get("dtype") for k, v in header.items() if k != "__metadata__" and isinstance(v, dict)}
    return None if not kinds else kinds.pop() if len(kinds) == 1 else "mixed"


def _native_bf16(device: torch.device) -> bool:
    """Whether ``device`` is a GPU that runs bfloat16 natively (not emulated)."""
    if device.type != "cuda":
        return False
    try:
        with torch.cuda.device(device):
            return bool(torch.cuda.is_bf16_supported(including_emulation=False))
    except TypeError:  # a torch without including_emulation
        return torch.cuda.get_device_capability(device)[0] >= 8


def resolve_dtype(dtype: Any, device: str | torch.device) -> torch.dtype:
    """The dtype to run in: None or "auto" = bfloat16 on a GPU with native bfloat16, float32 everywhere else;
    "bf16" / "bfloat16" / torch.bfloat16 or "fp32" / "float32" / torch.float32 as given. float16 raises ValueError."""
    device = torch.device(device)
    if dtype is None or (isinstance(dtype, str) and dtype.lower() == "auto"):
        return torch.bfloat16 if _native_bf16(device) else torch.float32
    if isinstance(dtype, str):
        key = dtype.lower().removeprefix("torch.")
        if key in ("float16", "fp16", "half"):
            raise ValueError(_FP16)
        if key not in _DTYPES:
            raise ValueError(f"unknown dtype {dtype!r}: use 'auto', 'bf16' or 'fp32'")
        return _DTYPES[key]
    if dtype == torch.float16:
        raise ValueError(_FP16)
    if dtype not in (torch.bfloat16, torch.float32):
        raise ValueError(f"unsupported dtype {dtype}: use torch.bfloat16 or torch.float32")
    return dtype


def _dtype_kwarg() -> str:
    """transformers 4.56 renamed from_pretrained(torch_dtype=...) to dtype=..."""
    import transformers

    major, minor = (int("".join(c for c in p if c.isdigit()) or 0) for p in transformers.__version__.split(".")[:2])
    return "dtype" if (major, minor) >= (4, 56) else "torch_dtype"


_REPO_ID = re.compile(r"[A-Za-z0-9][\w.-]*/[\w.-]+", re.ASCII)


def _resolve_dir(path_or_repo: str | Path, revision: str | None = None,
                 precision: str = DEFAULT_PRECISION) -> Path:
    """A local directory as is; a Hugging Face repo id ("owner/name") downloaded with huggingface_hub (only the files
    source1.py needs, with the weights of ``precision`` only, at ``revision`` when given). Anything else raises
    FileNotFoundError, never a download."""
    s = str(path_or_repo)
    path = Path(s).expanduser()
    if path.is_dir():
        if revision is not None:
            raise ValueError("revision= applies to a Hugging Face repo id, not to a local directory")
        return path
    if path.exists():
        raise FileNotFoundError(f"{s} is a file; pass the Source-1 directory that holds it")
    is_repo_id = (isinstance(path_or_repo, str) and _REPO_ID.fullmatch(s) is not None
                  and not s.startswith((".", "~", "/")) and not Path(s.split("/")[0]).exists())
    if not is_repo_id:
        raise FileNotFoundError(f"{s}: no such Source-1 directory")
    from huggingface_hub import snapshot_download  # installed with transformers

    return Path(snapshot_download(repo_id=s, revision=revision, allow_patterns=list(hub_files(precision))))


def _load_calibration(path: Path, drop_line: str | None, apply_offsets: bool) -> dict | None:
    """calibration.json of a Source-1 directory. It is required for the calibrated drop line (the default) and for
    ``apply_offsets``; a missing or incomplete file then raises FileNotFoundError instead of silently falling back."""
    cal_path = path / CALIBRATION
    calibration = json.loads(cal_path.read_text(encoding="utf-8")) if cal_path.exists() else None
    has_line = bool(((calibration or {}).get("drop_line") or {}).get("line"))
    if drop_line in (None, "calibrated") and not has_line:
        what = "has no drop line" if calibration is not None else "is missing"
        raise FileNotFoundError(
            f"{cal_path} {what}, and the default drop line comes from it: download it again, or pass "
            "drop_line='default' (the rubric's own hard filters) or your own drop line to load without it")
    if apply_offsets and not (calibration or {}).get("offsets"):
        raise FileNotFoundError(f"{cal_path} is missing or has no offsets, which apply_offsets=True needs")
    return calibration


def _make_drop_line(config: dict, calibration: dict | None, drop_line: str | None = None) -> DropLine:
    """The ``DropLine`` of ``drop_line`` for the schema of source1.json (``config``): None or "calibrated" =
    calibration.json's line, "default" = the schema's hard filters, anything else = the expression itself, checked
    against the schema's fields and label values (ValueError when it is not a valid drop line)."""
    schema = config["schema"]
    if drop_line in (None, "calibrated"):
        drop_line = ((calibration or {}).get("drop_line") or {}).get("line")
        if not drop_line:
            raise ValueError("no calibrated drop line (calibration.json missing or incomplete); pass "
                             "drop_line='default' or your own drop line")
    elif drop_line == "default":
        drop_line = " or ".join((schema.get("composite") or {}).get("hard_filters") or []) or "False"
    labels = {name: list(spec["values"]) for name, spec in schema["labels"].items()}
    numbers = [n for g in ("quality", "red_flags", "gated") for n in schema.get(g, {})] + ["overall", "tokens", "parts"]
    return DropLine(drop_line, [*labels, *numbers], labels)


class Source1(nn.Module):
    """Source-1: the mmBERT-base encoder, mean pooling and one linear head per field. Build it with
    ``Source1.from_pretrained``; score documents with ``score`` / ``score_batch``."""

    def __init__(self, backbone: nn.Module, tokenizer: Any, config: dict, calibration: dict | None = None,
                 drop_line: str | None = None, apply_offsets: bool = False, show_url: bool = False):
        super().__init__()
        self.backbone = backbone
        self.tok = tokenizer
        self.config = config
        self.schema = config["schema"]
        self.layout = [tuple(x) for x in config["layout"]]
        if any(kind not in ("label_single", "score") for _, kind, _ in self.layout):
            raise ValueError("this scorer handles single-choice labels and 0-5 scores only")
        if config.get("pooling", "mean") != "mean" or config.get("text_normalize") not in (None, "collapse_spaces"):
            raise ValueError("unexpected source1.json: this scorer expects mean pooling and collapse_spaces")
        self.normalize = collapse_spaces if config.get("text_normalize") else (lambda s: s)
        hidden = int(backbone.config.hidden_size)
        self.heads = nn.ModuleDict({name: nn.Linear(hidden, size) for name, _, size in self.layout})
        self.max_length = int(config.get("max_length") or MAX_LENGTH)
        self.labels = {name: list(spec["values"]) for name, spec in self.schema["labels"].items()}
        self.fields = list(self.labels) + [n for g in ("quality", "red_flags", "gated") for n in self.schema.get(g, {})]
        self.calibration = calibration or {}
        self.drop_line = _make_drop_line(config, self.calibration, drop_line)
        self.offsets = {k: float(v["offset"]) for k, v in (self.calibration.get("offsets") or {}).items()}
        self.apply_offsets = apply_offsets
        self.show_url = show_url
        self.precision: str | None = None  # set by from_pretrained: "bf16" or "fp32", the weights it loaded
        self.weights_file: str | None = None
        self.weights_dtype: str | None = None  # how that file stores its tensors ("BF16", "F32")
        ids = tokenizer.encode("", add_special_tokens=True).ids
        if len(ids) != 2:
            raise ValueError("expected the tokenizer to add exactly <bos> and <eos>")
        self.bos_id, self.eos_id = ids
        self.pad_id = int(tokenizer.token_to_id("<pad>") if tokenizer.token_to_id("<pad>") is not None else 0)

    # ----------------------------------------------------------------------------------------- loading

    @classmethod
    def from_pretrained(cls, path_or_dir: str | Path, device: str | None = None, dtype: Any = None, *,
                        precision: Any = DEFAULT_PRECISION, drop_line: str | None = None,
                        apply_offsets: bool = False, show_url: bool = False, revision: str | None = None,
                        **backbone_kwargs: Any) -> "Source1":
        """Load Source-1 from a directory, or from a Hugging Face repo id ("owner/Source-1", downloaded with
        huggingface_hub; ``revision`` pins a branch, tag or commit).

        device: "cuda", "cuda:1", "cpu", ...; default the GPU when there is one, else the CPU.
        precision: which weights to load: "bf16" (default) = model.safetensors, bfloat16; "fp32" =
            model.fp32.safetensors, the full-precision float32 copy (a Hub repo id downloads only the chosen file).
            float16 raises ValueError.
        dtype: what the model computes in. None or "auto" (default) = bfloat16 on a GPU with native bfloat16
            (Ampere and newer), float32 elsewhere (bfloat16 weights are then upcast to float32); or "bf16" / "fp32" /
            torch.bfloat16 / torch.float32. float16 raises ValueError (it overflows).
        drop_line: None or "calibrated" = calibration.json's line (the file is then required); "default" = the
            schema's own hard filters (``toxicity >= 4 or spam_seo >= 4 or boilerplate >= 4.5``); or any expression
            ``DropLine`` accepts.
        apply_offsets: add calibration.json's offsets to the quality scores (tiny; off by default).
        show_url: show ``url`` to the model in the header (it was never trained with one; off by default).
        backbone_kwargs: passed to transformers' AutoModel.from_pretrained (e.g. attn_implementation="sdpa")."""
        from safetensors.torch import load_file
        from tokenizers import Tokenizer
        from transformers import AutoModel

        precision = resolve_precision(precision)
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        dtype = resolve_dtype(dtype, device)
        if "variant" in backbone_kwargs:
            raise TypeError("pass precision='bf16' or precision='fp32' instead of variant=")
        path = _resolve_dir(path_or_dir, revision, precision)
        weights = WEIGHTS[precision]
        if not (path / weights).exists():
            other = "bf16" if precision == "fp32" else "fp32"
            have_other = (path / WEIGHTS[other]).exists()
            raise FileNotFoundError(
                f"{path} lacks {weights}, the {'float32' if precision == 'fp32' else 'bfloat16'} weights that "
                f"precision={precision!r} loads" + (f"; precision={other!r} loads {WEIGHTS[other]}, which is there"
                                                    if have_other else ""))
        missing = [f for f in FILES if not (path / f).exists()]
        if missing:
            raise FileNotFoundError(f"{path} lacks {', '.join(missing)}")
        config = json.loads((path / "source1.json").read_text(encoding="utf-8"))
        calibration = _load_calibration(path, drop_line, apply_offsets)
        _make_drop_line(config, calibration, drop_line)  # a bad drop line fails here, before the weights load
        tok = Tokenizer.from_file(str(path / "tokenizer.json"))
        tok.no_truncation()
        tok.no_padding()
        variant = {"variant": "fp32"} if precision == "fp32" else {}
        backbone_kwargs.pop("output_loading_info", None)  # always requested, to check the weights
        backbone, info = AutoModel.from_pretrained(str(path), **{_dtype_kwarg(): dtype}, **variant, **backbone_kwargs,
                                                   output_loading_info=True)
        keys = ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")
        bad = {k: info[k] for k in keys if info.get(k)}
        if bad:  # transformers would otherwise fill the missing weights with random values, without an error
            detail = "; ".join(f"{len(v)} {k.replace('_', ' ')} ({', '.join(sorted(map(str, v))[:3])}"
                               f"{', ...' if len(v) > 3 else ''})" for k, v in bad.items())
            raise RuntimeError(f"Source-1: the backbone weights in {path / weights} do not match the model: {detail}. "
                               "Download the weights again, and use the package versions in requirements.txt")
        model = cls(backbone, tok, config, calibration, drop_line, apply_offsets, show_url)
        model.heads.load_state_dict(load_file(str(path / "heads.safetensors")))
        model.heads.to(dtype)
        model.precision = precision
        model.weights_file = weights
        model.weights_dtype = stored_dtype(path / weights)
        return model.to(device).eval()

    @property
    def device(self) -> torch.device:
        return next(self.backbone.parameters()).device

    # ----------------------------------------------------------------------------------------- inputs

    def chunk(self, text: str, is_code: bool = False, max_tokens: int = CHUNK_TOKENS) -> list[tuple[int, int, int]]:
        """``split_text`` with this model's tokenizer: (start_char, end_char, tokens) per chunk of a cleaned text."""
        if not text:
            return []
        offsets = [s for s, _ in self.tok.encode(text, add_special_tokens=False).offsets]
        return split_text(text, offsets, max_tokens, is_code)

    def encode(self, texts: list[str]) -> list[tuple[list[int], bool]]:
        """Token ids of model inputs (``collapse_spaces`` applied, <bos> ... <eos>, at most max_length tokens: a
        longer input keeps its first max_length - 1 tokens and its <eos>), each with whether it had to be cut."""
        out = []
        for enc in self.tok.encode_batch([self.normalize(t) for t in texts], add_special_tokens=True):
            ids = enc.ids
            out.append((ids, False) if len(ids) <= self.max_length else (ids[: self.max_length - 1] + [self.eos_id], True))
        return out

    # ----------------------------------------------------------------------------------------- scoring

    @torch.inference_mode()
    def _logits(self, batch: list[list[int]]) -> dict[str, torch.Tensor]:
        width = max(len(ids) for ids in batch)
        x = torch.full((len(batch), width), self.pad_id, dtype=torch.long)
        m = torch.zeros((len(batch), width), dtype=torch.long)
        for row, ids in enumerate(batch):
            x[row, : len(ids)] = torch.tensor(ids, dtype=torch.long)
            m[row, : len(ids)] = 1
        x, m = x.to(self.device), m.to(self.device)
        out = self.backbone(input_ids=x, attention_mask=m)
        hidden = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
        mask = m.unsqueeze(-1).to(hidden.dtype)
        pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)  # mean over the real tokens
        pooled = pooled.to(next(iter(self.heads.values())).weight.dtype)
        return {name: head(pooled) for name, head in self.heads.items()}

    @torch.inference_mode()
    def _decode(self, logits: dict[str, torch.Tensor]) -> list[dict]:
        n = next(iter(logits.values())).shape[0]
        for name, x in logits.items():  # NaN would silently keep a document and break the JSON output
            finite = torch.isfinite(x).all(dim=-1)
            if not bool(finite.all()):
                raise FloatingPointError(
                    f"head {name!r} gave NaN or infinite outputs for {int((~finite).sum())} of {n} inputs; "
                    "this should not happen in bfloat16 or float32")
        rows: list[dict] = [{} for _ in range(n)]
        levels = torch.tensor(LEVELS, dtype=torch.float32, device=next(iter(logits.values())).device)
        for name, kind, _ in self.layout:
            x = logits[name].float()
            if kind == "label_single":
                for row, i in enumerate(x.argmax(dim=-1).tolist()):
                    rows[row][name] = self.labels[name][i]
            else:
                expected = (torch.softmax(x, dim=-1) * levels).sum(dim=-1).tolist()
                for row, e in enumerate(expected):
                    rows[row][name] = round(e, 3)
        return rows

    def _finish(self, scores: dict) -> dict:
        """Gating, optional offsets, overall and keep for one chunk's raw head outputs (a flat dict)."""
        labels = {k: scores[k] for k in self.labels}
        out = {**labels}
        for group in ("quality", "red_flags", "gated"):
            for name in self.schema.get(group, {}):
                v = scores[name]
                if group == "gated" and not gate_applies(self.schema, name, labels):
                    v = None
                elif group == "quality" and self.apply_offsets and name in self.offsets:
                    v = round(min(5.0, max(0.0, v + self.offsets[name])), 3)
                out[name] = v
        out["overall"] = composite(self.schema, out)
        return out

    def _judge(self, rec: dict) -> dict:
        reasons = self.drop_line.reasons(rec)
        rec["keep"] = not reasons
        rec["drop_reasons"] = reasons
        return rec

    def default_batch_tokens(self) -> int:
        """The padded tokens per forward pass when ``batch_tokens`` is None: 16,384 on a CPU, 65,536 on a GPU."""
        return DEFAULT_BATCH_TOKENS_CPU if self.device.type == "cpu" else DEFAULT_BATCH_TOKENS

    def score_inputs(self, inputs: list[str], batch_tokens: int | None = None) -> list[dict]:
        """Score ready-made model inputs (``build_input`` output: header, blank line, chunk text), one dict per
        input with the 13 fields, overall, keep, drop_reasons, input_tokens and truncated.

        Inputs are sorted by length and batched with at most ``batch_tokens`` padded tokens per forward pass
        (default: 16,384 on a CPU, 65,536 on a GPU); a batch that runs out of GPU memory is split in half and
        retried."""
        if batch_tokens is None:
            batch_tokens = self.default_batch_tokens()
        encoded = self.encode(inputs)
        ids = [e[0] for e in encoded]
        results: list[dict | None] = [None] * len(ids)

        def run(batch: list[int]) -> None:
            try:
                rows = self._decode(self._logits([ids[i] for i in batch]))
            except torch.cuda.OutOfMemoryError:
                torch.cuda.empty_cache()
                if len(batch) == 1:
                    raise
                run(batch[: len(batch) // 2])
                run(batch[len(batch) // 2:])
                return
            for i, raw in zip(batch, rows):
                rec = self._judge(self._finish(raw))
                rec["input_tokens"] = len(ids[i])
                rec["truncated"] = encoded[i][1]
                results[i] = rec

        order = sorted(range(len(ids)), key=lambda i: -len(ids[i]))
        start = 0
        while start < len(order):
            width = len(ids[order[start]])
            batch = order[start: start + max(1, batch_tokens // max(width, 1))]
            start += len(batch)
            run(batch)
        return results  # type: ignore[return-value]

    def score_batch(self, docs: Iterable[str | bytes | dict], batch_tokens: int | None = None, *,
                    text_field: str = "text", max_chunks: int = 0) -> list[dict]:
        """Score a list of documents: strings (bytes are decoded as UTF-8), or dicts with the text under
        ``text_field`` and optionally ``title``, ``url``, ``source_type`` and ``code_language`` (see ``score``).
        A None text is scored as an empty document (keep False, drop_reasons ["empty text"]). Chunks of all
        documents are batched together. ``max_chunks`` > 0 scores only that many evenly spaced chunks of a long
        document (the Part numbers still count every chunk). ``batch_tokens``: see ``score_inputs``.

        In bfloat16 a document's scores can shift slightly (up to about 0.04 on overall, 0.10 on a single field)
        depending on which other documents share its batch, because the batch shape changes the kernels' rounding;
        use dtype="fp32" or batch_tokens=1 when a document's scores must not depend on its neighbours."""
        if isinstance(docs, (str, bytes, bytearray, Mapping)):
            raise TypeError("score_batch takes a list of documents; use score() for one document")
        plans: list[dict] = []
        inputs: list[str] = []
        for n, doc in enumerate(docs):
            if isinstance(doc, Mapping):
                if text_field not in doc:
                    raise KeyError(f"document {n} has no {text_field!r} field")
                d = {k: _from_bytes(v) for k, v in doc.items()}
            else:
                d = {text_field: _from_bytes(doc)}
            raw = d[text_field]
            if raw is not None and not isinstance(raw, str):
                raise TypeError(f"document {n}: the text must be a str, bytes or None, not {type(raw).__name__}")
            text = clean_text(raw or "")
            code_language = d.get("code_language")
            source = source_description(d.get("source_type"), code_language)
            spans = self.chunk(text, is_code=bool(code_language))
            chosen = select_chunks(len(spans), max_chunks)
            first = len(inputs)
            for i in chosen:
                s, e, _ = spans[i]
                inputs.append(build_input(text[s:e].strip(), source_type=source, title=d.get("title"),
                                          url=d.get("url") if self.show_url else None, part=i + 1,
                                          parts=len(spans)))
            plans.append({"spans": spans, "chosen": chosen, "first": first})
        scored = self.score_inputs(inputs, batch_tokens)
        return [self._document(p, scored[p["first"]: p["first"] + len(p["chosen"])]) for p in plans]

    def score(self, text: str | bytes, title: str | None = None, url: str | None = None, *,
              source_type: str | None = None, code_language: str | None = None, max_chunks: int = 0,
              batch_tokens: int | None = None) -> dict:
        """Score one document (a str; bytes are decoded as UTF-8; anything else raises TypeError).

        title: shown to the model in the header when given (as in training, where about a quarter of inputs had one).
        url: kept out of the model's input unless the model was loaded with show_url=True.
        source_type: the header's Source value; default "dataset record" ("" leaves the Source line out).
        code_language: for source code, e.g. "Python": Source becomes "Python source file" and long files are split
            at definitions."""
        doc = {"text": text, "title": title, "url": url, "source_type": source_type, "code_language": code_language}
        return self.score_batch([doc], batch_tokens, max_chunks=max_chunks)[0]

    def _document(self, plan: dict, chunks: list[dict]) -> dict:
        spans, chosen = plan["spans"], plan["chosen"]
        if not spans:
            out = {name: None for name in self.fields}
            out.update(overall=None, keep=False, drop_reasons=["empty text"], parts=0, tokens=0, truncated=False)
            return out
        items = [(spans[i][2], c) for i, c in zip(chosen, chunks)]
        agg = aggregate(self.schema, items)
        out = {name: agg[name] for name in self.fields}
        out["overall"] = composite(self.schema, out)
        parts, tokens = len(spans), sum(s[2] for s in spans)
        reasons = self.drop_line.reasons({**out, "parts": parts, "tokens": tokens})
        out.update(keep=not reasons, drop_reasons=reasons, parts=parts, tokens=tokens,
                   truncated=any(c["truncated"] for c in chunks))
        if len(spans) > 1:
            out["chunks"] = [{"part": i + 1, "start": spans[i][0], "end": spans[i][1], "tokens": spans[i][2], **c}
                             for i, c in zip(chosen, chunks)]
            for key in ("label_dist", "ranges"):
                if key in agg:
                    out[key] = agg[key]
        return out


# --------------------------------------------------------------------------------------------- CLI


class BadInput(ValueError):
    """A record of the input file that cannot be scored (the message starts with file:line), or, with
    ``skippable=False``, an input file that cannot be read at all."""

    def __init__(self, message: str, skippable: bool = True):
        super().__init__(message)
        self.skippable = skippable


# Compressed inputs read transparently, recognized by their first bytes whatever their name.
_DECOMPRESS = {"gzip": gzip.open, "bzip2": bz2.open, "xz": lzma.open}
_COMPRESSED_SUFFIXES = (".gz", ".bz2", ".xz")  # docs.jsonl.gz is read as .jsonl
_UNREADABLE = {".zst": "is zstd-compressed: decompress it first (zstd -d)",
               ".zstd": "is zstd-compressed: decompress it first (zstd -d)",
               ".parquet": "is a Parquet file: convert it to JSON Lines first",
               ".zip": "is a zip archive: extract it first",
               ".7z": "is a 7z archive: extract it first"}
_READ_ERRORS = (OSError, EOFError, zlib.error, lzma.LZMAError)  # what damaged or truncated compressed files raise
_SNIFF_BYTES = 65536  # how much of an input file _input_problem looks at
_MAX_INVALID = 0.2  # share of invalid UTF-8 sequences among the characters above which an input file is refused
_MAX_NUL = 0.01  # share of NUL bytes above which an input file is refused as binary


def _compression(head: bytes) -> str | None:
    """The compression that a file's first bytes show: "gzip", "bzip2", "xz", "zstd" or None."""
    if head.startswith(b"\x1f\x8b"):
        return "gzip"
    if head[:3] == b"BZh" and b"1" <= head[3:4] <= b"9" and head[4:10] in (b"1AY&SY", b"\x17rE8P\x90"):
        return "bzip2"
    if head.startswith(b"\xfd7zXZ\x00"):
        return "xz"
    if head.startswith(b"\x28\xb5\x2f\xfd"):
        return "zstd"
    return None


def _open_input(path: Path) -> Any:
    """``path`` opened for reading bytes, decompressed when it is a gzip, bzip2 or xz file."""
    with open(path, "rb") as f:
        head = f.read(10)
    return _DECOMPRESS.get(_compression(head) or "", open)(path, "rb")


def _format_suffix(path: Path) -> str:
    """The suffix that says how to read ``path``: its last one, or the one before .gz, .bz2 or .xz."""
    suffixes = [s.lower() for s in path.suffixes]
    if suffixes and suffixes[-1] in _COMPRESSED_SUFFIXES:
        suffixes.pop()
    return suffixes[-1] if suffixes else ""


def _input_problem(path: Path) -> str | None:
    """Why ``path`` cannot be scored as UTF-8 text or JSON (zstd, Parquet, an archive, binary, UTF-16, mostly invalid
    UTF-8, damaged), judged from its name and its first 64 KiB once decompressed; None when it looks readable."""
    if path.suffix.lower() in _UNREADABLE:
        return _UNREADABLE[path.suffix.lower()]
    try:
        with _open_input(path) as f:
            head = f.read(_SNIFF_BYTES)
    except _READ_ERRORS as e:
        return f"cannot be read ({e})"
    if _compression(head) == "zstd":
        return _UNREADABLE[".zst"]
    if head.startswith((b"\xff\xfe", b"\xfe\xff")):
        return "is UTF-16 or UTF-32: save it as UTF-8"
    if head.count(b"\x00") > _MAX_NUL * len(head):  # a stray NUL in a text is dropped like other control characters
        return "holds NUL bytes, so it is not UTF-8 text (binary, UTF-16, or compressed in a format not read here)"
    text = head.decode("utf-8", errors="replace")
    invalid = text.count("\ufffd") - head.count("\ufffd".encode())  # a U+FFFD already in the text is valid UTF-8
    if invalid >= 8 and invalid > _MAX_INVALID * len(text):
        return f"is not UTF-8 text: {invalid / len(text):.0%} of its first {len(text):,} characters are invalid UTF-8"
    return None


def _json_error(e: BaseException) -> str:
    """Why json.loads failed, in a few words."""
    if isinstance(e, json.JSONDecodeError):
        return f"{e.msg} at column {e.colno}"
    if isinstance(e, RecursionError):
        return "nested too deeply"
    return str(e).split(";")[0]  # e.g. an integer of more than 4,300 digits


def _unwritable(value: Any) -> str | None:
    """Why ``value`` cannot be written to the JSON Lines output, or None when it can."""
    try:
        json.dumps(value, ensure_ascii=False, allow_nan=False).encode("utf-8")
    except UnicodeEncodeError:
        return "holds a lone surrogate, which UTF-8 cannot encode"
    except ValueError:
        return "holds NaN or an infinite number, which JSON cannot hold"
    except RecursionError:
        return "is nested too deeply"
    return None


def _read_docs(path: Path, text_field: str, skip_bad: bool = False) -> Iterator[dict]:
    """Documents of an input file: .jsonl / .ndjson (one JSON object per line), .json (a JSON array of objects, one
    object, or JSON Lines), or any other file as one plain-text document; gzip, bzip2 and xz files are decompressed
    (docs.jsonl.gz is read as .jsonl). Invalid UTF-8 is replaced, with a warning. A bad record (not a JSON object,
    no string text, an id that cannot be written back) raises BadInput, or with ``skip_bad`` is reported on stderr
    and skipped. A file that cannot be read at all (see ``_input_problem``) raises BadInput with skippable=False.
    A null text is kept and scored as an empty document."""

    def usable(rec: Any, where: str) -> bool:
        if not isinstance(rec, dict):
            problem = f"expected a JSON object, got {type(rec).__name__}"
        elif text_field not in rec:
            problem = f"no {text_field!r} field"
        elif rec[text_field] is not None and not isinstance(rec[text_field], str):
            problem = f"{text_field!r} is a {type(rec[text_field]).__name__}, not a string"
        elif "id" in rec and (why := _unwritable(rec["id"])):
            problem = f"the id {why}"
        else:
            return True
        if not skip_bad:
            raise BadInput(f"{where}: {problem}")
        print(f"source1: skipped {where}: {problem}", file=sys.stderr)
        return False

    def decode(raw: bytes, where: str) -> str:
        try:
            return raw.decode("utf-8")
        except UnicodeDecodeError as e:
            print(f"source1: {where}: invalid UTF-8 at byte {e.start}; invalid bytes replaced with U+FFFD",
                  file=sys.stderr)
            return raw.decode("utf-8", errors="replace")

    def read_all() -> bytes:
        try:
            with _open_input(path) as f:
                return f.read()
        except _READ_ERRORS as e:
            raise BadInput(f"{path} cannot be read ({e})", skippable=False) from None

    problem = _input_problem(path)
    if problem:
        raise BadInput(f"{path} {problem}", skippable=False)
    suffix = _format_suffix(path)
    if suffix not in (".jsonl", ".ndjson", ".json"):
        yield {text_field: decode(read_all(), str(path)), "id": path.name}
        return
    if suffix == ".json":
        try:
            data = json.loads(decode(read_all(), str(path)).lstrip(""))
        except (ValueError, RecursionError):  # JSONDecodeError is a ValueError
            data = None  # not one JSON value (or too large or too deep to read as one): read it as JSON Lines below
        if data is not None:
            for n, rec in enumerate(data if isinstance(data, list) else [data]):
                if usable(rec, f"{path}[{n}]"):
                    yield rec
            return
    n = 0
    try:
        with _open_input(path) as f:
            for n, raw in enumerate(f, 1):
                where = f"{path}:{n}"
                line = decode(raw, where)
                if n == 1:
                    line = line.lstrip("")
                if not line.strip():
                    continue
                try:
                    rec = json.loads(line)
                except (ValueError, RecursionError) as e:  # also integers too long to read, and very deep nesting
                    problem = f"invalid JSON ({_json_error(e)})"
                    if not skip_bad:
                        raise BadInput(f"{where}: {problem}") from None
                    print(f"source1: skipped {where}: {problem}", file=sys.stderr)
                    continue
                if usable(rec, where):
                    yield rec
    except _READ_ERRORS as e:
        raise BadInput(f"{path} cannot be read after line {n} ({e})", skippable=False) from None


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description="Score documents with Source-1 (one JSON object per document).")
    p.add_argument("--model", default=str(Path(__file__).resolve().parent),
                   help="Source-1 directory or Hugging Face repo id (default: this file's directory)")
    p.add_argument("--input", required=True, help=".jsonl (one document per line), .json (an array of objects) or "
                   "a text file (one document); gzip, bzip2 and xz files are decompressed (docs.jsonl.gz)")
    p.add_argument("--text-field", default="text", help="JSON field holding the text (default: text); "
                   "title, url, source_type and code_language fields are used when present")
    p.add_argument("--output", help="output .jsonl (default: standard output); written to <output>.tmp and renamed "
                   "at the end, so a run that fails leaves an older output as it was")
    p.add_argument("--skip-bad", action="store_true", help="skip (and report on stderr) records that are not valid "
                   "JSON objects with a string text, instead of stopping")
    p.add_argument("--revision", help="branch, tag or commit, when --model is a Hugging Face repo id")
    p.add_argument("--device", help="cpu, cuda, cuda:1, ... (default: cuda when available)")
    p.add_argument("--precision", choices=PRECISION_CHOICES, default=DEFAULT_PRECISION, help="weights to load: "
                   "bf16 (default, model.safetensors) or fp32 (model.fp32.safetensors, the full-precision copy)")
    p.add_argument("--dtype", choices=DTYPE_CHOICES, default="auto", help="what to compute in; auto (default): "
                   "bfloat16 on a GPU with native bfloat16, else float32 (bf16 weights upcast); float16 is not "
                   "supported")
    p.add_argument("--batch-tokens", type=int, help="padded tokens per forward pass (default: "
                   f"{DEFAULT_BATCH_TOKENS_CPU} on a CPU, {DEFAULT_BATCH_TOKENS} on a GPU)")
    p.add_argument("--max-chunks", type=int, default=0, help="score at most N evenly spaced chunks per document")
    p.add_argument("--drop-line", help='"calibrated" (default), "default" (the schema\'s hard filters) or an '
                   "expression such as 'toxicity >= 4 or spam_seo >= 3'")
    p.add_argument("--apply-offsets", action="store_true", help="add the calibration offsets to the quality scores")
    p.add_argument("--show-url", action="store_true", help="show the url field to the model (untrained)")
    p.add_argument("--no-chunks", action="store_true", help="leave out the per-chunk list of split documents")
    p.add_argument("--group", type=int, default=256, help="documents scored together")
    args = p.parse_args(argv)
    source = Path(args.input)
    if not source.is_file():
        p.error(f"--input {args.input}: no such file")
    problem = _input_problem(source)
    if problem:
        p.error(f"--input {args.input} {problem}")
    if args.batch_tokens is not None and args.batch_tokens < 1:
        p.error("--batch-tokens must be at least 1")
    # Checked before the model loads. A regular output file is written as <output>.tmp, which replaces the output
    # only when the run succeeds; a device or pipe (such as /dev/stdout) is written directly.
    target = tmp = None
    if args.output:
        out_path = Path(args.output)
        if out_path.is_dir():
            p.error(f"--output {args.output} is a folder; give a file name")
        if out_path.exists() and out_path.samefile(source):
            p.error("--output is the same file as --input; refusing to overwrite it")
        if not out_path.resolve().parent.is_dir():
            p.error(f"--output {args.output}: the folder does not exist")
        if out_path.is_file() or not out_path.exists():
            target = out_path.resolve()
            tmp = target.with_name(target.name + ".tmp")
            if tmp.is_dir() or (tmp.exists() and tmp.samefile(source)):
                p.error(f"--output {args.output}: the temporary file it is written to first, {tmp}, is "
                        + ("a folder" if tmp.is_dir() else "the --input file"))

    t0 = time.time()
    try:
        model = Source1.from_pretrained(args.model, device=args.device, dtype=args.dtype, precision=args.precision,
                                        drop_line=args.drop_line, apply_offsets=args.apply_offsets,
                                        show_url=args.show_url, revision=args.revision)
    except DropLineError as e:
        p.error(str(e))
    compute = str(next(model.parameters()).dtype).replace("torch.", "")
    print(f"source1: loaded {model.weights_file} ({model.weights_dtype or '?'} weights) on {model.device}, computing "
          f"in {compute}, in {time.time() - t0:.1f} s; drop line: {model.drop_line.source}", file=sys.stderr)
    out = open(tmp or args.output, "w", encoding="utf-8") if args.output else sys.stdout
    done = seen = 0
    t0 = time.time()

    def flush(group: list[dict]) -> None:
        nonlocal done, seen
        for rec, res in zip(group, model.score_batch(group, args.batch_tokens, text_field=args.text_field,
                                                     max_chunks=args.max_chunks)):
            seen += 1
            if args.no_chunks:
                res.pop("chunks", None)
            if "id" in rec:
                res = {"id": rec["id"], **res}
            try:  # the reader lets no unwritable id through; this keeps a half-written line out of the output
                line = json.dumps(res, ensure_ascii=False, allow_nan=False)
                line.encode("utf-8")
            except (ValueError, RecursionError) as e:
                where = f"document {seen:,}" + (f" (id {rec['id']!r})" if "id" in rec else "")
                if not args.skip_bad:
                    raise BadInput(f"{where}: its scores cannot be written as JSON ({e})") from None
                print(f"source1: skipped {where}: its scores cannot be written as JSON ({e})", file=sys.stderr)
                continue
            out.write(line + "\n")
            done += 1
        print(f"source1: {done:,} documents, {done / max(time.time() - t0, 1e-9):.1f}/s", file=sys.stderr)

    def written() -> str:
        """What a run that stopped early left behind."""
        what = f"The {done:,} documents before it were" if done != 1 else "The document before it was"
        if tmp is None:
            return f"{what} written"
        if not done:
            tmp.unlink(missing_ok=True)
            return f"Nothing was written, and {args.output} was not changed"
        return f"{what} written to {tmp}, and {args.output} was not changed"

    docs = _read_docs(source, args.text_field, args.skip_bad)
    group: list[dict] = []
    try:
        try:
            while True:
                try:
                    rec = next(docs)
                except StopIteration:
                    break
                except Exception:
                    if group:
                        flush(group)  # the documents read before a bad record or a read error are still scored
                    raise
                group.append(rec)
                if len(group) >= args.group:
                    flush(group)
                    group = []
            if group:
                flush(group)
        finally:
            if out is not sys.stdout:
                out.close()
    except BadInput as e:
        raise SystemExit(f"source1: {e}. {written()}"
                         + ("; --skip-bad skips bad records" if e.skippable else "")) from None
    except BaseException as e:
        if tmp is not None:
            print(f"source1: stopped by {type(e).__name__}. {written()}", file=sys.stderr)
        raise
    if tmp is not None:
        tmp.replace(target)
    return 0


if __name__ == "__main__":
    sys.exit(main())