#!/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 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", control characters 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: | 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 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 json import math import operator import re import sys import time import unicodedata 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 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]") _TRAILING_WS = re.compile(r"[ \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), control characters, zero-width spaces and byte-order marks dropped, 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 DropLine: """A drop line such as ``toxicity >= 4 or spam_seo >= 3.5 or boilerplate >= 4.5``: field names, numbers, comparisons, ``and`` / ``or`` / ``not`` and parentheses. A comparison with a missing score (None) is false. A chunk or document is kept when the line does not match.""" _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]): self.source = source.strip() tree = ast.parse(self.source, mode="eval") for node in ast.walk(tree): if not isinstance(node, self._NODES): raise ValueError(f"{type(node).__name__} is not allowed in a drop line: {source!r}") if isinstance(node, ast.Name) and node.id not in set(names): raise ValueError(f"unknown name {node.id!r} in drop line {source!r}") self.tree = tree.body 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] 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 not isinstance(op, (ast.Eq, ast.NotEq)) and (left is None or right is None): return False try: if not _CMP[type(op)](left, right): return False except TypeError: return False left = right return True raise ValueError(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 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 {} default_line = " or ".join((self.schema.get("composite") or {}).get("hard_filters") or []) or "False" if drop_line in (None, "calibrated"): drop_line = (self.calibration.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 = default_line self.drop_line = DropLine(drop_line, set(self.fields) | {"overall", "tokens", "parts"}) 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) tok = Tokenizer.from_file(str(path / "tokenizer.json")) tok.no_truncation() tok.no_padding() variant = {"variant": "fp32"} if precision == "fp32" else {} backbone = AutoModel.from_pretrained(str(path), **{_dtype_kwarg(): dtype}, **variant, **backbone_kwargs) 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 score_inputs(self, inputs: list[str], batch_tokens: int = DEFAULT_BATCH_TOKENS) -> 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; a batch that runs out of GPU memory is split in half and retried.""" 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 = DEFAULT_BATCH_TOKENS, *, 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). 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 = DEFAULT_BATCH_TOKENS) -> 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).""" 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. Invalid UTF-8 is replaced (with a warning for JSON). A bad record raises BadInput, or with ``skip_bad`` is reported on stderr and skipped. 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" 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") suffix = path.suffix.lower() if suffix not in (".jsonl", ".ndjson", ".json"): yield {text_field: path.read_text(encoding="utf-8", errors="replace"), "id": path.name} return if suffix == ".json": try: data = json.loads(decode(path.read_bytes(), str(path)).lstrip("")) except json.JSONDecodeError: data = None # not a single JSON value: read the file 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 with open(path, "rb") 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 json.JSONDecodeError as e: if not skip_bad: raise BadInput(f"{where}: invalid JSON ({e.msg} at column {e.colno})") from None print(f"source1: skipped {where}: invalid JSON ({e.msg} at column {e.colno})", file=sys.stderr) continue if usable(rec, where): yield rec 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)") 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)") 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, default=DEFAULT_BATCH_TOKENS, help="padded tokens per forward pass") 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) if not Path(args.input).is_file(): p.error(f"--input {args.input}: no such file") t0 = time.time() 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) 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(args.output, "w", encoding="utf-8") if args.output else sys.stdout done = 0 t0 = time.time() def flush(group: list[dict]) -> None: nonlocal done for rec, res in zip(group, model.score_batch(group, args.batch_tokens, text_field=args.text_field, max_chunks=args.max_chunks)): if args.no_chunks: res.pop("chunks", None) if "id" in rec: res = {"id": rec["id"], **res} out.write(json.dumps(res, ensure_ascii=False, allow_nan=False) + "\n") done += len(group) print(f"source1: {done:,} documents, {done / max(time.time() - t0, 1e-9):.1f}/s", file=sys.stderr) try: group: list[dict] = [] try: for rec in _read_docs(Path(args.input), args.text_field, args.skip_bad): group.append(rec) if len(group) >= args.group: flush(group) group = [] except BadInput as e: if group: flush(group) # the documents read before the bad record are still scored and written raise SystemExit(f"source1: {e}. The {done:,} documents before it were written; --skip-bad skips " "bad records") from None if group: flush(group) finally: if out is not sys.stdout: out.close() return 0 if __name__ == "__main__": sys.exit(main())