From 5153ff86b8cddfaea7afa45afc30f1fbf1134e0e Mon Sep 17 00:00:00 2001 From: Kaushik Kumar Date: Mon, 31 Aug 2026 20:14:18 +0530 Subject: [PATCH 1/5] Add model training and export commands Package the DeBERTa BIOES and CRF training pipeline with strict dataset and checkpoint validation. Add ONNX export, focused tests, documentation, and installed command entry points. Signed-off-by: Kaushik Kumar --- CHANGELOG.rst | 1 + MANIFEST.in | 1 + README.rst | 14 + azure-pipelines.yml | 23 + conftest.py | 15 + docs/source/index.rst | 1 + docs/source/training.rst | 59 ++ pyproject.toml | 2 +- setup.cfg | 16 + src/scancode_required_phrases/export.py | 233 +++++++ src/scancode_required_phrases/model.py | 228 +++++++ src/scancode_required_phrases/training.py | 752 ++++++++++++++++++++++ tests/test_export.py | 42 ++ tests/test_model.py | 139 ++++ tests/test_training.py | 332 ++++++++++ 15 files changed, 1857 insertions(+), 1 deletion(-) create mode 100644 conftest.py create mode 100644 docs/source/training.rst create mode 100644 src/scancode_required_phrases/export.py create mode 100644 src/scancode_required_phrases/model.py create mode 100644 src/scancode_required_phrases/training.py create mode 100644 tests/test_export.py create mode 100644 tests/test_model.py create mode 100644 tests/test_training.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ed95994..2f482ee 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -9,3 +9,4 @@ v0.0.0 - Add required phrase dataset extraction. - Add composite rule required phrase updates. +- Add required phrase model training and ONNX export. diff --git a/MANIFEST.in b/MANIFEST.in index 8767f1e..842e574 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -12,6 +12,7 @@ include *.rst include *.png include setup.* include configure* +include conftest.py include requirements* include .dockerignore include .gitignore diff --git a/README.rst b/README.rst index c2e65c0..3aadee4 100644 --- a/README.rst +++ b/README.rst @@ -35,6 +35,20 @@ The command uses existing required phrases from single-key rules to update composite rules. A rule is updated only when every relevant license key has a non-overlapping match. It operates on the installed ScanCode rules directory. +Train and export a model +======================== + +Install the optional training dependencies and train from a generated dataset: + +.. code-block:: console + + python -m pip install ".[training,training-8bit]" + train-required-phrase-model --data-dir dataset-output --output-dir model-output \ + --optimizer adamw-8bit + +Test-set evaluation is opt-in and should only be used for a final selected run. +See ``docs/source/training.rst`` for training and ONNX export details. + Development =========== diff --git a/azure-pipelines.yml b/azure-pipelines.yml index 3b17abc..1619850 100644 --- a/azure-pipelines.yml +++ b/azure-pipelines.yml @@ -7,6 +7,29 @@ jobs: + - job: ubuntu24_ml_tests + pool: + vmImage: ubuntu-24.04 + steps: + - checkout: self + fetchDepth: 10 + - task: UsePythonVersion@0 + inputs: + versionSpec: '3.12' + architecture: x64 + displayName: Install Python 3.12 + - script: | + ./configure --clean + ./configure --dev + venv/bin/pip install -e ".[training]" + displayName: Configure training test environment + - script: | + USE_TF=0 venv/bin/pytest -q \ + tests/test_training.py \ + tests/test_model.py \ + tests/test_export.py + displayName: Run training unit tests + - template: etc/ci/azure-posix.yml parameters: job_name: run_code_checks diff --git a/conftest.py b/conftest.py new file mode 100644 index 0000000..e5538d9 --- /dev/null +++ b/conftest.py @@ -0,0 +1,15 @@ +# -*- coding: utf-8 -*- +# +# Copyright (c) nexB Inc. and others. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Pytest collection settings for optional training dependencies.""" + +import importlib.util + + +training_dependencies = ("torch", "torchcrf", "transformers") +if any(importlib.util.find_spec(name) is None for name in training_dependencies): + collect_ignore = ["src/scancode_required_phrases/model.py"] +else: + collect_ignore = [] diff --git a/docs/source/index.rst b/docs/source/index.rst index daf7c60..779611d 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -10,6 +10,7 @@ license rules. dataset composite_rules + training contribute/contrib_doc Indices and tables diff --git a/docs/source/training.rst b/docs/source/training.rst new file mode 100644 index 0000000..4594e95 --- /dev/null +++ b/docs/source/training.rst @@ -0,0 +1,59 @@ +Train and export a required phrase model +========================================= + +The training command fine-tunes a DeBERTa token classifier with BIOES labels +and an optional word-level CRF. It reads the JSONL files produced by +``build-required-phrases-dataset``. + +Installation +------------ + +Install the training dependencies before running the command. The 8-bit +optimizer is a separate extra because it is only needed on constrained GPUs. + +.. code-block:: console + + python -m pip install ".[training,training-8bit]" + +Install the ONNX dependencies when an export is needed: + +.. code-block:: console + + python -m pip install ".[onnx]" + +Training +-------- + +.. code-block:: console + + train-required-phrase-model \ + --data-dir dataset-output \ + --output-dir model-output \ + --optimizer adamw-8bit + +The default configuration uses FP32 and standard AdamW. The 8-bit optimizer is +recommended for DeBERTa-large on a 16 GB GPU. Test-set evaluation is disabled by +default so repeated experiments select models using validation data only. Add +``--evaluate-test`` only for the final selected run. + +The command validates the dataset, skips examples whose required phrase is cut +by tokenizer truncation, and saves the checkpoint with the best validation +strict span F1. The output directory must be empty unless ``--resume`` is used. +It writes ``train_config.json`` for inference compatibility and +``run_manifest.json`` for reproducibility. + +Export +------ + +.. code-block:: console + + export-required-phrase-model \ + --model-dir model-output \ + --output-dir onnx-output + +The export command strictly loads the saved model, exports the DeBERTa emissions +to ONNX, stores the CRF transitions separately, and verifies that NumPy and +PyTorch Viterbi decoding agree. + +Generated datasets, checkpoints, model files, and export artifacts must not be +committed to this repository. diff --git a/pyproject.toml b/pyproject.toml index 39bc302..f451b33 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -103,7 +103,7 @@ force-single-line = true lines-after-imports = 1 default-section = "first-party" known-first-party = ["scancode_required_phrases", "tests", "etc/scripts/**/*.py"] -known-third-party = ["click", "pytest"] +known-third-party = ["click", "numpy", "pytest", "safetensors", "torch", "torchcrf", "transformers"] sections = { django = ["django"] } section-order = [ diff --git a/setup.cfg b/setup.cfg index fe37819..5045e8a 100644 --- a/setup.cfg +++ b/setup.cfg @@ -56,9 +56,25 @@ where = src console_scripts = add-composite-required-phrases = scancode_required_phrases.composite_rules:add_composite_required_phrases build-required-phrases-dataset = scancode_required_phrases.dataset:main + export-required-phrase-model = scancode_required_phrases.export:main + train-required-phrase-model = scancode_required_phrases.training:main [options.extras_require] +training = + accelerate >= 0.33 + protobuf >= 3.20 + pytorch-crf >= 0.7.2 + safetensors >= 0.4 + sentencepiece >= 0.2 + torch >= 2.0 + transformers >= 4.44 +training-8bit = + bitsandbytes >= 0.43 +onnx = + numpy + onnx >= 1.16 + onnxruntime >= 1.18 dev = pytest >= 7.0.1 pytest-xdist >= 2 diff --git a/src/scancode_required_phrases/export.py b/src/scancode_required_phrases/export.py new file mode 100644 index 0000000..17a55c8 --- /dev/null +++ b/src/scancode_required_phrases/export.py @@ -0,0 +1,233 @@ +# -*- coding: utf-8 -*- +# +# Copyright (c) nexB Inc. and others. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Export a trained required phrase tagger for CPU inference.""" + +import hashlib +import json +import os +from pathlib import Path +from types import SimpleNamespace + +import click + +os.environ.setdefault("USE_TF", "0") + + +def viterbi_decode(emissions, start_transitions, transitions, end_transitions): + """Return the best tag path for one sequence.""" + sequence_length = emissions.shape[0] + score = start_transitions + emissions[0] + backpointers = [] + + for step in range(1, sequence_length): + candidates = score[:, None] + transitions + best_source = candidates.argmax(axis=0) + score = candidates.max(axis=0) + emissions[step] + backpointers.append(best_source) + + score = score + end_transitions + best = int(score.argmax()) + path = [best] + for sources in reversed(backpointers): + best = int(sources[best]) + path.append(best) + path.reverse() + return path + + +def sha256(path): + """Return the hexadecimal SHA256 digest of a file.""" + digest = hashlib.sha256() + with open(path, "rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def build_emissions_module(tagger): + """Wrap the trained backbone and classifier for ONNX export.""" + import torch.nn as nn + + class EmissionsModule(nn.Module): + def __init__(self): + super().__init__() + self.backbone = tagger.backbone + self.classifier = tagger.classifier + + def forward(self, input_ids, attention_mask): + hidden = self.backbone( + input_ids=input_ids, + attention_mask=attention_mask, + ).last_hidden_state + return self.classifier(hidden) + + return EmissionsModule().eval() + + +def load_tagger(model_dir, train_config): + """Rebuild a tagger and strictly load its saved weights.""" + import torch + from safetensors.torch import load_file + + from scancode_required_phrases.model import PhraseTagger + + config = SimpleNamespace( + model_name=train_config["model_name"], + model_revision=train_config.get("model_revision"), + use_crf=train_config["use_crf"], + aux_ce_weight=0.0, + label_weights=[1.0] * len(train_config["labels"]), + ) + tagger = PhraseTagger(config) + + model_dir = Path(model_dir) + safetensors_file = model_dir / "model.safetensors" + pytorch_file = model_dir / "pytorch_model.bin" + if safetensors_file.exists(): + state = load_file(str(safetensors_file)) + elif pytorch_file.exists(): + state = torch.load(pytorch_file, map_location="cpu", weights_only=True) + else: + raise FileNotFoundError(f"No model weights found in {model_dir}") + + # Checkpoints created before class weights became non-persistent contain + # this training-only tensor. + state.pop("class_weights", None) + for name, tensor in state.items(): + if not torch.isfinite(tensor).all(): + raise ValueError(f"Checkpoint tensor {name!r} contains non-finite values") + + tagger.load_state_dict(state, strict=True) + return tagger.eval() + + +def check_viterbi_matches_crf(tagger, num_tags): + """Verify NumPy and pytorch-crf return identical paths.""" + import torch + + start = tagger.crf.start_transitions.detach().cpu().numpy() + transitions = tagger.crf.transitions.detach().cpu().numpy() + end = tagger.crf.end_transitions.detach().cpu().numpy() + + emissions = torch.randn(3, 14, num_tags) + mask = torch.ones(3, 14, dtype=torch.bool) + crf_paths = tagger.crf.decode(emissions, mask=mask) + for row in range(emissions.size(0)): + numpy_path = viterbi_decode(emissions[row].numpy(), start, transitions, end) + if numpy_path != crf_paths[row]: + raise AssertionError("NumPy Viterbi disagrees with pytorch-crf decoding") + + return start, transitions, end + + +def export(model_dir, output_dir, opset): + """Export ONNX emissions, CRF transitions, and a checksum manifest.""" + import numpy as np + import torch + from transformers import AutoTokenizer + + model_dir = Path(model_dir) + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + config_path = model_dir / "train_config.json" + if not config_path.exists(): + raise FileNotFoundError(f"No train_config.json found in {model_dir}") + train_config = json.loads(config_path.read_text(encoding="utf-8")) + labels = train_config["labels"] + use_crf = train_config["use_crf"] + + tagger = load_tagger(model_dir, train_config) + emissions_module = build_emissions_module(tagger) + tokenizer = AutoTokenizer.from_pretrained(str(model_dir), use_fast=True) + + sample = tokenizer( + "Licensed under the Apache License Version 2.0", + return_tensors="pt", + ) + inputs = sample["input_ids"], sample["attention_mask"] + + onnx_path = output_dir / "model.onnx" + torch.onnx.export( + emissions_module, + inputs, + str(onnx_path), + input_names=["input_ids", "attention_mask"], + output_names=["emissions"], + dynamic_axes={ + "input_ids": {0: "batch", 1: "sequence"}, + "attention_mask": {0: "batch", 1: "sequence"}, + "emissions": {0: "batch", 1: "sequence"}, + }, + opset_version=opset, + do_constant_folding=True, + ) + + manifest = { + "labels": labels, + "onnx_model": sha256(onnx_path), + } + + if use_crf: + start, transitions, end = check_viterbi_matches_crf(tagger, len(labels)) + transitions_path = output_dir / "crf_transitions.npz" + np.savez( + transitions_path, + start=start, + transitions=transitions, + end=end, + ) + manifest["crf_transitions"] = sha256(transitions_path) + + import onnxruntime + + session = onnxruntime.InferenceSession( + str(onnx_path), + providers=["CPUExecutionProvider"], + ) + feeds = { + "input_ids": sample["input_ids"].numpy(), + "attention_mask": sample["attention_mask"].numpy(), + } + onnx_emissions = session.run(["emissions"], feeds)[0] + torch_emissions = emissions_module(*inputs).detach().numpy() + if not np.allclose(onnx_emissions, torch_emissions, atol=1e-3): + raise AssertionError("ONNX emissions differ from PyTorch emissions") + + manifest_path = output_dir / "manifest.json" + manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8") + click.echo(f"wrote {onnx_path}") + if use_crf: + click.echo(f"wrote {output_dir / 'crf_transitions.npz'}") + click.echo(f"wrote {manifest_path}") + + +@click.command() +@click.option( + "--model-dir", + required=True, + type=click.Path(exists=True, file_okay=False, path_type=Path), + help="Directory containing the trained model and train_config.json.", +) +@click.option( + "--output-dir", + default=None, + type=click.Path(file_okay=False, path_type=Path), + help="Output directory; defaults to the model directory.", +) +@click.option("--opset", default=14, type=int, show_default=True) +def main(model_dir, output_dir, opset): + """Export a trained required phrase tagger to ONNX.""" + try: + export(model_dir, output_dir or model_dir, opset) + except ImportError as error: + raise click.ClickException( + f"{error}; install scancode-required-phrases[training,onnx]" + ) from error + + +if __name__ == "__main__": + main() diff --git a/src/scancode_required_phrases/model.py b/src/scancode_required_phrases/model.py new file mode 100644 index 0000000..309475d --- /dev/null +++ b/src/scancode_required_phrases/model.py @@ -0,0 +1,228 @@ +# -*- coding: utf-8 -*- +# +# Copyright (c) nexB Inc. and others. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""DeBERTa model and Trainer support for required phrase tagging.""" + +import os + +os.environ.setdefault("USE_TF", "0") + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.optim import AdamW +from torchcrf import CRF +from transformers import AutoModel +from transformers import Trainer + +from scancode_required_phrases.training import first_subword_positions +from scancode_required_phrases.training import IGNORE_INDEX +from scancode_required_phrases.training import LABELS + + +class PhraseTagger(nn.Module): + """DeBERTa backbone with a word-level token classifier and optional CRF.""" + + def __init__(self, config): + super().__init__() + self.use_crf = config.use_crf + self.aux_ce_weight = config.aux_ce_weight + self.num_labels = len(LABELS) + + self.backbone = AutoModel.from_pretrained( + config.model_name, + revision=config.model_revision, + ).float() + self.backbone.gradient_checkpointing_enable() + self.backbone.enable_input_require_grads() + + hidden_size = self.backbone.config.hidden_size + dropout = getattr(self.backbone.config, "hidden_dropout_prob", 0.1) + self.dropout = nn.Dropout(dropout) + self.classifier = nn.Linear(hidden_size, self.num_labels) + + if self.use_crf: + self.crf = CRF(self.num_labels, batch_first=True) + + if self.aux_ce_weight > 0: + self.register_buffer( + "class_weights", + torch.tensor(config.label_weights, dtype=torch.float), + persistent=False, + ) + else: + self.class_weights = None + + def emissions(self, input_ids, attention_mask): + """Return per-subword label scores.""" + hidden = self.backbone( + input_ids=input_ids, + attention_mask=attention_mask, + ).last_hidden_state + return self.classifier(self.dropout(hidden)) + + def token_cross_entropy(self, emissions, labels): + """Return weighted cross entropy over labeled subwords.""" + return F.cross_entropy( + emissions.reshape(-1, self.num_labels), + labels.reshape(-1), + weight=self.class_weights, + ignore_index=IGNORE_INDEX, + ) + + def gather_words(self, emissions, labels): + """Pack first-subword emissions and labels into word-level sequences.""" + batch, _, num_labels = emissions.shape + is_word = labels.ne(IGNORE_INDEX) + lengths = is_word.sum(dim=1) + width = int(lengths.max().item()) + + word_emissions = emissions.new_zeros((batch, width, num_labels)) + crf_tags = labels.new_zeros((batch, width)) + eval_tags = labels.new_full((batch, width), IGNORE_INDEX) + mask = torch.zeros((batch, width), dtype=torch.bool, device=emissions.device) + + for row in range(batch): + positions = is_word[row].nonzero(as_tuple=True)[0] + count = positions.numel() + word_emissions[row, :count] = emissions[row, positions] + tags = labels[row, positions] + crf_tags[row, :count] = tags + eval_tags[row, :count] = tags + mask[row, :count] = True + + return word_emissions, crf_tags, eval_tags, mask + + def forward(self, input_ids, attention_mask, labels=None): + emissions = self.emissions(input_ids, attention_mask) + result = {} + + if not self.use_crf: + if labels is not None: + result["loss"] = self.token_cross_entropy(emissions, labels) + result["word_labels"] = labels + if not self.training: + result["predictions"] = emissions.argmax(dim=-1) + return result + + if labels is None: + raise ValueError("CRF head needs labels to locate words") + + word_emissions, crf_tags, eval_tags, mask = self.gather_words(emissions, labels) + word_emissions = word_emissions.float() + + log_likelihood = self.crf(word_emissions, crf_tags, mask=mask, reduction="mean") + loss = -log_likelihood + if self.aux_ce_weight > 0: + loss = loss + self.aux_ce_weight * self.token_cross_entropy(emissions, labels) + + result["loss"] = loss + result["word_labels"] = eval_tags + + if not self.training: + decoded = self.crf.decode(word_emissions, mask=mask) + result["predictions"] = self.pad_decoded(decoded, mask.size(1), emissions.device) + + return result + + def predict_words(self, input_ids, attention_mask, word_ids): + """Return one label ID per word for a single rule.""" + positions = first_subword_positions(word_ids) + if not positions: + return [] + + emissions = self.emissions(input_ids, attention_mask) + word_emissions = emissions[0, positions].unsqueeze(0).float() + if not self.use_crf: + return word_emissions.argmax(dim=-1)[0].tolist() + + mask = torch.ones(word_emissions.shape[:2], dtype=torch.bool, device=emissions.device) + return self.crf.decode(word_emissions, mask=mask)[0] + + @staticmethod + def pad_decoded(decoded, width, device): + """Return variable-length decoded paths as a padded tensor.""" + predictions = torch.full( + (len(decoded), width), + IGNORE_INDEX, + dtype=torch.long, + device=device, + ) + for row, path in enumerate(decoded): + if path: + predictions[row, : len(path)] = torch.tensor( + path, + dtype=torch.long, + device=device, + ) + return predictions + + +def build_optimizer(config, model): + """Return the configured AdamW optimizer with layer-wise learning rates.""" + num_layers = model.backbone.config.num_hidden_layers + no_decay = ("bias", "LayerNorm.weight", "layer_norm.weight") + + def rate_for(name): + if name.startswith("classifier") or name.startswith("crf"): + return config.head_lr + if ".encoder.layer." in name: + layer = int(name.split(".encoder.layer.")[1].split(".")[0]) + return config.base_lr * (config.layer_decay ** (num_layers - layer)) + return config.base_lr * (config.layer_decay ** (num_layers + 1)) + + groups = [] + for name, parameter in model.named_parameters(): + if not parameter.requires_grad: + continue + decay = 0.0 if any(part in name for part in no_decay) else config.weight_decay + groups.append( + { + "params": [parameter], + "lr": rate_for(name), + "weight_decay": decay, + } + ) + + optimizer_args = { + "lr": config.base_lr, + "eps": config.adam_epsilon, + "betas": (0.9, 0.999), + } + if config.optimizer == "adamw": + return AdamW(groups, **optimizer_args) + + if config.optimizer == "adamw-8bit": + try: + from bitsandbytes.optim import AdamW8bit + except ImportError as error: + raise RuntimeError( + "adamw-8bit requires bitsandbytes; install the training dependencies" + ) from error + return AdamW8bit(groups, **optimizer_args) + + raise ValueError(f"Unsupported optimizer: {config.optimizer}") + + +class PhraseTrainer(Trainer): + """Trainer adapter for PhraseTagger output dictionaries.""" + + def compute_loss(self, model, inputs, return_outputs=False, **kwargs): + outputs = model(**inputs) + loss = outputs["loss"] + if not torch.isfinite(loss): + raise FloatingPointError("Training produced a non-finite loss") + return (loss, outputs) if return_outputs else loss + + def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None): + inputs = self._prepare_inputs(inputs) + with torch.no_grad(): + outputs = model(**inputs) + loss = outputs.get("loss") + if loss is not None: + loss = loss.detach() + if prediction_loss_only: + return loss, None, None + return loss, outputs["predictions"], outputs["word_labels"] diff --git a/src/scancode_required_phrases/training.py b/src/scancode_required_phrases/training.py new file mode 100644 index 0000000..123b465 --- /dev/null +++ b/src/scancode_required_phrases/training.py @@ -0,0 +1,752 @@ +# -*- coding: utf-8 -*- +# +# Copyright (c) nexB Inc. and others. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Train a DeBERTa BIOES tagger for required phrase spans.""" + +from dataclasses import asdict +from dataclasses import dataclass +from dataclasses import field +import hashlib +import importlib.metadata +import inspect +import json +import os +from pathlib import Path +import platform +import random + +import click + +os.environ.setdefault("USE_TF", "0") + + +LABELS = ["O", "B-REQ", "I-REQ", "E-REQ", "S-REQ"] +LABEL2ID = {label: index for index, label in enumerate(LABELS)} +ID2LABEL = {index: label for index, label in enumerate(LABELS)} +IGNORE_INDEX = -100 + +MODEL_NAME = "microsoft/deberta-v3-large" +MAX_LENGTH = 512 +RECORD_FIELDS = ( + "identifier", + "license_expression", + "rule_type", + "text", + "tokens", + "bioes_labels", +) +START_LABELS = {"O", "B-REQ", "S-REQ"} +END_LABELS = {"O", "E-REQ", "S-REQ"} +VALID_TRANSITIONS = { + "O": {"O", "B-REQ", "S-REQ"}, + "B-REQ": {"I-REQ", "E-REQ"}, + "I-REQ": {"I-REQ", "E-REQ"}, + "E-REQ": {"O", "B-REQ", "S-REQ"}, + "S-REQ": {"O", "B-REQ", "S-REQ"}, +} + + +@dataclass +class Config: + """Settings for one training run.""" + + data_dir: Path + output_dir: Path + model_name: str = MODEL_NAME + model_revision: str | None = None + max_length: int = MAX_LENGTH + + epochs: int = 8 + batch_size: int = 1 + grad_accum: int = 16 + base_lr: float = 2e-5 + head_lr: float = 1e-4 + layer_decay: float = 0.98 + weight_decay: float = 0.01 + warmup_ratio: float = 0.1 + max_grad_norm: float = 0.5 + adam_epsilon: float = 1e-6 + early_stopping_patience: int = 3 + optimizer: str = "adamw" + precision: str = "fp32" + + limit: int = 0 + resume: bool = False + use_crf: bool = True + aux_ce_weight: float = 0.3 + evaluate_test: bool = False + with_isr: bool = False + seed: int = 42 + + label_weights: list = field(default_factory=lambda: [1.0, 2.0, 1.5, 1.5, 2.0]) + + +def validate_config(config): + """Validate settings before loading the tokenizer or model.""" + if config.optimizer not in {"adamw", "adamw-8bit"}: + raise ValueError(f"Unsupported optimizer: {config.optimizer}") + if config.precision not in {"fp32", "bf16"}: + raise ValueError(f"Unsupported precision: {config.precision}") + if config.epochs < 1 or config.batch_size < 1 or config.grad_accum < 1: + raise ValueError("Epochs, batch size, and gradient accumulation must be positive") + if config.base_lr <= 0 or config.head_lr <= 0: + raise ValueError("Learning rates must be positive") + if config.aux_ce_weight < 0: + raise ValueError("Auxiliary loss weight cannot be negative") + if len(config.label_weights) != len(LABELS): + raise ValueError(f"Expected {len(LABELS)} label weights") + + +def prepare_output_dir(output_dir, resume): + """Create a new output directory or validate a resumable one.""" + output_dir = Path(output_dir) + checkpoints = list(output_dir.glob("checkpoint-*")) if output_dir.exists() else [] + if resume and not checkpoints: + raise ValueError(f"No checkpoint found in {output_dir}") + if not resume and output_dir.exists() and any(output_dir.iterdir()): + raise ValueError(f"Output directory is not empty: {output_dir}") + output_dir.mkdir(parents=True, exist_ok=True) + + +def set_seed(seed): + """Seed Python, NumPy, and PyTorch.""" + import numpy as np + import torch + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def load_jsonl(path): + """Yield parsed records from a JSONL file.""" + with open(path, encoding="utf-8") as lines: + for line_number, line in enumerate(lines, 1): + line = line.strip() + if line: + yield line_number, json.loads(line) + + +def validate_bioes(labels): + """Return an error for an invalid BIOES sequence, or None.""" + if not labels: + return "has no labels" + unknown = sorted(set(labels) - set(LABELS)) + if unknown: + return f"contains unknown labels: {unknown}" + if labels[0] not in START_LABELS: + return f"starts with {labels[0]}" + for previous, current in zip(labels, labels[1:]): + if current not in VALID_TRANSITIONS[previous]: + return f"contains invalid transition {previous} -> {current}" + if labels[-1] not in END_LABELS: + return f"ends with {labels[-1]}" + + +def validate_record(record, path, line_number): + """Validate one dataset record and return it.""" + location = f"{path} line {line_number}" + for field_name in RECORD_FIELDS: + if field_name not in record: + raise ValueError(f"{location}: missing {field_name!r}") + + identifier = record["identifier"] + if not identifier: + raise ValueError(f"{location}: empty identifier") + if not record["license_expression"]: + raise ValueError(f"{location} ({identifier}): empty license expression") + if not record["rule_type"]: + raise ValueError(f"{location} ({identifier}): empty rule type") + + tokens = record["tokens"] + labels = record["bioes_labels"] + if not tokens: + raise ValueError(f"{location} ({identifier}): no tokens") + if len(tokens) != len(labels): + raise ValueError( + f"{location} ({identifier}): {len(tokens)} tokens and {len(labels)} labels" + ) + + unknown = sorted(set(labels) - set(LABELS)) + if unknown: + raise ValueError(f"{location} ({identifier}): unknown labels: {unknown}") + + error = validate_bioes(labels) + if error: + raise ValueError(f"{location} ({identifier}): {error}") + return record + + +def align_labels(tokens, word_labels, tokenizer, max_length): + """Tokenize words and label only the first subword of each word.""" + encoding = tokenizer( + tokens, + is_split_into_words=True, + truncation=True, + max_length=max_length, + ) + + word_ids = encoding.word_ids() + label_ids = [] + previous_word = None + for word_id in word_ids: + if word_id is None: + label_ids.append(IGNORE_INDEX) + elif word_id != previous_word: + label_ids.append(LABEL2ID[word_labels[word_id]]) + else: + label_ids.append(IGNORE_INDEX) + previous_word = word_id + + encoding["labels"] = label_ids + kept_word_ids = [word_id for word_id in word_ids if word_id is not None] + kept_words = max(kept_word_ids) + 1 if kept_word_ids else 0 + truncated = kept_words < len(tokens) + cut_phrase = ( + truncated + and kept_words + and word_labels[kept_words - 1] + in { + "B-REQ", + "I-REQ", + } + ) + return encoding, truncated, bool(cut_phrase) + + +def first_subword_positions(word_ids): + """Return positions that start a tokenizer word.""" + positions = [] + previous = None + for index, word_id in enumerate(word_ids): + if word_id is None: + previous = None + continue + if word_id != previous: + positions.append(index) + previous = word_id + return positions + + +class PhraseDataset: + """Read and encode one BIOES JSONL split.""" + + def __init__(self, path, tokenizer, max_length, limit=0): + self.examples = [] + self.identifiers = [] + self.truncated = 0 + self.cut_phrases = 0 + + for line_number, unvalidated_record in load_jsonl(path): + if limit and len(self.examples) >= limit: + break + record = validate_record(unvalidated_record, path, line_number) + encoding, truncated, cut_phrase = align_labels( + record["tokens"], + record["bioes_labels"], + tokenizer, + max_length, + ) + if truncated: + self.truncated += 1 + if cut_phrase: + self.cut_phrases += 1 + continue + + self.identifiers.append(record["identifier"]) + self.examples.append( + { + "input_ids": encoding["input_ids"], + "attention_mask": encoding["attention_mask"], + "labels": encoding["labels"], + } + ) + + def __len__(self): + return len(self.examples) + + def __getitem__(self, index): + return self.examples[index] + + +def validate_splits(datasets): + """Require non-empty splits with disjoint rule identifiers.""" + seen = {} + for split_name, dataset in datasets.items(): + if not dataset: + raise ValueError(f"{split_name} split has no usable examples") + for identifier in dataset.identifiers: + previous_split = seen.get(identifier) + if previous_split: + raise ValueError( + f"Rule {identifier!r} occurs in both {previous_split} and {split_name}" + ) + seen[identifier] = split_name + + +def extract_spans(tags): + """Return the set of inclusive word spans in a BIOES sequence.""" + spans = [] + start = None + for index, tag in enumerate(tags): + if tag == "S-REQ": + spans.append((index, index)) + start = None + elif tag == "B-REQ": + if start is not None: + spans.append((start, index - 1)) + start = index + elif tag == "I-REQ": + if start is None: + start = index + elif tag == "E-REQ": + if start is None: + start = index + spans.append((start, index)) + start = None + elif start is not None: + spans.append((start, index - 1)) + start = None + if start is not None: + spans.append((start, len(tags) - 1)) + return set(spans) + + +def decode_row(pred_row, label_row): + """Drop ignored positions and map label IDs to BIOES tags.""" + predicted = [] + actual = [] + for prediction, label in zip(pred_row, label_row): + if int(label) == IGNORE_INDEX: + continue + actual.append(ID2LABEL[int(label)]) + predicted.append(ID2LABEL.get(int(prediction), "O")) + return predicted, actual + + +def compute_metrics(eval_pred): + """Return strict span-level micro metrics and rule-level exact match.""" + predictions, labels = eval_pred + true_positive = false_positive = false_negative = 0 + exact = 0 + + for pred_row, label_row in zip(predictions, labels): + predicted, actual = decode_row(pred_row, label_row) + predicted_spans = extract_spans(predicted) + actual_spans = extract_spans(actual) + true_positive += len(predicted_spans & actual_spans) + false_positive += len(predicted_spans - actual_spans) + false_negative += len(actual_spans - predicted_spans) + if predicted_spans == actual_spans: + exact += 1 + + precision_denominator = true_positive + false_positive + recall_denominator = true_positive + false_negative + precision = true_positive / precision_denominator if precision_denominator else 0.0 + recall = true_positive / recall_denominator if recall_denominator else 0.0 + f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0 + return { + "f1": f1, + "precision": precision, + "recall": recall, + "exact_match": exact / len(predictions) if len(predictions) else 0.0, + "predicted_spans": true_positive + false_positive, + "gold_spans": true_positive + false_negative, + } + + +def evaluate_isr(records, model, tokenizer, max_length): + """Return the fraction of predicted phrases ScanCode can locate.""" + import torch + from licensedcode.required_phrases import find_phrase_spans_in_text + + device = next(model.parameters()).device + model.eval() + total = 0 + injectable = 0 + for record in records: + encoding, _, cut_phrase = align_labels( + record["tokens"], + record["bioes_labels"], + tokenizer, + max_length, + ) + if cut_phrase: + continue + inputs = { + "input_ids": torch.tensor([encoding["input_ids"]], device=device), + "attention_mask": torch.tensor([encoding["attention_mask"]], device=device), + "labels": torch.tensor([encoding["labels"]], device=device), + } + with torch.no_grad(): + output = model(**inputs) + tags, _ = decode_row( + output["predictions"][0].tolist(), + output["word_labels"][0].tolist(), + ) + for start, end in extract_spans(tags): + if end >= len(record["tokens"]): + continue + phrase = " ".join(record["tokens"][start : end + 1]) + total += 1 + if find_phrase_spans_in_text(record["text"], phrase): + injectable += 1 + + return injectable / total if total else 0.0 + + +def sha256(path): + """Return the hexadecimal SHA256 digest of a file.""" + digest = hashlib.sha256() + with open(path, "rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def installed_version(package_name): + """Return an installed package version, or None.""" + try: + return importlib.metadata.version(package_name) + except importlib.metadata.PackageNotFoundError: + return None + + +def serializable_config(config): + """Return the training configuration with paths converted to strings.""" + values = asdict(config) + values["data_dir"] = str(values["data_dir"]) + values["output_dir"] = str(values["output_dir"]) + return values + + +def validate_saved_state(model, model_path): + """Validate the keys, shapes, and values in a saved safetensors model.""" + import torch + from safetensors.torch import load_file + + saved = load_file(str(model_path)) + expected = model.state_dict() + if set(saved) != set(expected): + missing = sorted(set(expected) - set(saved)) + unexpected = sorted(set(saved) - set(expected)) + raise ValueError(f"Saved model key mismatch; missing={missing}, unexpected={unexpected}") + + for name, tensor in saved.items(): + if tensor.shape != expected[name].shape: + raise ValueError( + f"Saved tensor {name!r} has shape {tuple(tensor.shape)}, " + f"expected {tuple(expected[name].shape)}" + ) + if not torch.isfinite(tensor).all(): + raise ValueError(f"Saved tensor {name!r} contains non-finite values") + + +def validate_precision(precision): + """Validate the selected training precision.""" + import torch + + if precision == "bf16" and not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()): + raise ValueError("bf16 requires a CUDA device with BF16 support") + + +def run_training(config): + """Train, validate, and save a required phrase model.""" + import torch + import transformers + from transformers import AutoTokenizer + from transformers import DataCollatorForTokenClassification + from transformers import EarlyStoppingCallback + from transformers import TrainingArguments + + from scancode_required_phrases.model import PhraseTagger + from scancode_required_phrases.model import PhraseTrainer + from scancode_required_phrases.model import build_optimizer + + if config.with_isr and not config.evaluate_test: + raise ValueError("ISR evaluation requires --evaluate-test") + + validate_config(config) + validate_precision(config.precision) + prepare_output_dir(config.output_dir, config.resume) + set_seed(config.seed) + + tokenizer = AutoTokenizer.from_pretrained( + config.model_name, + revision=config.model_revision, + use_fast=True, + ) + if not tokenizer.is_fast: + raise RuntimeError("Training requires a fast tokenizer with word IDs") + + paths = { + "train": config.data_dir / "train.jsonl", + "validation": config.data_dir / "val.jsonl", + "test": config.data_dir / "test.jsonl", + } + datasets = { + name: PhraseDataset(path, tokenizer, config.max_length, config.limit) + for name, path in paths.items() + } + validate_splits(datasets) + + for name, dataset in datasets.items(): + click.echo( + f"{name}: {len(dataset)} examples, {dataset.truncated} truncated, " + f"{dataset.cut_phrases} skipped with a cut phrase" + ) + + model = PhraseTagger(config) + collator = DataCollatorForTokenClassification( + tokenizer, + label_pad_token_id=IGNORE_INDEX, + ) + + arguments = TrainingArguments( + output_dir=str(config.output_dir), + num_train_epochs=config.epochs, + per_device_train_batch_size=config.batch_size, + per_device_eval_batch_size=config.batch_size, + gradient_accumulation_steps=config.grad_accum, + learning_rate=config.base_lr, + weight_decay=config.weight_decay, + warmup_ratio=config.warmup_ratio, + lr_scheduler_type="cosine", + max_grad_norm=config.max_grad_norm, + eval_strategy="epoch", + save_strategy="epoch", + save_total_limit=2, + load_best_model_at_end=True, + metric_for_best_model="f1", + greater_is_better=True, + bf16=config.precision == "bf16", + fp16=False, + logging_steps=50, + report_to="none", + seed=config.seed, + data_seed=config.seed, + dataloader_num_workers=2, + save_safetensors=True, + ) + + trainer_kwargs = { + "model": model, + "args": arguments, + "train_dataset": datasets["train"], + "eval_dataset": datasets["validation"], + "data_collator": collator, + "compute_metrics": compute_metrics, + "optimizers": (build_optimizer(config, model), None), + "callbacks": [ + EarlyStoppingCallback( + early_stopping_patience=config.early_stopping_patience, + ) + ], + } + if "processing_class" in inspect.signature(PhraseTrainer.__init__).parameters: + trainer_kwargs["processing_class"] = tokenizer + else: + trainer_kwargs["tokenizer"] = tokenizer + + trainer = PhraseTrainer(**trainer_kwargs) + trainer.train(resume_from_checkpoint=config.resume or None) + + trainer.save_model(str(config.output_dir)) + tokenizer.save_pretrained(str(config.output_dir)) + + resolved_revision = getattr(model.backbone.config, "_commit_hash", None) + train_config = { + "model_name": config.model_name, + "model_revision": resolved_revision or config.model_revision, + "use_crf": config.use_crf, + "max_length": config.max_length, + "labels": LABELS, + } + (config.output_dir / "train_config.json").write_text( + json.dumps(train_config, indent=2), + encoding="utf-8", + ) + + model_path = config.output_dir / "model.safetensors" + validate_saved_state(model, model_path) + + validation_metrics = trainer.evaluate( + datasets["validation"], + metric_key_prefix="validation", + ) + test_metrics = None + if config.evaluate_test: + test_metrics = trainer.evaluate( + datasets["test"], + metric_key_prefix="test", + ) + if config.with_isr: + test_records = [ + validate_record(record, paths["test"], line_number) + for line_number, record in load_jsonl(paths["test"]) + ] + test_metrics["test_isr"] = evaluate_isr( + test_records, + model, + tokenizer, + config.max_length, + ) + + manifest = { + "config": serializable_config(config), + "dataset": { + name: { + "path": str(path), + "sha256": sha256(path), + "examples": len(datasets[name]), + "truncated": datasets[name].truncated, + "cut_phrases": datasets[name].cut_phrases, + } + for name, path in paths.items() + }, + "model_revision": resolved_revision or config.model_revision, + "best_checkpoint": trainer.state.best_model_checkpoint, + "best_validation_f1": trainer.state.best_metric, + "validation_metrics": validation_metrics, + "test_metrics": test_metrics, + "log_history": trainer.state.log_history, + "versions": { + "python": platform.python_version(), + "torch": torch.__version__, + "transformers": transformers.__version__, + "scancode_toolkit": installed_version("scancode-toolkit"), + "pytorch_crf": installed_version("pytorch-crf"), + }, + "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None, + } + (config.output_dir / "run_manifest.json").write_text( + json.dumps(manifest, indent=2), + encoding="utf-8", + ) + + click.echo(f"best checkpoint: {trainer.state.best_model_checkpoint}") + click.echo(f"best validation F1: {trainer.state.best_metric}") + click.echo(f"validation: {validation_metrics}") + if test_metrics is not None: + click.echo(f"test: {test_metrics}") + + return { + "validation": validation_metrics, + "test": test_metrics, + } + + +@click.command() +@click.option( + "--data-dir", + required=True, + type=click.Path(exists=True, file_okay=False, path_type=Path), + help="Directory containing train.jsonl, val.jsonl, and test.jsonl.", +) +@click.option( + "--output-dir", + default="model-output", + type=click.Path(file_okay=False, path_type=Path), + help="Directory for checkpoints and the final model.", +) +@click.option("--model-name", default=MODEL_NAME, help="Base model to fine-tune.") +@click.option("--model-revision", default=None, help="Optional model revision to pin.") +@click.option( + "--max-length", + default=MAX_LENGTH, + type=click.IntRange(min=3, max=MAX_LENGTH), + show_default=True, +) +@click.option("--epochs", default=8, type=click.IntRange(min=1), show_default=True) +@click.option("--batch-size", default=1, type=int, show_default=True) +@click.option("--grad-accum", default=16, type=int, show_default=True) +@click.option("--base-lr", default=2e-5, type=float, show_default=True) +@click.option("--head-lr", default=1e-4, type=float, show_default=True) +@click.option("--aux-ce-weight", default=0.3, type=float, show_default=True) +@click.option( + "--optimizer", + type=click.Choice(["adamw", "adamw-8bit"]), + default="adamw", + show_default=True, +) +@click.option( + "--precision", + type=click.Choice(["fp32", "bf16"]), + default="fp32", + show_default=True, +) +@click.option("--no-crf", is_flag=True, default=False, help="Train without the CRF head.") +@click.option( + "--evaluate-test", + is_flag=True, + default=False, + help="Evaluate the test split after training.", +) +@click.option( + "--with-isr", + is_flag=True, + default=False, + help="Report injection success rate with final test evaluation.", +) +@click.option("--limit", default=0, type=int, help="Limit examples per split; zero uses all.") +@click.option("--resume", is_flag=True, default=False, help="Resume from the latest checkpoint.") +@click.option("--seed", default=42, type=int, show_default=True) +def main( + data_dir, + output_dir, + model_name, + model_revision, + max_length, + epochs, + batch_size, + grad_accum, + base_lr, + head_lr, + aux_ce_weight, + optimizer, + precision, + no_crf, + evaluate_test, + with_isr, + limit, + resume, + seed, +): + """Train the required phrase tagger from a BIOES dataset.""" + if with_isr and not evaluate_test: + raise click.UsageError("--with-isr requires --evaluate-test") + + config = Config( + data_dir=data_dir, + output_dir=output_dir, + model_name=model_name, + model_revision=model_revision, + max_length=max_length, + epochs=epochs, + batch_size=batch_size, + grad_accum=grad_accum, + base_lr=base_lr, + head_lr=head_lr, + aux_ce_weight=aux_ce_weight, + optimizer=optimizer, + precision=precision, + use_crf=not no_crf, + evaluate_test=evaluate_test, + with_isr=with_isr, + limit=limit, + resume=resume, + seed=seed, + ) + try: + run_training(config) + except ImportError as error: + raise click.ClickException( + f"{error}; install scancode-required-phrases[training]" + ) from error + + +if __name__ == "__main__": + main() diff --git a/tests/test_export.py b/tests/test_export.py new file mode 100644 index 0000000..174ebd4 --- /dev/null +++ b/tests/test_export.py @@ -0,0 +1,42 @@ +# -*- coding: utf-8 -*- +# +# Copyright (c) nexB Inc. and others. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +np = pytest.importorskip("numpy") + +from scancode_required_phrases.export import sha256 +from scancode_required_phrases.export import viterbi_decode + + +def test_viterbi_with_zero_transitions_is_argmax(): + emissions = np.array( + [ + [0.1, 0.9, 0.0, 0.0, 0.0], + [0.7, 0.2, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.9], + ] + ) + transitions = np.zeros((5, 5)) + edges = np.zeros(5) + + assert viterbi_decode(emissions, edges, transitions, edges) == [1, 0, 4] + + +def test_viterbi_obeys_transition_scores(): + emissions = np.array([[0.0, 1.0], [1.0, 0.0]]) + transitions = np.array([[0.0, 0.0], [-100.0, 0.0]]) + edges = np.zeros(2) + + path = viterbi_decode(emissions, edges, transitions, edges) + + assert path[0] == path[1] + + +def test_sha256_is_stable(tmp_path): + path = tmp_path / "model.onnx" + path.write_bytes(b"model") + + assert sha256(path) == "9372c470eeadd5ecd9c3c74c2b3cb633f8e2f2fad799250a0f70d652b6b825e4" diff --git a/tests/test_model.py b/tests/test_model.py new file mode 100644 index 0000000..b7ae65f --- /dev/null +++ b/tests/test_model.py @@ -0,0 +1,139 @@ +# -*- coding: utf-8 -*- +# +# Copyright (c) nexB Inc. and others. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +from types import SimpleNamespace + +import pytest + +os.environ.setdefault("USE_TF", "0") + +torch = pytest.importorskip("torch") +pytest.importorskip("torchcrf") +pytest.importorskip("transformers") + +from scancode_required_phrases import model as model_module +from scancode_required_phrases.model import build_optimizer +from scancode_required_phrases.model import PhraseTagger +from scancode_required_phrases.training import LABEL2ID +from scancode_required_phrases.training import LABELS + + +class FakeBackbone(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace( + hidden_size=4, + hidden_dropout_prob=0.1, + num_hidden_layers=1, + ) + self.encoder = torch.nn.Module() + self.encoder.layer = torch.nn.ModuleList([torch.nn.Linear(4, 4)]) + self.embeddings = torch.nn.Linear(4, 4) + + def gradient_checkpointing_enable(self): + pass + + def enable_input_require_grads(self): + pass + + +@pytest.fixture +def config(): + return SimpleNamespace( + model_name="fake-model", + model_revision="revision", + use_crf=True, + aux_ce_weight=0.3, + label_weights=[1.0] * len(LABELS), + optimizer="adamw", + base_lr=2e-5, + head_lr=1e-4, + layer_decay=0.98, + weight_decay=0.01, + adam_epsilon=1e-6, + ) + + +def test_class_weights_are_not_saved(monkeypatch, config): + monkeypatch.setattr( + model_module.AutoModel, + "from_pretrained", + lambda *args, **kwargs: FakeBackbone(), + ) + + tagger = PhraseTagger(config) + + assert tagger.class_weights is not None + assert "class_weights" not in tagger.state_dict() + + +def test_model_revision_is_passed_to_the_backbone(monkeypatch, config): + calls = [] + + def from_pretrained(*args, **kwargs): + calls.append((args, kwargs)) + return FakeBackbone() + + monkeypatch.setattr(model_module.AutoModel, "from_pretrained", from_pretrained) + PhraseTagger(config) + + assert calls == [(("fake-model",), {"revision": "revision"})] + + +def test_build_optimizer_uses_explicit_adamw(monkeypatch, config): + monkeypatch.setattr( + model_module.AutoModel, + "from_pretrained", + lambda *args, **kwargs: FakeBackbone(), + ) + tagger = PhraseTagger(config) + + optimizer = build_optimizer(config, tagger) + + assert isinstance(optimizer, torch.optim.AdamW) + learning_rates = {group["lr"] for group in optimizer.param_groups} + assert config.head_lr in learning_rates + assert any(rate < config.base_lr for rate in learning_rates) + + +def make_crf_tagger(): + from torchcrf import CRF + + tagger = PhraseTagger.__new__(PhraseTagger) + torch.nn.Module.__init__(tagger) + tagger.use_crf = True + tagger.num_labels = len(LABELS) + tagger.crf = CRF(len(LABELS), batch_first=True) + with torch.no_grad(): + for parameter in tagger.crf.parameters(): + parameter.zero_() + return tagger + + +def test_predict_words_uses_first_subwords(): + tagger = make_crf_tagger() + emissions = torch.zeros((1, 5, len(LABELS))) + emissions[0, 1, LABEL2ID["B-REQ"]] = 9.0 + emissions[0, 2, LABEL2ID["E-REQ"]] = 9.0 + emissions[0, 3, LABEL2ID["S-REQ"]] = 9.0 + tagger.emissions = lambda input_ids, attention_mask: emissions + + input_ids = torch.zeros((1, 5), dtype=torch.long) + tags = tagger.predict_words( + input_ids, + input_ids, + [None, 0, 1, 1, None], + ) + + assert tags == [LABEL2ID["B-REQ"], LABEL2ID["E-REQ"]] + + +def test_predict_words_does_not_run_for_an_empty_sequence(): + tagger = make_crf_tagger() + tagger.emissions = lambda *args: pytest.fail("emissions should not be computed") + input_ids = torch.zeros((1, 2), dtype=torch.long) + + assert tagger.predict_words(input_ids, input_ids, [None, None]) == [] diff --git a/tests/test_training.py b/tests/test_training.py new file mode 100644 index 0000000..ceb9801 --- /dev/null +++ b/tests/test_training.py @@ -0,0 +1,332 @@ +# -*- coding: utf-8 -*- +# +# Copyright (c) nexB Inc. and others. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +from click.testing import CliRunner +import pytest + +from scancode_required_phrases.training import align_labels +from scancode_required_phrases.training import compute_metrics +from scancode_required_phrases.training import Config +from scancode_required_phrases.training import decode_row +from scancode_required_phrases.training import extract_spans +from scancode_required_phrases.training import first_subword_positions +from scancode_required_phrases.training import IGNORE_INDEX +from scancode_required_phrases.training import LABEL2ID +from scancode_required_phrases.training import main +from scancode_required_phrases.training import PhraseDataset +from scancode_required_phrases.training import prepare_output_dir +from scancode_required_phrases.training import serializable_config +from scancode_required_phrases.training import sha256 +from scancode_required_phrases.training import validate_bioes +from scancode_required_phrases.training import validate_config +from scancode_required_phrases.training import validate_record +from scancode_required_phrases.training import validate_saved_state +from scancode_required_phrases.training import validate_splits + + +class FakeEncoding(dict): + def __init__(self, word_ids): + super().__init__() + self._word_ids = word_ids + self["input_ids"] = [0] * len(word_ids) + self["attention_mask"] = [1] * len(word_ids) + + def word_ids(self): + return self._word_ids + + +class FakeTokenizer: + def __call__(self, tokens, max_length=512, **kwargs): + word_ids = [None, *range(len(tokens)), None] + return FakeEncoding(word_ids[:max_length]) + + +def make_record(identifier="mit_1.RULE", tokens=None, labels=None): + return { + "identifier": identifier, + "license_expression": "mit", + "rule_type": "is_license_notice", + "text": "MIT License terms apply", + "tokens": tokens or ["MIT", "License", "terms", "apply"], + "bioes_labels": labels or ["B-REQ", "E-REQ", "O", "O"], + } + + +def write_jsonl(path, records): + path.write_text( + "".join(json.dumps(record) + "\n" for record in records), + encoding="utf-8", + ) + + +@pytest.mark.parametrize( + "labels", + [ + ["O"], + ["S-REQ"], + ["B-REQ", "E-REQ"], + ["B-REQ", "I-REQ", "E-REQ", "O", "S-REQ"], + ], +) +def test_validate_bioes_accepts_valid_sequences(labels): + assert validate_bioes(labels) is None + + +@pytest.mark.parametrize( + "labels", + [ + [], + ["I-REQ"], + ["E-REQ"], + ["B-REQ", "O"], + ["O", "I-REQ"], + ["B-REQ", "I-REQ"], + ], +) +def test_validate_bioes_rejects_invalid_sequences(labels): + assert validate_bioes(labels) + + +def test_validate_record_accepts_a_complete_record(): + record = make_record() + assert validate_record(record, "train.jsonl", 1) is record + + +@pytest.mark.parametrize( + "field_name", + ["identifier", "license_expression", "rule_type", "text", "tokens", "bioes_labels"], +) +def test_validate_record_rejects_a_missing_field(field_name): + record = make_record() + del record[field_name] + with pytest.raises(ValueError, match=field_name): + validate_record(record, "train.jsonl", 4) + + +def test_validate_record_rejects_mismatched_tokens_and_labels(): + record = make_record(labels=["S-REQ"]) + with pytest.raises(ValueError, match="tokens and"): + validate_record(record, "train.jsonl", 2) + + +def test_validate_record_rejects_unknown_labels(): + record = make_record(labels=["B-REQ", "BAD", "O", "O"]) + with pytest.raises(ValueError, match="unknown labels"): + validate_record(record, "train.jsonl", 2) + + +def test_align_labels_keeps_only_first_subwords(): + tokenizer = lambda tokens, **kwargs: FakeEncoding([None, 0, 1, 1, None]) + encoding, truncated, cut_phrase = align_labels( + ["MIT", "License"], + ["B-REQ", "E-REQ"], + tokenizer, + 512, + ) + assert encoding["labels"] == [ + IGNORE_INDEX, + LABEL2ID["B-REQ"], + LABEL2ID["E-REQ"], + IGNORE_INDEX, + IGNORE_INDEX, + ] + assert not truncated + assert not cut_phrase + + +def test_align_labels_detects_a_phrase_cut_by_truncation(): + encoding, truncated, cut_phrase = align_labels( + ["prefix", "GNU", "General", "Public", "License"], + ["O", "B-REQ", "I-REQ", "I-REQ", "E-REQ"], + FakeTokenizer(), + 4, + ) + assert encoding["labels"][-1] == LABEL2ID["I-REQ"] + assert truncated + assert cut_phrase + + +def test_align_labels_allows_safe_truncation(): + _, truncated, cut_phrase = align_labels( + ["MIT", "License", "terms", "apply"], + ["B-REQ", "E-REQ", "O", "O"], + FakeTokenizer(), + 4, + ) + assert truncated + assert not cut_phrase + + +def test_phrase_dataset_skips_a_cut_phrase(tmp_path): + path = tmp_path / "train.jsonl" + write_jsonl( + path, + [ + make_record(identifier="safe.RULE"), + make_record( + identifier="cut.RULE", + tokens=["prefix", "GNU", "General", "Public", "License"], + labels=["O", "B-REQ", "I-REQ", "I-REQ", "E-REQ"], + ), + ], + ) + + dataset = PhraseDataset(path, FakeTokenizer(), max_length=4) + + assert dataset.identifiers == ["safe.RULE"] + assert dataset.truncated == 2 + assert dataset.cut_phrases == 1 + + +def test_validate_splits_with_real_dataset_objects(tmp_path): + paths = {} + for name, identifier in (("train", "a.RULE"), ("validation", "b.RULE")): + path = tmp_path / f"{name}.jsonl" + write_jsonl(path, [make_record(identifier=identifier)]) + paths[name] = PhraseDataset(path, FakeTokenizer(), 512) + validate_splits(paths) + + +def test_validate_splits_rejects_duplicate_identifiers(tmp_path): + datasets = {} + for name in ("train", "validation"): + path = tmp_path / f"{name}.jsonl" + write_jsonl(path, [make_record(identifier="same.RULE")]) + datasets[name] = PhraseDataset(path, FakeTokenizer(), 512) + with pytest.raises(ValueError, match="both train and validation"): + validate_splits(datasets) + + +@pytest.mark.parametrize( + "tags, expected", + [ + (["O", "B-REQ", "I-REQ", "E-REQ", "O"], {(1, 3)}), + (["O", "S-REQ", "O"], {(1, 1)}), + (["S-REQ", "O", "B-REQ", "E-REQ"], {(0, 0), (2, 3)}), + (["O", "O"], set()), + ], +) +def test_extract_spans(tags, expected): + assert extract_spans(tags) == expected + + +def test_decode_row_drops_ignored_positions(): + predicted, actual = decode_row( + [LABEL2ID["B-REQ"], 0, LABEL2ID["E-REQ"]], + [LABEL2ID["B-REQ"], IGNORE_INDEX, LABEL2ID["E-REQ"]], + ) + assert predicted == ["B-REQ", "E-REQ"] + assert actual == ["B-REQ", "E-REQ"] + + +def test_compute_metrics_scores_exact_spans(): + predictions = [[LABEL2ID["B-REQ"], LABEL2ID["E-REQ"], LABEL2ID["O"]]] + labels = [[LABEL2ID["B-REQ"], LABEL2ID["E-REQ"], LABEL2ID["O"]]] + scores = compute_metrics((predictions, labels)) + assert scores["f1"] == 1.0 + assert scores["exact_match"] == 1.0 + assert scores["predicted_spans"] == scores["gold_spans"] == 1 + + +def test_first_subword_positions_skips_specials_and_continuations(): + assert first_subword_positions([None, 0, 1, 1, 2, None]) == [1, 2, 4] + + +def test_sha256_is_stable(tmp_path): + path = tmp_path / "data.jsonl" + path.write_bytes(b"required phrase\n") + assert sha256(path) == "792616e2062f96efb6ae2f69e8637b834e2db74354eb4f51e78eda329038cc70" + + +def test_serializable_config_converts_paths(tmp_path): + config = Config(data_dir=tmp_path / "data", output_dir=tmp_path / "model") + values = serializable_config(config) + assert values["data_dir"] == str(tmp_path / "data") + assert values["output_dir"] == str(tmp_path / "model") + + +def test_validate_config_accepts_the_default_training_settings(tmp_path): + config = Config(data_dir=tmp_path / "data", output_dir=tmp_path / "model") + validate_config(config) + + +def test_validate_config_rejects_invalid_settings(tmp_path): + config = Config( + data_dir=tmp_path / "data", + output_dir=tmp_path / "model", + aux_ce_weight=-1, + ) + with pytest.raises(ValueError, match="cannot be negative"): + validate_config(config) + + +def test_prepare_output_dir_refuses_to_overwrite_a_run(tmp_path): + output_dir = tmp_path / "model" + output_dir.mkdir() + (output_dir / "model.safetensors").write_bytes(b"model") + + with pytest.raises(ValueError, match="not empty"): + prepare_output_dir(output_dir, resume=False) + + +def test_prepare_output_dir_requires_a_checkpoint_to_resume(tmp_path): + with pytest.raises(ValueError, match="No checkpoint"): + prepare_output_dir(tmp_path / "model", resume=True) + + +def test_validate_saved_state_accepts_matching_tensors(tmp_path): + torch = pytest.importorskip("torch") + safetensors = pytest.importorskip("safetensors.torch") + + model = torch.nn.Linear(2, 1) + path = tmp_path / "model.safetensors" + safetensors.save_file(model.state_dict(), str(path)) + validate_saved_state(model, path) + + +def test_validate_saved_state_rejects_unexpected_keys(tmp_path): + torch = pytest.importorskip("torch") + safetensors = pytest.importorskip("safetensors.torch") + + model = torch.nn.Linear(2, 1) + state = dict(model.state_dict()) + state["unexpected"] = torch.ones(1) + path = tmp_path / "model.safetensors" + safetensors.save_file(state, str(path)) + with pytest.raises(ValueError, match="key mismatch"): + validate_saved_state(model, path) + + +def test_validate_saved_state_rejects_wrong_shapes(tmp_path): + torch = pytest.importorskip("torch") + safetensors = pytest.importorskip("safetensors.torch") + + model = torch.nn.Linear(2, 1) + state = dict(model.state_dict()) + state["weight"] = torch.ones((1, 3)) + path = tmp_path / "model.safetensors" + safetensors.save_file(state, str(path)) + with pytest.raises(ValueError, match="has shape"): + validate_saved_state(model, path) + + +def test_validate_saved_state_rejects_non_finite_values(tmp_path): + torch = pytest.importorskip("torch") + safetensors = pytest.importorskip("safetensors.torch") + + model = torch.nn.Linear(2, 1) + state = dict(model.state_dict()) + state["weight"] = torch.full_like(state["weight"], float("nan")) + path = tmp_path / "model.safetensors" + safetensors.save_file(state, str(path)) + with pytest.raises(ValueError, match="non-finite"): + validate_saved_state(model, path) + + +def test_cli_requires_test_evaluation_for_isr(tmp_path): + result = CliRunner().invoke(main, ["--data-dir", str(tmp_path), "--with-isr"]) + assert result.exit_code == 2 + assert "--with-isr requires --evaluate-test" in result.output From 7b45ad16226d47ee1496942c4737bb12351d670a Mon Sep 17 00:00:00 2001 From: Kaushik Date: Mon, 7 Sep 2026 12:23:12 +0000 Subject: [PATCH 2/5] Harden required phrase model training Signed-off-by: Kaushik --- README.rst | 4 +- docs/source/training.rst | 97 +- setup.cfg | 5 +- src/scancode_required_phrases/export.py | 666 +++++-- src/scancode_required_phrases/model.py | 241 ++- src/scancode_required_phrases/training.py | 1916 +++++++++++++++++---- tests/test_export.py | 474 ++++- tests/test_model.py | 241 ++- tests/test_training.py | 1224 +++++++++++-- 9 files changed, 4141 insertions(+), 727 deletions(-) diff --git a/README.rst b/README.rst index 3aadee4..2e09892 100644 --- a/README.rst +++ b/README.rst @@ -44,10 +44,12 @@ Install the optional training dependencies and train from a generated dataset: python -m pip install ".[training,training-8bit]" train-required-phrase-model --data-dir dataset-output --output-dir model-output \ + --model-revision 64a8c8eab3e352a784c658aef62be1662607476f \ --optimizer adamw-8bit Test-set evaluation is opt-in and should only be used for a final selected run. -See ``docs/source/training.rst`` for training and ONNX export details. +The completed model is written to ``model-output/final-model`` only after local +reload and validation succeed. See ``docs/source/training.rst`` for details. Development =========== diff --git a/docs/source/training.rst b/docs/source/training.rst index 4594e95..dd41f8f 100644 --- a/docs/source/training.rst +++ b/docs/source/training.rst @@ -1,59 +1,82 @@ Train and export a required phrase model ========================================= -The training command fine-tunes a DeBERTa token classifier with BIOES labels -and an optional word-level CRF. It reads the JSONL files produced by -``build-required-phrases-dataset``. +The training command fine-tunes the reviewed DeBERTa word-input classifier with +BIOES labels and an optional constrained word-level CRF. It consumes the hybrid +JSONL splits produced by ``build-required-phrases-dataset`` without moving, +deduplicating, retokenizing, relabeling, or repairing records. Installation ------------ -Install the training dependencies before running the command. The 8-bit -optimizer is a separate extra because it is only needed on constrained GPUs. +Install training support with ``python -m pip install ".[training]"``. Add the +``training-8bit`` extra only when using ``adamw-8bit``. ONNX support remains +optional and is installed separately with ``.[onnx]``. The training extra pins +Transformers 4.57.3 and Hugging Face Hub 0.36.2, the versions validated by the +single-T4 Kaggle smoke run. -.. code-block:: console - - python -m pip install ".[training,training-8bit]" - -Install the ONNX dependencies when an export is needed: - -.. code-block:: console - - python -m pip install ".[onnx]" +Training contract +----------------- -Training --------- +A run requires a full immutable model-host commit in ``--model-revision`` and a +new or empty ``--output-dir``. For the current DeBERTa base model, run: .. code-block:: console - train-required-phrase-model \ - --data-dir dataset-output \ - --output-dir model-output \ + train-required-phrase-model --data-dir dataset-output --output-dir model-output \ + --model-revision 64a8c8eab3e352a784c658aef62be1662607476f \ --optimizer adamw-8bit -The default configuration uses FP32 and standard AdamW. The 8-bit optimizer is -recommended for DeBERTa-large on a 16 GB GPU. Test-set evaluation is disabled by -default so repeated experiments select models using validation data only. Add -``--evaluate-test`` only for the final selected run. - -The command validates the dataset, skips examples whose required phrase is cut -by tokenizer truncation, and saves the checkpoint with the best validation -strict span F1. The output directory must be empty unless ``--resume`` is used. -It writes ``train_config.json`` for inference compatibility and -``run_manifest.json`` for reproducibility. +Resume is unsupported. Every non-empty line in +all three splits is validated before tokenizer loading or ``--limit`` selection. +Records must use exact field types, contain a positive valid BIOES sequence, and +have unique identifiers across all splits. Exact token and label duplicates are +reported but never removed. + +Alignment uses the dataset's existing words. Full untruncated fast-tokenizer +coverage proves the complete retained word prefix. Coverage gaps, zero-subword +words, partial words, and truncation that omits any non-``O`` label reject the +record. ``--limit`` is applied only after complete validation and alignment. + +The run records versioned raw-byte (H0), validated-record (H1), and effective +example (H2) hashes. ``run_manifest.json`` is atomically replaced through +pre-run, failure, or success state and contains source/runtime provenance from +runtime APIs, never an environment dump. + +Final model +----------- + +Selection remains based only on best validation strict span F1. The selected +checkpoint, selected in-memory model, and staged model must have exactly equal +states. ``final-model.tmp`` is then reloaded using only its local configuration, +tokenizer, and full weights; ordered validation predictions, labels, invalid +path count, and metrics must match exactly. The stage is atomically promoted to +``final-model`` and ``SUCCESS.json`` is written last. A model is publishable +only while that marker and every recorded file hash validate. + +Strict metrics reject malformed gold paths and malformed CRF predictions. +Malformed non-CRF predictions are not repaired: they produce no predicted spans, +count all gold spans as false negatives, and increment ``invalid_paths``. + +ISR retains its existing name but measures predicted-phrase locatability only. +It does not exercise injection gates or rule mutation and is not evidence of +injection success. Export ------ +``export-required-phrase-model`` validates a publishable local final model and +exports constrained CRF matrices by default: + .. code-block:: console - export-required-phrase-model \ - --model-dir model-output \ - --output-dir onnx-output + export-required-phrase-model --model-dir model-output/final-model \ + --output-dir model-export -The export command strictly loads the saved model, exports the DeBERTa emissions -to ONNX, stores the CRF transitions separately, and verifies that NumPy and -PyTorch Viterbi decoding agree. +Run a separate export with ``--operation onnx`` and a different new output +directory for optional ONNX emissions; missing ONNX packages cannot disable CRF +export, training, finalization, or offline reload. -Generated datasets, checkpoints, model files, and export artifacts must not be -committed to this repository. +Generated datasets, reports, manifests, checkpoints, final models, matrix +files, ONNX files, smoke outputs, and caches must remain outside commits and the +pull request. diff --git a/setup.cfg b/setup.cfg index 5045e8a..14d02c4 100644 --- a/setup.cfg +++ b/setup.cfg @@ -63,12 +63,13 @@ console_scripts = [options.extras_require] training = accelerate >= 0.33 + huggingface-hub == 0.36.2 protobuf >= 3.20 - pytorch-crf >= 0.7.2 + pytorch-crf == 0.7.2 safetensors >= 0.4 sentencepiece >= 0.2 torch >= 2.0 - transformers >= 4.44 + transformers == 4.57.3 training-8bit = bitsandbytes >= 0.43 onnx = diff --git a/src/scancode_required_phrases/export.py b/src/scancode_required_phrases/export.py index 17a55c8..504c40d 100644 --- a/src/scancode_required_phrases/export.py +++ b/src/scancode_required_phrases/export.py @@ -3,32 +3,103 @@ # Copyright (c) nexB Inc. and others. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Export a trained required phrase tagger for CPU inference.""" +"""Export a completed required phrase model for CPU inference.""" import hashlib import json import os from pathlib import Path -from types import SimpleNamespace +import shutil import click os.environ.setdefault("USE_TF", "0") +EXPORT_SCHEMA = "scancode-required-phrases-export-v1" +PARITY_CASES = ( + "one-word-tie", + "normal", + "variable-length", + "boundary", + "extreme-forbidden", +) -def viterbi_decode(emissions, start_transitions, transitions, end_transitions): - """Return the best tag path for one sequence.""" - sequence_length = emissions.shape[0] - score = start_transitions + emissions[0] - backpointers = [] +class OnnxDependencyError(RuntimeError): + """Raised when an explicitly requested ONNX export cannot run.""" + + +def _require_float_array(name, value, shape=None): + """Return a validated floating-point NumPy array.""" + import numpy as np + + if not isinstance(value, np.ndarray): + raise TypeError(f"{name} must be a NumPy array") + if not np.issubdtype(value.dtype, np.floating): + raise TypeError(f"{name} must have a floating-point dtype") + if shape is not None and value.shape != shape: + raise ValueError(f"{name} must have shape {shape}, not {value.shape}") + if np.isnan(value).any(): + raise ValueError(f"{name} contains NaN scores") + return value + + +def viterbi_decode( + emissions, + start_transitions, + transitions, + end_transitions, + num_labels=None, +): + """Return the best tag path for one validated, non-empty sequence.""" + import numpy as np + + emissions = _require_float_array("emissions", emissions) + if emissions.ndim != 2: + raise ValueError("emissions must have shape (sequence_length, labels)") + sequence_length, inferred_labels = emissions.shape + if sequence_length == 0 or inferred_labels == 0: + raise ValueError("emissions must contain at least one token and label") + if num_labels is not None: + if isinstance(num_labels, bool) or not isinstance(num_labels, int): + raise TypeError("num_labels must be an integer") + if num_labels <= 0 or num_labels != inferred_labels: + raise ValueError(f"emissions have {inferred_labels} labels, expected {num_labels}") + + start_transitions = _require_float_array( + "start_transitions", start_transitions, (inferred_labels,) + ) + transitions = _require_float_array( + "transitions", transitions, (inferred_labels, inferred_labels) + ) + end_transitions = _require_float_array("end_transitions", end_transitions, (inferred_labels,)) + + with np.errstate(invalid="ignore"): + score = start_transitions + emissions[0] + if np.isnan(score).any(): + raise ValueError("start and emission scores produce NaN") + if np.isneginf(score).all(): + raise ValueError("no valid Viterbi path at token 0") + + backpointers = [] for step in range(1, sequence_length): - candidates = score[:, None] + transitions - best_source = candidates.argmax(axis=0) - score = candidates.max(axis=0) + emissions[step] + with np.errstate(invalid="ignore"): + candidates = score[:, None] + transitions + best_source = candidates.argmax(axis=0) + score = candidates.max(axis=0) + emissions[step] + if np.isnan(score).any(): + raise ValueError(f"scores produce NaN at token {step}") + if np.isneginf(score).all(): + raise ValueError(f"no valid Viterbi path at token {step}") backpointers.append(best_source) - score = score + end_transitions + with np.errstate(invalid="ignore"): + score = score + end_transitions + if np.isnan(score).any(): + raise ValueError("end and path scores produce NaN") + if np.isneginf(score).all(): + raise ValueError("no valid Viterbi path reaches an end label") + best = int(score.argmax()) path = [best] for sources in reversed(backpointers): @@ -47,162 +118,443 @@ def sha256(path): return digest.hexdigest() -def build_emissions_module(tagger): - """Wrap the trained backbone and classifier for ONNX export.""" - import torch.nn as nn - - class EmissionsModule(nn.Module): - def __init__(self): - super().__init__() - self.backbone = tagger.backbone - self.classifier = tagger.classifier +def _training_interfaces(): + """Import final-artifact interfaces without importing optional ONNX code.""" + from scancode_required_phrases.training import ARTIFACT_SCHEMA + from scancode_required_phrases.training import CONSTRAINT_CONTRACT + from scancode_required_phrases.training import LABELS + from scancode_required_phrases.training import load_final_model + from scancode_required_phrases.training import validate_bioes + from scancode_required_phrases.training import validate_publishable_model + from scancode_required_phrases.training import write_json_atomic + + return ( + ARTIFACT_SCHEMA, + CONSTRAINT_CONTRACT, + tuple(LABELS), + load_final_model, + validate_bioes, + validate_publishable_model, + write_json_atomic, + ) - def forward(self, input_ids, attention_mask): - hidden = self.backbone( - input_ids=input_ids, - attention_mask=attention_mask, - ).last_hidden_state - return self.classifier(hidden) - return EmissionsModule().eval() +def _load_publishable_artifact(model_dir, require_crf=False): + """Validate and load one supported Final_Model entirely from local files.""" + ( + artifact_schema, + constraint_contract, + labels, + load_final_model, + _validate_bioes, + validate_publishable_model, + _write_json_atomic, + ) = _training_interfaces() + model_dir = Path(model_dir) + validate_publishable_model(model_dir) -def load_tagger(model_dir, train_config): - """Rebuild a tagger and strictly load its saved weights.""" + config_path = model_dir / "train_config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + if config.get("artifact_schema") != artifact_schema: + raise ValueError("Final_Model uses an unsupported artifact schema") + if tuple(config.get("labels", ())) != labels: + raise ValueError("Final_Model labels do not match the supported label order") + if config.get("constraint_contract") != constraint_contract: + raise ValueError("Final_Model uses an unsupported constraint contract") + if type(require_crf) is not bool: + raise TypeError("require_crf must be a boolean") + + loaded = load_final_model(model_dir, offline=True) + if not isinstance(loaded, tuple) or len(loaded) < 2: + raise TypeError("load_final_model must return (model, tokenizer)") + tagger, tokenizer = loaded[:2] + if require_crf and ( + config.get("use_crf") is not True + or not getattr(tagger, "use_crf", False) + or not hasattr(tagger, "crf") + ): + raise ValueError("CRF matrix export requires a Final_Model with a CRF") + + marker_path = model_dir / "SUCCESS.json" + marker = json.loads(marker_path.read_text(encoding="utf-8")) + return tagger.eval(), tokenizer, config, marker + + +def _effective_crf_matrices(tagger): + """Return effective constrained CRF matrices without mutating learned values.""" + import numpy as np import torch - from safetensors.torch import load_file - - from scancode_required_phrases.model import PhraseTagger - config = SimpleNamespace( - model_name=train_config["model_name"], - model_revision=train_config.get("model_revision"), - use_crf=train_config["use_crf"], - aux_ce_weight=0.0, - label_weights=[1.0] * len(train_config["labels"]), + crf = tagger.crf + values = ( + ( + "start_transitions", + "start_mask", + "effective_start_transitions", + (tagger.num_labels,), + ), + ( + "transitions", + "transition_mask", + "effective_transitions", + (tagger.num_labels, tagger.num_labels), + ), + ( + "end_transitions", + "end_mask", + "effective_end_transitions", + (tagger.num_labels,), + ), ) - tagger = PhraseTagger(config) - - model_dir = Path(model_dir) - safetensors_file = model_dir / "model.safetensors" - pytorch_file = model_dir / "pytorch_model.bin" - if safetensors_file.exists(): - state = load_file(str(safetensors_file)) - elif pytorch_file.exists(): - state = torch.load(pytorch_file, map_location="cpu", weights_only=True) - else: - raise FileNotFoundError(f"No model weights found in {model_dir}") - - # Checkpoints created before class weights became non-persistent contain - # this training-only tensor. - state.pop("class_weights", None) - for name, tensor in state.items(): - if not torch.isfinite(tensor).all(): - raise ValueError(f"Checkpoint tensor {name!r} contains non-finite values") - - tagger.load_state_dict(state, strict=True) - return tagger.eval() - + effective = [] + for parameter_name, mask_name, effective_name, shape in values: + parameter = getattr(crf, parameter_name, None) + mask = getattr(crf, mask_name, None) + constrained = getattr(crf, effective_name, None) + if parameter is None or mask is None or constrained is None: + raise ValueError(f"CRF is missing {parameter_name}, {mask_name}, or {effective_name}") + if ( + tuple(parameter.shape) != shape + or tuple(mask.shape) != shape + or tuple(constrained.shape) != shape + ): + raise ValueError(f"CRF effective {parameter_name} shapes are invalid") + if mask.dtype != torch.bool: + raise TypeError(f"CRF {mask_name} must be boolean") + if not torch.isfinite(parameter).all(): + raise ValueError(f"CRF {parameter_name} contains non-finite learned values") + expected = parameter.masked_fill(~mask, -torch.inf) + if not torch.equal(constrained, expected): + raise ValueError(f"CRF {effective_name} does not apply the supported masks") + array = constrained.detach().cpu().numpy().copy() + if np.isnan(array).any(): + raise ValueError(f"effective CRF {parameter_name} contains NaN") + effective.append(array) + return tuple(effective) + + +def _parity_batches(num_labels): + """Return deterministic normal, variable, boundary, and adversarial batches.""" + import numpy as np -def check_viterbi_matches_crf(tagger, num_tags): - """Verify NumPy and pytorch-crf return identical paths.""" - import torch + normal = np.zeros((3, num_labels), dtype=np.float32) + normal[0, 1] = 4.0 + normal[1, 2] = 4.0 + normal[2, 3] = 4.0 + + variable = np.arange(6 * num_labels, dtype=np.float32).reshape(6, num_labels) + variable = (variable % 7.0) - 3.0 + variable_batch = np.stack((np.roll(variable, 1, axis=1), variable)) + + boundary = np.zeros((4, num_labels), dtype=np.float32) + boundary[0, 2] = 1_000.0 + boundary[1, 3] = 1_000.0 + boundary[-1, 1] = 1_000.0 + + adversarial = np.full((5, num_labels), -10_000.0, dtype=np.float32) + adversarial[:, 2] = 10_000.0 + adversarial[0, 3] = 20_000.0 + adversarial[-1, 1] = 20_000.0 + + return ( + (PARITY_CASES[1], normal[None, :, :], (3,)), + (PARITY_CASES[2], variable_batch, (2, 6)), + (PARITY_CASES[3], boundary[None, :, :], (4,)), + (PARITY_CASES[4], adversarial[None, :, :], (5,)), + ) - start = tagger.crf.start_transitions.detach().cpu().numpy() - transitions = tagger.crf.transitions.detach().cpu().numpy() - end = tagger.crf.end_transitions.detach().cpu().numpy() - emissions = torch.randn(3, 14, num_tags) - mask = torch.ones(3, 14, dtype=torch.bool) - crf_paths = tagger.crf.decode(emissions, mask=mask) - for row in range(emissions.size(0)): - numpy_path = viterbi_decode(emissions[row].numpy(), start, transitions, end) - if numpy_path != crf_paths[row]: - raise AssertionError("NumPy Viterbi disagrees with pytorch-crf decoding") +def _exact_one_word_tie(start, end): + """Return emissions that produce an exact tie under exported edge scores.""" + import numpy as np - return start, transitions, end + legal = np.flatnonzero(np.isfinite(start) & np.isfinite(end)) + if legal.size < 2: + raise ValueError("CRF constraints must permit at least two one-word paths") + + candidates_by_label = {} + for label in legal: + base = np.float32(-(np.float64(start[label]) + np.float64(end[label]))) + candidates = {} + lower = upper = base + for _step in range(16_385): + for emission in (lower, upper): + score = np.float32(np.float32(start[label] + emission) + end[label]) + candidates.setdefault(score.tobytes(), (emission, score)) + lower = np.float32(np.nextafter(lower, np.float32(-np.inf))) + upper = np.float32(np.nextafter(upper, np.float32(np.inf))) + candidates_by_label[int(label)] = candidates + + for left_index, left in enumerate(legal[:-1]): + left = int(left) + for right_value in legal[left_index + 1 :]: + right = int(right_value) + shared = set(candidates_by_label[left]).intersection(candidates_by_label[right]) + if not shared: + continue + score_key = min(shared) + left_emission, tied_score = candidates_by_label[left][score_key] + right_emission, right_score = candidates_by_label[right][score_key] + if tied_score.tobytes() != right_score.tobytes(): + raise AssertionError("Constructed tie scores are not exactly equal") + emissions = np.zeros((1, len(start)), dtype=np.float32) + for label in legal: + label = int(label) + baseline = -( + np.float64(start[label]) + np.float64(end[label]) + 10_000.0 + ) + emissions[0, label] = np.float32(baseline) + emissions[0, left] = left_emission + emissions[0, right] = right_emission + return emissions, tied_score + raise AssertionError("Could not construct an exact tie from exported CRF edge scores") -def export(model_dir, output_dir, opset): - """Export ONNX emissions, CRF transitions, and a checksum manifest.""" +def check_viterbi_matches_crf(tagger, num_tags): + """Verify exact NumPy/PyTorch constrained paths on deterministic cases.""" import numpy as np import torch - from transformers import AutoTokenizer - model_dir = Path(model_dir) - output_dir = Path(output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - - config_path = model_dir / "train_config.json" - if not config_path.exists(): - raise FileNotFoundError(f"No train_config.json found in {model_dir}") - train_config = json.loads(config_path.read_text(encoding="utf-8")) - labels = train_config["labels"] - use_crf = train_config["use_crf"] - - tagger = load_tagger(model_dir, train_config) - emissions_module = build_emissions_module(tagger) - tokenizer = AutoTokenizer.from_pretrained(str(model_dir), use_fast=True) - - sample = tokenizer( - "Licensed under the Apache License Version 2.0", - return_tensors="pt", - ) - inputs = sample["input_ids"], sample["attention_mask"] - - onnx_path = output_dir / "model.onnx" - torch.onnx.export( - emissions_module, - inputs, - str(onnx_path), - input_names=["input_ids", "attention_mask"], - output_names=["emissions"], - dynamic_axes={ - "input_ids": {0: "batch", 1: "sequence"}, - "attention_mask": {0: "batch", 1: "sequence"}, - "emissions": {0: "batch", 1: "sequence"}, - }, - opset_version=opset, - do_constant_folding=True, + ( + _artifact_schema, + _constraint_contract, + labels, + _load_final_model, + validate_bioes, + _validate_publishable_model, + _write_json_atomic, + ) = _training_interfaces() + if isinstance(num_tags, bool) or not isinstance(num_tags, int): + raise TypeError("num_tags must be an integer") + if num_tags != len(labels) or num_tags != tagger.num_labels: + raise ValueError("Parity label count does not match the supported labels") + + start, transitions, end = _effective_crf_matrices(tagger) + tie_emissions, tied_score = _exact_one_word_tie(start, end) + terminal_scores = np.float32(np.float32(start + tie_emissions[0]) + end) + if np.count_nonzero(terminal_scores == tied_score) < 2: + raise AssertionError("Exported matrices did not produce the constructed exact tie") + tie_tensor = torch.from_numpy(tie_emissions).unsqueeze(0) + tie_mask = torch.ones((1, 1), dtype=torch.bool) + tie_path = tagger.crf.decode(tie_tensor, mask=tie_mask)[0] + numpy_tie_path = viterbi_decode( + tie_emissions, + start, + transitions, + end, + num_labels=num_tags, ) + if numpy_tie_path != tie_path: + raise AssertionError( + f"NumPy Viterbi disagrees with PyTorch for exact tie: " + f"{numpy_tie_path} != {tie_path}" + ) + tie_error = validate_bioes([labels[tag] for tag in tie_path]) + if tie_error: + raise AssertionError(f"Exact tie parity decoded {tie_error}") + + start, transitions, end = _effective_crf_matrices(tagger) + for case_name, emissions, lengths in _parity_batches(num_tags): + tensor = torch.from_numpy(emissions) + mask = torch.arange(tensor.shape[1]).unsqueeze(0) < torch.tensor(lengths).unsqueeze(1) + pytorch_paths = tagger.crf.decode(tensor, mask=mask) + if len(pytorch_paths) != len(lengths): + raise AssertionError(f"PyTorch returned an invalid batch for {case_name}") + for row, length in enumerate(lengths): + numpy_path = viterbi_decode( + emissions[row, :length], + start, + transitions, + end, + num_labels=num_tags, + ) + if numpy_path != pytorch_paths[row]: + raise AssertionError( + f"NumPy Viterbi disagrees with PyTorch for {case_name} row {row}: " + f"{numpy_path} != {pytorch_paths[row]}" + ) + path_error = validate_bioes([labels[tag] for tag in numpy_path]) + if path_error: + raise AssertionError(f"Parity case {case_name} row {row} decoded {path_error}") + return start, transitions, end + +def _export_manifest(marker, matrix_path, onnx_path=None): + """Return the export manifest for one validated Final_Model.""" + artifact_schema, constraint_contract, labels, *_unused = _training_interfaces() manifest = { - "labels": labels, - "onnx_model": sha256(onnx_path), + "schema": EXPORT_SCHEMA, + "artifact_schema": artifact_schema, + "constraint_contract": constraint_contract, + "labels": list(labels), + "source_final_model": marker, + "crf_transitions": sha256(matrix_path), + "parity": {"passed": True, "cases": list(PARITY_CASES)}, } + if onnx_path is not None: + manifest["onnx_model"] = sha256(onnx_path) + return manifest + + +def _prepare_export_stage(model_dir, output_dir): + """Return a new staging directory and absent final export destination.""" + model_dir = Path(model_dir).resolve() + output_dir = Path(output_dir).resolve() + if output_dir == model_dir or model_dir in output_dir.parents: + raise ValueError("Export output must be outside the Final_Model directory") + if output_dir.exists(): + if not output_dir.is_dir(): + raise ValueError(f"Export output is not a directory: {output_dir}") + if any(output_dir.iterdir()): + raise ValueError(f"Export output directory is not empty: {output_dir}") + output_dir.rmdir() + stage = output_dir.with_name(f"{output_dir.name}.tmp") + if stage.exists(): + raise ValueError(f"Export staging directory already exists: {stage}") + stage.parent.mkdir(parents=True, exist_ok=True) + stage.mkdir() + return stage, output_dir + + +def _promote_export(stage, output_dir): + """Atomically publish one completely verified export directory.""" + os.replace(stage, output_dir) + + +def export_crf_matrices(model_dir, output_dir): + """Export effective constrained matrices without importing ONNX packages.""" + import numpy as np - if use_crf: - start, transitions, end = check_viterbi_matches_crf(tagger, len(labels)) - transitions_path = output_dir / "crf_transitions.npz" - np.savez( - transitions_path, - start=start, - transitions=transitions, - end=end, + tagger, _tokenizer, config, marker = _load_publishable_artifact( + model_dir, require_crf=True + ) + start, transitions, end = check_viterbi_matches_crf(tagger, len(config["labels"])) + stage, output_dir = _prepare_export_stage(model_dir, output_dir) + try: + matrix_path = stage / "crf_transitions.npz" + np.savez(matrix_path, start=start, transitions=transitions, end=end) + manifest_path = stage / "manifest.json" + *_interfaces, write_json_atomic = _training_interfaces() + write_json_atomic( + manifest_path, + _export_manifest(marker, matrix_path), ) - manifest["crf_transitions"] = sha256(transitions_path) + _promote_export(stage, output_dir) + except Exception: + shutil.rmtree(stage, ignore_errors=True) + raise + return output_dir / matrix_path.name, output_dir / manifest_path.name - import onnxruntime - session = onnxruntime.InferenceSession( - str(onnx_path), - providers=["CPUExecutionProvider"], - ) - feeds = { - "input_ids": sample["input_ids"].numpy(), - "attention_mask": sample["attention_mask"].numpy(), - } - onnx_emissions = session.run(["emissions"], feeds)[0] - torch_emissions = emissions_module(*inputs).detach().numpy() - if not np.allclose(onnx_emissions, torch_emissions, atol=1e-3): - raise AssertionError("ONNX emissions differ from PyTorch emissions") +def build_emissions_module(tagger): + """Wrap only emissions computation for ONNX export.""" + import torch.nn as nn - manifest_path = output_dir / "manifest.json" - manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8") - click.echo(f"wrote {onnx_path}") - if use_crf: - click.echo(f"wrote {output_dir / 'crf_transitions.npz'}") - click.echo(f"wrote {manifest_path}") + class EmissionsModule(nn.Module): + def __init__(self): + super().__init__() + self.tagger = tagger + + def forward(self, input_ids, attention_mask): + return self.tagger.emissions(input_ids, attention_mask) + + return EmissionsModule().eval() + + +def export_onnx_emissions(model_dir, output_dir, opset=14): + """Export and verify optional ONNX emissions for a publishable model.""" + tagger, tokenizer, config, marker = _load_publishable_artifact(model_dir) + try: + import numpy as np + import onnx + import onnxruntime + import torch + except ImportError as error: + raise OnnxDependencyError( + "ONNX export requires scancode-required-phrases[training,onnx]" + ) from error + _ = onnx + + stage, output_dir = _prepare_export_stage(model_dir, output_dir) + try: + emissions_module = build_emissions_module(tagger) + sample = tokenizer( + "Licensed under the Apache License Version 2.0", + return_tensors="pt", + ) + inputs = sample["input_ids"], sample["attention_mask"] + onnx_path = stage / "model.onnx" + torch.onnx.export( + emissions_module, + inputs, + str(onnx_path), + input_names=["input_ids", "attention_mask"], + output_names=["emissions"], + dynamic_axes={ + "input_ids": {0: "batch", 1: "sequence"}, + "attention_mask": {0: "batch", 1: "sequence"}, + "emissions": {0: "batch", 1: "sequence"}, + }, + opset_version=opset, + do_constant_folding=True, + ) + + session = onnxruntime.InferenceSession( + str(onnx_path), providers=["CPUExecutionProvider"] + ) + feeds = { + "input_ids": sample["input_ids"].numpy(), + "attention_mask": sample["attention_mask"].numpy(), + } + onnx_emissions = session.run(["emissions"], feeds)[0] + torch_emissions = emissions_module(*inputs).detach().cpu().numpy() + if not isinstance(onnx_emissions, np.ndarray): + raise TypeError("ONNX emissions output must be a NumPy array") + if onnx_emissions.shape != torch_emissions.shape: + raise ValueError( + f"ONNX emissions shape {onnx_emissions.shape} differs from " + f"PyTorch {torch_emissions.shape}" + ) + if onnx_emissions.dtype != torch_emissions.dtype: + raise TypeError( + f"ONNX emissions dtype {onnx_emissions.dtype} differs from " + f"PyTorch {torch_emissions.dtype}" + ) + if not np.issubdtype(onnx_emissions.dtype, np.floating): + raise TypeError("ONNX emissions must have a floating-point dtype") + if not np.isfinite(onnx_emissions).all() or not np.isfinite(torch_emissions).all(): + raise ValueError("ONNX and PyTorch emissions must contain only finite values") + if not np.allclose(onnx_emissions, torch_emissions, atol=1e-3, rtol=1e-5): + raise AssertionError("ONNX emissions differ from PyTorch emissions") + manifest_path = stage / "manifest.json" + ( + artifact_schema, + constraint_contract, + labels, + _load_final_model, + _validate_bioes, + _validate_publishable_model, + write_json_atomic, + ) = _training_interfaces() + write_json_atomic( + manifest_path, + { + "schema": EXPORT_SCHEMA, + "artifact_schema": artifact_schema, + "constraint_contract": constraint_contract, + "labels": list(labels), + "source_final_model": marker, + "onnx_model": sha256(onnx_path), + "opset": opset, + "source_model_revision": config["resolved_model_revision"], + }, + ) + _promote_export(stage, output_dir) + except Exception: + shutil.rmtree(stage, ignore_errors=True) + raise + return output_dir / onnx_path.name, output_dir / manifest_path.name @click.command() @@ -210,23 +562,33 @@ def export(model_dir, output_dir, opset): "--model-dir", required=True, type=click.Path(exists=True, file_okay=False, path_type=Path), - help="Directory containing the trained model and train_config.json.", + help="Completed local Final_Model directory.", ) @click.option( "--output-dir", - default=None, + required=True, type=click.Path(file_okay=False, path_type=Path), - help="Output directory; defaults to the model directory.", + help="New or empty directory for this export operation.", +) +@click.option( + "--operation", + type=click.Choice(["crf", "onnx"]), + default="crf", + show_default=True, + help="Artifact operation to run.", ) @click.option("--opset", default=14, type=int, show_default=True) -def main(model_dir, output_dir, opset): - """Export a trained required phrase tagger to ONNX.""" +def main(model_dir, output_dir, operation, opset): + """Export constrained matrices or optional ONNX emissions.""" try: - export(model_dir, output_dir or model_dir, opset) - except ImportError as error: - raise click.ClickException( - f"{error}; install scancode-required-phrases[training,onnx]" - ) from error + if operation == "onnx": + paths = export_onnx_emissions(model_dir, output_dir, opset) + else: + paths = export_crf_matrices(model_dir, output_dir) + except OnnxDependencyError as error: + raise click.ClickException(str(error)) from error + for path in paths: + click.echo(f"wrote {path}") if __name__ == "__main__": diff --git a/src/scancode_required_phrases/model.py b/src/scancode_required_phrases/model.py index 309475d..7a19909 100644 --- a/src/scancode_required_phrases/model.py +++ b/src/scancode_required_phrases/model.py @@ -22,21 +22,235 @@ from scancode_required_phrases.training import LABELS +def build_constraint_masks(): + """Return deterministic BIOES start, transition, and end masks.""" + label_ids = {label: index for index, label in enumerate(LABELS)} + expected = {"O", "B-REQ", "I-REQ", "E-REQ", "S-REQ"} + if set(label_ids) != expected or len(label_ids) != len(LABELS): + raise ValueError(f"Unsupported BIOES labels: {LABELS!r}") + + start_mask = torch.zeros(len(LABELS), dtype=torch.bool) + transition_mask = torch.zeros((len(LABELS), len(LABELS)), dtype=torch.bool) + end_mask = torch.zeros(len(LABELS), dtype=torch.bool) + + for label in ("O", "B-REQ", "S-REQ"): + start_mask[label_ids[label]] = True + for label in ("O", "E-REQ", "S-REQ"): + end_mask[label_ids[label]] = True + + allowed_transitions = { + "O": ("O", "B-REQ", "S-REQ"), + "B-REQ": ("I-REQ", "E-REQ"), + "I-REQ": ("I-REQ", "E-REQ"), + "E-REQ": ("O", "B-REQ", "S-REQ"), + "S-REQ": ("O", "B-REQ", "S-REQ"), + } + for previous, following_labels in allowed_transitions.items(): + for following in following_labels: + transition_mask[label_ids[previous], label_ids[following]] = True + + return start_mask, transition_mask, end_mask + + +class ConstrainedCRF(CRF): + """CRF with hard BIOES constraints and finite learned parameters.""" + + def __init__(self, num_tags, batch_first=False): + if num_tags != len(LABELS): + raise ValueError( + f"ConstrainedCRF needs {len(LABELS)} tags, received {num_tags}" + ) + super().__init__(num_tags, batch_first=batch_first) + start_mask, transition_mask, end_mask = build_constraint_masks() + self.register_buffer("start_mask", start_mask, persistent=False) + self.register_buffer("transition_mask", transition_mask, persistent=False) + self.register_buffer("end_mask", end_mask, persistent=False) + self.o_tag_id = LABELS.index("O") + + @property + def effective_start_transitions(self): + return self.start_transitions.masked_fill(~self.start_mask, -torch.inf) + + @property + def effective_transitions(self): + return self.transitions.masked_fill(~self.transition_mask, -torch.inf) + + @property + def effective_end_transitions(self): + return self.end_transitions.masked_fill(~self.end_mask, -torch.inf) + + def _validate_learned_scores(self): + for name, scores in ( + ("start_transitions", self.start_transitions), + ("transitions", self.transitions), + ("end_transitions", self.end_transitions), + ): + if not torch.isfinite(scores).all(): + raise FloatingPointError(f"CRF {name} contains a non-finite value") + + def _validate_inputs(self, emissions, tags=None, mask=None): + if emissions.dim() != 3: + raise ValueError( + f"emissions must have rank 3, received shape {tuple(emissions.shape)}" + ) + if emissions.size(2) != self.num_tags: + raise ValueError( + f"emissions last dimension must be {self.num_tags}, " + f"received {emissions.size(2)}" + ) + + batch_size = emissions.size(0 if self.batch_first else 1) + sequence_length = emissions.size(1 if self.batch_first else 0) + if batch_size == 0 or sequence_length == 0: + raise ValueError("CRF emissions must contain a non-empty sequence batch") + if not emissions.is_floating_point() or not torch.isfinite(emissions).all(): + raise FloatingPointError("CRF emissions must contain only finite scores") + + expected_shape = emissions.shape[:2] + if mask is None: + mask = torch.ones(expected_shape, dtype=torch.bool, device=emissions.device) + elif mask.shape != expected_shape: + raise ValueError( + f"mask shape {tuple(mask.shape)} does not match {tuple(expected_shape)}" + ) + elif mask.dtype != torch.bool: + raise TypeError("CRF mask must be a boolean tensor") + elif mask.device != emissions.device: + raise ValueError("CRF mask and emissions must be on the same device") + + batch_mask = mask if self.batch_first else mask.transpose(0, 1) + if not batch_mask[:, 0].all(): + raise ValueError("CRF mask must include the first position of every row") + if (batch_mask[:, 1:] & ~batch_mask[:, :-1]).any(): + raise ValueError("CRF mask must be left aligned") + + if tags is not None: + if tags.shape != expected_shape: + raise ValueError( + f"tag shape {tuple(tags.shape)} does not match {tuple(expected_shape)}" + ) + if tags.dtype != torch.long: + raise TypeError("CRF tags must be a torch.long tensor") + if tags.device != emissions.device: + raise ValueError("CRF tags and emissions must be on the same device") + active_tags = tags[mask] + if ((active_tags < 0) | (active_tags >= self.num_tags)).any(): + raise ValueError("CRF active tags must be valid label IDs") + self._validate_gold_paths(tags, mask) + + self._validate_learned_scores() + return mask + + def _validate_gold_paths(self, tags, mask): + batch_tags = tags if self.batch_first else tags.transpose(0, 1) + batch_mask = mask if self.batch_first else mask.transpose(0, 1) + for row in range(batch_tags.size(0)): + length = int(batch_mask[row].sum().item()) + path = batch_tags[row, :length] + if not self.start_mask[path[0]] or not self.end_mask[path[-1]]: + raise ValueError(f"Invalid BIOES gold path in row {row}: {path.tolist()}") + if length > 1 and not self.transition_mask[path[:-1], path[1:]].all(): + raise ValueError(f"Invalid BIOES gold path in row {row}: {path.tolist()}") + + def forward(self, emissions, tags, mask=None, reduction="sum"): + """Return constrained conditional log likelihood for valid gold paths.""" + if reduction not in ("none", "sum", "mean", "token_mean"): + raise ValueError(f"invalid reduction: {reduction}") + mask = self._validate_inputs(emissions, tags=tags, mask=mask) + tags = tags.masked_fill(~mask, self.o_tag_id) + + if self.batch_first: + emissions = emissions.transpose(0, 1) + tags = tags.transpose(0, 1) + mask = mask.transpose(0, 1) + + numerator = self._compute_constrained_score(emissions, tags, mask) + denominator = self._compute_constrained_normalizer(emissions, mask) + likelihood = numerator - denominator + + if reduction == "none": + return likelihood + if reduction == "sum": + return likelihood.sum() + if reduction == "mean": + return likelihood.mean() + return likelihood.sum() / mask.sum() + + def decode(self, emissions, mask=None): + """Return the highest-scoring valid BIOES path for each batch row.""" + mask = self._validate_inputs(emissions, mask=mask) + if self.batch_first: + emissions = emissions.transpose(0, 1) + mask = mask.transpose(0, 1) + return self._constrained_viterbi_decode(emissions, mask) + + def _compute_constrained_score(self, emissions, tags, mask): + sequence_length, batch_size = tags.shape + batch_index = torch.arange(batch_size, device=tags.device) + score = self.effective_start_transitions[tags[0]] + score = score + emissions[0, batch_index, tags[0]] + + for position in range(1, sequence_length): + step_score = self.effective_transitions[tags[position - 1], tags[position]] + step_score = step_score + emissions[position, batch_index, tags[position]] + score = torch.where(mask[position], score + step_score, score) + + sequence_ends = mask.long().sum(dim=0) - 1 + last_tags = tags.gather(0, sequence_ends.unsqueeze(0)).squeeze(0) + return score + self.effective_end_transitions[last_tags] + + def _compute_constrained_normalizer(self, emissions, mask): + score = self.effective_start_transitions + emissions[0] + for position in range(1, emissions.size(0)): + next_score = score.unsqueeze(2) + self.effective_transitions.unsqueeze(0) + next_score = next_score + emissions[position].unsqueeze(1) + next_score = torch.logsumexp(next_score, dim=1) + score = torch.where(mask[position].unsqueeze(1), next_score, score) + return torch.logsumexp(score + self.effective_end_transitions, dim=1) + + def _constrained_viterbi_decode(self, emissions, mask): + score = self.effective_start_transitions + emissions[0] + history = [] + + for position in range(1, emissions.size(0)): + next_score = score.unsqueeze(2) + self.effective_transitions.unsqueeze(0) + next_score, indices = next_score.max(dim=1) + next_score = next_score + emissions[position] + score = torch.where(mask[position].unsqueeze(1), next_score, score) + history.append(indices) + + score = score + self.effective_end_transitions + sequence_ends = mask.long().sum(dim=0) - 1 + best_paths = [] + for row in range(emissions.size(1)): + best_last_tag = int(score[row].argmax().item()) + best_path = [best_last_tag] + for indices in reversed(history[: int(sequence_ends[row].item())]): + best_last_tag = int(indices[row, best_last_tag].item()) + best_path.append(best_last_tag) + best_path.reverse() + best_paths.append(best_path) + return best_paths + + class PhraseTagger(nn.Module): """DeBERTa backbone with a word-level token classifier and optional CRF.""" - def __init__(self, config): + def __init__(self, config, backbone=None): super().__init__() self.use_crf = config.use_crf self.aux_ce_weight = config.aux_ce_weight self.num_labels = len(LABELS) - self.backbone = AutoModel.from_pretrained( - config.model_name, - revision=config.model_revision, - ).float() - self.backbone.gradient_checkpointing_enable() - self.backbone.enable_input_require_grads() + if backbone is None: + backbone = AutoModel.from_pretrained( + config.model_name, + revision=config.model_revision, + ) + self.backbone = backbone.float() + self.backbone.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) hidden_size = self.backbone.config.hidden_size dropout = getattr(self.backbone.config, "hidden_dropout_prob", 0.1) @@ -44,7 +258,7 @@ def __init__(self, config): self.classifier = nn.Linear(hidden_size, self.num_labels) if self.use_crf: - self.crf = CRF(self.num_labels, batch_first=True) + self.crf = ConstrainedCRF(self.num_labels, batch_first=True) if self.aux_ce_weight > 0: self.register_buffer( @@ -77,10 +291,13 @@ def gather_words(self, emissions, labels): batch, _, num_labels = emissions.shape is_word = labels.ne(IGNORE_INDEX) lengths = is_word.sum(dim=1) + if (lengths == 0).any(): + raise ValueError("CRF word sequences must not be empty") width = int(lengths.max().item()) word_emissions = emissions.new_zeros((batch, width, num_labels)) - crf_tags = labels.new_zeros((batch, width)) + o_tag_id = LABELS.index("O") + crf_tags = labels.new_full((batch, width), o_tag_id) eval_tags = labels.new_full((batch, width), IGNORE_INDEX) mask = torch.zeros((batch, width), dtype=torch.bool, device=emissions.device) @@ -131,7 +348,7 @@ def predict_words(self, input_ids, attention_mask, word_ids): """Return one label ID per word for a single rule.""" positions = first_subword_positions(word_ids) if not positions: - return [] + raise ValueError("Inference word sequence must not be empty") emissions = self.emissions(input_ids, attention_mask) word_emissions = emissions[0, positions].unsqueeze(0).float() @@ -212,7 +429,7 @@ class PhraseTrainer(Trainer): def compute_loss(self, model, inputs, return_outputs=False, **kwargs): outputs = model(**inputs) loss = outputs["loss"] - if not torch.isfinite(loss): + if not torch.isfinite(loss).all(): raise FloatingPointError("Training produced a non-finite loss") return (loss, outputs) if return_outputs else loss @@ -222,7 +439,7 @@ def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None) outputs = model(**inputs) loss = outputs.get("loss") if loss is not None: - loss = loss.detach() + loss = loss.mean().detach() if prediction_loss_only: return loss, None, None return loss, outputs["predictions"], outputs["word_labels"] diff --git a/src/scancode_required_phrases/training.py b/src/scancode_required_phrases/training.py index 123b465..0ac3ecb 100644 --- a/src/scancode_required_phrases/training.py +++ b/src/scancode_required_phrases/training.py @@ -8,27 +8,52 @@ from dataclasses import asdict from dataclasses import dataclass from dataclasses import field +from functools import partial import hashlib import importlib.metadata import inspect import json +import math +from numbers import Integral import os from pathlib import Path import platform import random +import re +import subprocess +import tempfile +from types import SimpleNamespace +from urllib.parse import urlsplit +from urllib.parse import urlunsplit import click os.environ.setdefault("USE_TF", "0") -LABELS = ["O", "B-REQ", "I-REQ", "E-REQ", "S-REQ"] +LABELS = ("O", "B-REQ", "I-REQ", "E-REQ", "S-REQ") LABEL2ID = {label: index for index, label in enumerate(LABELS)} ID2LABEL = {index: label for index, label in enumerate(LABELS)} IGNORE_INDEX = -100 +ARTIFACT_SCHEMA = "scancode-required-phrases-model-v2" +CONSTRAINT_CONTRACT = "bioes-hard-v1" +H0_SERIALIZER = "raw-bytes-v1" +H1_SERIALIZER = "validated-records-canonical-json-v1" +H2_SERIALIZER = "effective-examples-canonical-json-v1" +MANIFEST_SCHEMA = "scancode-required-phrases-run-v2" +REQUIRED_ARTIFACT_FILES = { + "config.json", + "model.safetensors", + "special_tokens_map.json", + "tokenizer_config.json", + "train_config.json", + "run_manifest.json", +} +TOKENIZER_MODEL_FILES = {"tokenizer.json", "spiece.model", "sentencepiece.bpe.model"} MODEL_NAME = "microsoft/deberta-v3-large" MAX_LENGTH = 512 +IMMUTABLE_REVISION = re.compile(r"^[0-9a-fA-F]{40}$") RECORD_FIELDS = ( "identifier", "license_expression", @@ -83,29 +108,143 @@ class Config: label_weights: list = field(default_factory=lambda: [1.0, 2.0, 1.5, 1.5, 2.0]) -def validate_config(config): - """Validate settings before loading the tokenizer or model.""" +@dataclass(frozen=True) +class RecordLocation: + split: str + path: Path + line: int + + def __str__(self): + return f"{self.split} split, {self.path} line {self.line}" + + +@dataclass(frozen=True) +class ValidatedRecord: + record: dict + location: RecordLocation + + @property + def identifier(self): + return self.record["identifier"] + + +class AlignmentError(ValueError): + """Raised when tokenizer coverage cannot preserve a record safely.""" + + def __init__(self, reason, detail): + super().__init__(detail) + self.reason = reason + + +def _require_exact_type(name, value, expected): + if type(value) is not expected: + raise TypeError(f"{name} must be {expected.__name__}, not {type(value).__name__}") + + +def _require_finite_number(name, value, minimum=None, maximum=None, minimum_open=False): + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f"{name} must be a number") + if not math.isfinite(value): + raise ValueError(f"{name} must be finite") + if minimum is not None: + invalid = value <= minimum if minimum_open else value < minimum + if invalid: + comparator = "greater than" if minimum_open else "at least" + raise ValueError(f"{name} must be {comparator} {minimum}") + if maximum is not None and value > maximum: + raise ValueError(f"{name} must be at most {maximum}") + + +def validate_config(config, check_paths=True): + """Validate every setting before loading a tokenizer or model.""" + if type(config) is not Config: + raise TypeError("config must be a Config instance") + if type(check_paths) is not bool: + raise TypeError("check_paths must be a boolean") + for name in ("data_dir", "output_dir"): + value = getattr(config, name) + if not isinstance(value, Path): + raise TypeError(f"{name} must be a pathlib.Path") + if check_paths and not config.data_dir.is_dir(): + raise ValueError(f"Data directory does not exist: {config.data_dir}") + if check_paths: + for filename in ("train.jsonl", "val.jsonl", "test.jsonl"): + path = config.data_dir / filename + if not path.is_file(): + raise ValueError(f"Required split file does not exist: {path}") + if check_paths and config.output_dir.exists(): + if not config.output_dir.is_dir(): + raise ValueError(f"Output path is not a directory: {config.output_dir}") + if any(config.output_dir.iterdir()): + raise ValueError(f"Output directory is not empty: {config.output_dir}") + + for name in ("model_name", "model_revision", "optimizer", "precision"): + value = getattr(config, name) + if not isinstance(value, str) or not value: + raise ValueError(f"{name} must be a non-empty string") + if not IMMUTABLE_REVISION.fullmatch(config.model_revision): + raise ValueError("model_revision must be a full immutable 40-character commit") if config.optimizer not in {"adamw", "adamw-8bit"}: raise ValueError(f"Unsupported optimizer: {config.optimizer}") if config.precision not in {"fp32", "bf16"}: raise ValueError(f"Unsupported precision: {config.precision}") - if config.epochs < 1 or config.batch_size < 1 or config.grad_accum < 1: - raise ValueError("Epochs, batch size, and gradient accumulation must be positive") - if config.base_lr <= 0 or config.head_lr <= 0: - raise ValueError("Learning rates must be positive") - if config.aux_ce_weight < 0: - raise ValueError("Auxiliary loss weight cannot be negative") - if len(config.label_weights) != len(LABELS): - raise ValueError(f"Expected {len(LABELS)} label weights") - - -def prepare_output_dir(output_dir, resume): - """Create a new output directory or validate a resumable one.""" + + for name in ( + "max_length", "epochs", "batch_size", "grad_accum", + "early_stopping_patience", "limit", "seed", + ): + value = getattr(config, name) + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer") + if not 3 <= config.max_length <= MAX_LENGTH: + raise ValueError(f"max_length must be between 3 and {MAX_LENGTH}") + count_bounds = { + "epochs": 1_000, + "batch_size": 1_024, + "grad_accum": 65_536, + "early_stopping_patience": 1_000, + } + for name, maximum in count_bounds.items(): + value = getattr(config, name) + if not 1 <= value <= maximum: + raise ValueError(f"{name} must be between 1 and {maximum}") + if not 0 <= config.limit <= 100_000_000: + raise ValueError("limit must be between 0 and 100000000") + if config.seed < 0 or config.seed > 2**32 - 1: + raise ValueError("seed must be between 0 and 2**32 - 1") + + _require_finite_number("base_lr", config.base_lr, 0, 1, minimum_open=True) + _require_finite_number("head_lr", config.head_lr, 0, 1, minimum_open=True) + _require_finite_number("layer_decay", config.layer_decay, 0, 1, minimum_open=True) + _require_finite_number("weight_decay", config.weight_decay, 0, 1) + _require_finite_number("warmup_ratio", config.warmup_ratio, 0, 1) + _require_finite_number("max_grad_norm", config.max_grad_norm, 0, 1_000_000, minimum_open=True) + _require_finite_number("adam_epsilon", config.adam_epsilon, 0, 1, minimum_open=True) + _require_finite_number("aux_ce_weight", config.aux_ce_weight, 0, 100) + + for name in ("resume", "use_crf", "evaluate_test", "with_isr"): + if type(getattr(config, name)) is not bool: + raise TypeError(f"{name} must be a boolean") + if config.resume: + raise ValueError("Resume is unsupported for final hardened training") + if config.with_isr and not config.evaluate_test: + raise ValueError("ISR evaluation requires --evaluate-test") + if type(config.label_weights) is not list or len(config.label_weights) != len(LABELS): + raise ValueError(f"label_weights must contain exactly {len(LABELS)} values") + for index, weight in enumerate(config.label_weights): + _require_finite_number( + f"label_weights[{index}]", weight, 0, 1_000_000, minimum_open=True + ) + + +def prepare_output_dir(output_dir, resume=False): + """Create an absent or empty output directory; resume is unsupported.""" + if type(resume) is not bool: + raise TypeError("resume must be a boolean") + if resume: + raise ValueError("Resume is unsupported for final hardened training") output_dir = Path(output_dir) - checkpoints = list(output_dir.glob("checkpoint-*")) if output_dir.exists() else [] - if resume and not checkpoints: - raise ValueError(f"No checkpoint found in {output_dir}") - if not resume and output_dir.exists() and any(output_dir.iterdir()): + if output_dir.exists() and (not output_dir.is_dir() or any(output_dir.iterdir())): raise ValueError(f"Output directory is not empty: {output_dir}") output_dir.mkdir(parents=True, exist_ok=True) @@ -122,18 +261,27 @@ def set_seed(seed): def load_jsonl(path): - """Yield parsed records from a JSONL file.""" - with open(path, encoding="utf-8") as lines: - for line_number, line in enumerate(lines, 1): - line = line.strip() - if line: - yield line_number, json.loads(line) + """Yield parsed non-empty records with contextual malformed-JSON errors.""" + path = Path(path) + try: + with path.open(encoding="utf-8") as lines: + for line_number, line in enumerate(lines, 1): + if not line.strip(): + continue + try: + yield line_number, json.loads(line) + except json.JSONDecodeError as error: + raise ValueError(f"{path} line {line_number}: malformed JSON: {error.msg}") from error + except (OSError, UnicodeError) as error: + raise ValueError(f"Cannot read split file {path}: {error}") from error def validate_bioes(labels): """Return an error for an invalid BIOES sequence, or None.""" if not labels: return "has no labels" + if any(type(label) is not str for label in labels): + return "contains a non-string label" unknown = sorted(set(labels) - set(LABELS)) if unknown: return f"contains unknown labels: {unknown}" @@ -144,52 +292,336 @@ def validate_bioes(labels): return f"contains invalid transition {previous} -> {current}" if labels[-1] not in END_LABELS: return f"ends with {labels[-1]}" + return None def validate_record(record, path, line_number): - """Validate one dataset record and return it.""" + """Validate one positive raw record exactly without changing its content.""" location = f"{path} line {line_number}" + if type(record) is not dict: + raise TypeError(f"{location}: record must be an object") for field_name in RECORD_FIELDS: if field_name not in record: raise ValueError(f"{location}: missing {field_name!r}") + for field_name in ("identifier", "license_expression", "rule_type", "text"): + value = record[field_name] + if type(value) is not str: + raise TypeError(f"{location}: {field_name} must be a string") + if not value: + raise ValueError(f"{location}: empty {field_name.replace('_', ' ')}") identifier = record["identifier"] - if not identifier: - raise ValueError(f"{location}: empty identifier") - if not record["license_expression"]: - raise ValueError(f"{location} ({identifier}): empty license expression") - if not record["rule_type"]: - raise ValueError(f"{location} ({identifier}): empty rule type") - tokens = record["tokens"] labels = record["bioes_labels"] + if type(tokens) is not list: + raise TypeError(f"{location} ({identifier}): tokens must be a list") + if type(labels) is not list: + raise TypeError(f"{location} ({identifier}): bioes_labels must be a list") if not tokens: raise ValueError(f"{location} ({identifier}): no tokens") + for index, token in enumerate(tokens): + if type(token) is not str or not token: + raise ValueError( + f"{location} ({identifier}): token {index} must be a non-empty string" + ) + for index, label in enumerate(labels): + if type(label) is not str or not label: + raise ValueError( + f"{location} ({identifier}): label {index} must be a non-empty string" + ) if len(tokens) != len(labels): raise ValueError( f"{location} ({identifier}): {len(tokens)} tokens and {len(labels)} labels" ) - - unknown = sorted(set(labels) - set(LABELS)) - if unknown: - raise ValueError(f"{location} ({identifier}): unknown labels: {unknown}") - error = validate_bioes(labels) if error: raise ValueError(f"{location} ({identifier}): {error}") + if all(label == "O" for label in labels): + raise ValueError(f"{location} ({identifier}): record has no required phrase labels") return record +def _canonical_bytes(value): + return json.dumps( + value, + ensure_ascii=False, + allow_nan=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + + +def _versioned_hash(version, value): + digest = hashlib.sha256() + digest.update(version.encode("ascii") + b"\n") + digest.update(_canonical_bytes(value)) + return digest.hexdigest() + + +def validate_raw_split_hashes(paths, expected_h0): + """Require every raw split to retain its initial byte hash.""" + if set(paths) != set(expected_h0): + raise ValueError("Raw split paths and initial hashes differ") + for split, path in paths.items(): + expected = expected_h0[split].get("sha256") + actual = sha256(path) + if actual != expected: + raise ValueError( + f"{split} split changed during the run: {actual} != {expected}" + ) + + +def report_content_duplicates(records_by_split): + """Return exact token and label duplicate groups without changing records.""" + token_groups = {} + label_groups = {} + for records in records_by_split.values(): + for validated in records: + item = { + "identifier": validated.identifier, + "split": validated.location.split, + "path": str(validated.location.path), + "line": validated.location.line, + } + token_groups.setdefault(tuple(validated.record["tokens"]), []).append(item) + label_groups.setdefault(tuple(validated.record["bioes_labels"]), []).append(item) + + def duplicates(groups): + return [ + {"value": list(value), "records": locations} + for value, locations in groups.items() + if len(locations) > 1 + ] + + return {"tokens": duplicates(token_groups), "labels": duplicates(label_groups)} + + +def load_and_validate_splits(paths): + """Load every raw split, validate all records, and return pre-selection hashes.""" + if set(paths) != {"train", "validation", "test"}: + raise ValueError("paths must contain train, validation, and test splits") + records_by_split = {} + h0 = {} + for split, path_value in paths.items(): + path = Path(path_value) + if not path.is_file(): + raise ValueError(f"Missing {split} split file: {path}") + try: + raw = path.read_bytes() + except OSError as error: + raise ValueError(f"Cannot read {split} split file {path}: {error}") from error + if not raw: + raise ValueError(f"{split} split file is empty: {path}") + h0[split] = { + "serializer": H0_SERIALIZER, + "path": str(path), + "sha256": hashlib.sha256(raw).hexdigest(), + } + records = [] + for line_number, unvalidated in load_jsonl(path): + record = validate_record(unvalidated, path, line_number) + location = RecordLocation(split, path, line_number) + records.append(ValidatedRecord(record, location)) + if not records: + raise ValueError(f"{split} split has no Current_Record: {path}") + records_by_split[split] = records + + seen = {} + for records in records_by_split.values(): + for validated in records: + identifier = validated.identifier + if identifier in seen: + raise ValueError( + f"Duplicate identifier {identifier!r}: " + f"{seen[identifier]} and {validated.location}" + ) + seen[identifier] = validated.location + + h1 = {} + for split, records in records_by_split.items(): + material = { + "split": split, + "count": len(records), + "records": [validated.record for validated in records], + } + h1[split] = { + "serializer": H1_SERIALIZER, + "count": len(records), + "sha256": _versioned_hash(H1_SERIALIZER, material), + } + report = { + "h0": h0, + "h1": h1, + "duplicates": report_content_duplicates(records_by_split), + "raw_counts": {split: len(records) for split, records in records_by_split.items()}, + } + return records_by_split, report + + +def _validated_word_ids(encoding, word_count, context, special_ids, vocab_size=None): + try: + word_ids = list(encoding.word_ids()) + input_ids = list(encoding["input_ids"]) + attention_mask = list(encoding["attention_mask"]) + except (AttributeError, KeyError, TypeError) as error: + raise AlignmentError( + "missing-word-ids", f"{context}: tokenizer returned incomplete model inputs" + ) from error + if len(word_ids) != len(input_ids) or len(word_ids) != len(attention_mask): + raise AlignmentError( + "shape-mismatch", f"{context}: word IDs, input IDs, and attention mask lengths differ" + ) + if any(isinstance(value, bool) or not isinstance(value, int) or value < 0 for value in input_ids): + raise AlignmentError("invalid-input-id", f"{context}: input IDs must be non-negative integers") + if vocab_size is not None and any(value >= vocab_size for value in input_ids): + raise AlignmentError("invalid-input-id", f"{context}: input ID is outside tokenizer vocabulary") + if any(type(value) is not int or value != 1 for value in attention_mask): + raise AlignmentError("invalid-attention-mask", f"{context}: attention mask must be active and binary") + covered_positions = [ + index for index, word_id in enumerate(word_ids) if word_id is not None + ] + if not covered_positions: + raise AlignmentError("zero-coverage", f"{context}: tokenizer covered no dataset words") + first_covered = covered_positions[0] + last_covered = covered_positions[-1] + if first_covered == 0 or last_covered == len(word_ids) - 1: + raise AlignmentError( + "missing-special-token", f"{context}: required boundary special tokens are absent" + ) + boundary_positions = list(range(first_covered)) + list( + range(last_covered + 1, len(word_ids)) + ) + if not special_ids or any(input_ids[index] not in special_ids for index in boundary_positions): + raise AlignmentError( + "missing-special-token", f"{context}: boundary IDs are not tokenizer special tokens" + ) + if any(word_ids[index] is None for index in range(first_covered, last_covered + 1)): + raise AlignmentError( + "coverage-gap", f"{context}: special-token gap occurs inside word coverage" + ) + covered = [word_ids[index] for index in covered_positions] + if any(type(word_id) is not int for word_id in covered): + raise AlignmentError("invalid-word-id", f"{context}: tokenizer returned a non-integer word ID") + if any(word_id < 0 or word_id >= word_count for word_id in covered): + raise AlignmentError("out-of-range-word-id", f"{context}: tokenizer returned an out-of-range word ID") + distinct = [] + previous = None + for word_id in covered: + if previous is None or word_id != previous: + distinct.append(word_id) + if previous is not None and word_id != previous + 1: + raise AlignmentError( + "noncontiguous-coverage", + f"{context}: tokenizer word IDs are missing, decreasing, or noncontiguous", + ) + previous = word_id + if distinct[0] != 0: + raise AlignmentError("coverage-gap", f"{context}: tokenizer coverage does not start at word 0") + return word_ids, covered + + +def _coverage_counts(word_ids): + counts = {} + for word_id in word_ids: + if word_id is not None: + counts[word_id] = counts.get(word_id, 0) + 1 + return counts + + def align_labels(tokens, word_labels, tokenizer, max_length): - """Tokenize words and label only the first subword of each word.""" + """Align unchanged labels to a tokenizer-verified complete word prefix.""" + if type(tokens) is not list or type(word_labels) is not list: + raise TypeError("tokens and word_labels must be lists") + if len(tokens) != len(word_labels) or not tokens: + raise ValueError("tokens and word_labels must have equal non-zero lengths") + if validate_bioes(word_labels): + raise ValueError("word_labels must be a valid BIOES sequence") + if getattr(tokenizer, "is_fast", True) is not True: + raise ValueError("Training requires a fast tokenizer with word IDs") + + call = { + "is_split_into_words": True, + "add_special_tokens": True, + } + special_ids = set(getattr(tokenizer, "all_special_ids", ())) + vocab_size = getattr(tokenizer, "vocab_size", None) + if isinstance(vocab_size, bool) or ( + vocab_size is not None and (not isinstance(vocab_size, int) or vocab_size <= 0) + ): + raise ValueError("Tokenizer vocab_size must be a positive integer") + full = tokenizer(tokens, truncation=False, **call) + full_word_ids, _full_covered = _validated_word_ids( + full, len(tokens), "full encoding", special_ids, vocab_size + ) + full_counts = _coverage_counts(full_word_ids) + expected_ids = list(range(len(tokens))) + if sorted(full_counts) != expected_ids or any(full_counts[index] < 1 for index in expected_ids): + missing = [index for index in expected_ids if index not in full_counts] + raise AlignmentError( + "zero-coverage", + f"full encoding: dataset words have zero subwords at positions {missing}", + ) + encoding = tokenizer( tokens, - is_split_into_words=True, truncation=True, max_length=max_length, + **call, + ) + word_ids, covered = _validated_word_ids( + encoding, len(tokens), "retained encoding", special_ids, vocab_size ) + retained_counts = _coverage_counts(word_ids) + covered_words = max(covered) + 1 + complete_words = covered_words + for word_id in range(covered_words): + if retained_counts.get(word_id, 0) != full_counts[word_id]: + if word_id != covered_words - 1: + raise AlignmentError( + "partial-word-coverage", + f"retained encoding partially covers dataset word {word_id} before a later word", + ) + complete_words = word_id + + omitted_positions = list(range(complete_words, len(tokens))) + omitted_required = [ + position for position in omitted_positions if word_labels[position] != "O" + ] + if omitted_required: + raise AlignmentError( + "omitted-non-o", + f"truncation omits non-O labels at positions {omitted_required}", + ) + if not complete_words: + raise AlignmentError( + "zero-complete-word-prefix", + "retained encoding contains no complete dataset word", + ) + + if complete_words != covered_words: + encoding = tokenizer(tokens[:complete_words], truncation=False, **call) + word_ids, _covered = _validated_word_ids( + encoding, + complete_words, + "complete-prefix encoding", + special_ids, + vocab_size, + ) + prefix_counts = _coverage_counts(word_ids) + if any( + prefix_counts.get(word_id, 0) != full_counts[word_id] + for word_id in range(complete_words) + ): + raise AlignmentError( + "partial-word-coverage", + "complete-prefix encoding changed retained word coverage", + ) + + if len(encoding["input_ids"]) > max_length: + raise AlignmentError("length-overflow", "complete-prefix encoding exceeds max_length") + if len(encoding["attention_mask"]) != len(encoding["input_ids"]): + raise AlignmentError("shape-mismatch", "tokenizer input and attention lengths differ") - word_ids = encoding.word_ids() label_ids = [] previous_word = None for word_id in word_ids: @@ -200,25 +632,12 @@ def align_labels(tokens, word_labels, tokenizer, max_length): else: label_ids.append(IGNORE_INDEX) previous_word = word_id - encoding["labels"] = label_ids - kept_word_ids = [word_id for word_id in word_ids if word_id is not None] - kept_words = max(kept_word_ids) + 1 if kept_word_ids else 0 - truncated = kept_words < len(tokens) - cut_phrase = ( - truncated - and kept_words - and word_labels[kept_words - 1] - in { - "B-REQ", - "I-REQ", - } - ) - return encoding, truncated, bool(cut_phrase) + return encoding, bool(omitted_positions), False def first_subword_positions(word_ids): - """Return positions that start a tokenizer word.""" + """Return positions that start each contiguous tokenizer word.""" positions = [] previous = None for index, word_id in enumerate(word_ids): @@ -232,37 +651,68 @@ def first_subword_positions(word_ids): class PhraseDataset: - """Read and encode one BIOES JSONL split.""" + """Hold selected effective examples for one validated split.""" - def __init__(self, path, tokenizer, max_length, limit=0): + def __init__(self, records, tokenizer, max_length, limit=0): self.examples = [] self.identifiers = [] + self.effective_inventory = [] + self.rejections = [] + self.truncations = [] self.truncated = 0 self.cut_phrases = 0 - for line_number, unvalidated_record in load_jsonl(path): - if limit and len(self.examples) >= limit: - break - record = validate_record(unvalidated_record, path, line_number) - encoding, truncated, cut_phrase = align_labels( - record["tokens"], - record["bioes_labels"], - tokenizer, - max_length, - ) + for validated in records: + if type(validated) is not ValidatedRecord: + raise TypeError("PhraseDataset requires ValidatedRecord instances") + record = validated.record + try: + encoding, truncated, _unused = align_labels( + record["tokens"], record["bioes_labels"], tokenizer, max_length + ) + except AlignmentError as error: + rejection = { + "identifier": validated.identifier, + "split": validated.location.split, + "path": str(validated.location.path), + "line": validated.location.line, + "reason": error.reason, + "message": str(error), + } + self.rejections.append(rejection) + if error.reason == "omitted-non-o": + self.cut_phrases += 1 + continue + + example = { + "input_ids": list(encoding["input_ids"]), + "attention_mask": list(encoding["attention_mask"]), + "labels": list(encoding["labels"]), + } + inventory = { + "identifier": validated.identifier, + **example, + "truncated": truncated, + "location": { + "split": validated.location.split, + "path": str(validated.location.path), + "line": validated.location.line, + }, + "selected": False, + } + self.effective_inventory.append(inventory) if truncated: self.truncated += 1 - if cut_phrase: - self.cut_phrases += 1 - continue + self.truncations.append( + {"identifier": validated.identifier, **inventory["location"]} + ) - self.identifiers.append(record["identifier"]) + selected_count = limit or len(self.effective_inventory) + for inventory in self.effective_inventory[:selected_count]: + inventory["selected"] = True + self.identifiers.append(inventory["identifier"]) self.examples.append( - { - "input_ids": encoding["input_ids"], - "attention_mask": encoding["attention_mask"], - "labels": encoding["labels"], - } + {name: inventory[name] for name in ("input_ids", "attention_mask", "labels")} ) def __len__(self): @@ -273,70 +723,167 @@ def __getitem__(self, index): def validate_splits(datasets): - """Require non-empty splits with disjoint rule identifiers.""" - seen = {} + """Require every selected split to contain an effective example.""" for split_name, dataset in datasets.items(): if not dataset: - raise ValueError(f"{split_name} split has no usable examples") - for identifier in dataset.identifiers: - previous_split = seen.get(identifier) - if previous_split: - raise ValueError( - f"Rule {identifier!r} occurs in both {previous_split} and {split_name}" - ) - seen[identifier] = split_name + raise ValueError(f"{split_name} split has no selected Effective_Example") + + +def build_effective_datasets(records_by_split, tokenizer, max_length, limit=0): + """Build all effective inventories before applying the smoke-run limit.""" + datasets = { + split: PhraseDataset(records, tokenizer, max_length, limit) + for split, records in records_by_split.items() + } + validate_splits(datasets) + h2 = {} + for split, dataset in datasets.items(): + material = { + "split": split, + "limit": limit, + "effective": dataset.effective_inventory, + "rejections": dataset.rejections, + } + h2[split] = { + "serializer": H2_SERIALIZER, + "effective_count": len(dataset.effective_inventory), + "selected_count": len(dataset), + "rejected_count": len(dataset.rejections), + "sha256": _versioned_hash(H2_SERIALIZER, material), + } + return datasets, h2 def extract_spans(tags): - """Return the set of inclusive word spans in a BIOES sequence.""" - spans = [] + """Return inclusive spans from an already valid BIOES sequence.""" + error = validate_bioes(tags) + if error: + raise ValueError(f"Cannot extract spans from invalid BIOES: {error}") + spans = set() start = None for index, tag in enumerate(tags): if tag == "S-REQ": - spans.append((index, index)) - start = None + spans.add((index, index)) elif tag == "B-REQ": - if start is not None: - spans.append((start, index - 1)) start = index - elif tag == "I-REQ": - if start is None: - start = index elif tag == "E-REQ": - if start is None: - start = index - spans.append((start, index)) - start = None - elif start is not None: - spans.append((start, index - 1)) + spans.add((start, index)) start = None - if start is not None: - spans.append((start, len(tags) - 1)) - return set(spans) + return spans + + +def _integer_id(name, value, row, column, allowed): + if isinstance(value, bool) or not isinstance(value, Integral): + raise TypeError(f"{name} at row {row}, column {column} is not an integer ID") + converted = int(value) + if converted not in allowed: + raise ValueError(f"Unknown {name} ID {value!r} at row {row}, column {column}") + return converted -def decode_row(pred_row, label_row): - """Drop ignored positions and map label IDs to BIOES tags.""" +def decode_row(pred_row, label_row, row=0): + """Strictly map equal-length active prediction and gold IDs to BIOES tags.""" + if len(pred_row) != len(label_row): + raise ValueError(f"Prediction and label row {row} lengths differ") predicted = [] actual = [] - for prediction, label in zip(pred_row, label_row): - if int(label) == IGNORE_INDEX: + for column in range(len(label_row)): + label = _integer_id( + "gold label", label_row[column], row, column, set(ID2LABEL) | {IGNORE_INDEX} + ) + if label == IGNORE_INDEX: continue - actual.append(ID2LABEL[int(label)]) - predicted.append(ID2LABEL.get(int(prediction), "O")) + prediction = _integer_id( + "prediction", pred_row[column], row, column, set(ID2LABEL) + ) + actual.append(ID2LABEL[label]) + predicted.append(ID2LABEL[prediction]) + if not actual: + raise ValueError(f"Metric row {row} has no active gold labels") return predicted, actual -def compute_metrics(eval_pred): - """Return strict span-level micro metrics and rule-level exact match.""" +def _metric_matrix(name, value): + if hasattr(value, "detach"): + value = value.detach().cpu().tolist() + elif hasattr(value, "tolist"): + value = value.tolist() + if not isinstance(value, (list, tuple)): + raise TypeError(f"{name} must be a rank-two integer array") + if not value: + raise ValueError("Metrics require at least one row") + rows = [] + width = None + for row, values in enumerate(value): + if not isinstance(values, (list, tuple)): + raise ValueError(f"{name} must have rank 2; row {row} is not a row") + if width is None: + width = len(values) + elif len(values) != width: + raise ValueError(f"{name} rows must have equal lengths") + rows.append(list(values)) + if width == 0: + raise ValueError(f"{name} rows must not be empty") + return rows + + +def _validate_crf_metric_padding(pred_row, label_row, row): + """Require packed CRF rows to use one left-aligned active prefix.""" + padding_started = False + for column in range(len(label_row)): + label = _integer_id( + "gold label", label_row[column], row, column, set(ID2LABEL) | {IGNORE_INDEX} + ) + prediction = _integer_id( + "prediction", pred_row[column], row, column, set(ID2LABEL) | {IGNORE_INDEX} + ) + if label == IGNORE_INDEX: + padding_started = True + if prediction != IGNORE_INDEX: + raise ValueError( + f"CRF prediction padding at row {row}, column {column} must be IGNORE_INDEX" + ) + elif padding_started: + raise ValueError(f"CRF metric row {row} padding must be left aligned") + elif prediction == IGNORE_INDEX: + raise ValueError( + f"CRF active prediction at row {row}, column {column} is IGNORE_INDEX" + ) + + +def compute_metrics(eval_pred, use_crf=False): + """Return strict span metrics with explicit invalid-path accounting.""" + if type(use_crf) is not bool: + raise TypeError("use_crf must be a boolean") predictions, labels = eval_pred + predictions = _metric_matrix("Predictions", predictions) + labels = _metric_matrix("Labels", labels) + if len(predictions) != len(labels): + raise ValueError("Prediction and label batch sizes differ") + if len(predictions[0]) != len(labels[0]): + raise ValueError("Prediction and label row lengths differ") + true_positive = false_positive = false_negative = 0 - exact = 0 + exact = invalid_paths = 0 + for row in range(len(predictions)): + if use_crf: + _validate_crf_metric_padding(predictions[row], labels[row], row) + predicted, actual = decode_row(predictions[row], labels[row], row) + actual_error = validate_bioes(actual) + if actual_error: + raise ValueError(f"Invalid gold BIOES path in row {row}: {actual}; {actual_error}") + actual_spans = extract_spans(actual) + prediction_error = validate_bioes(predicted) + if prediction_error: + if use_crf: + raise ValueError( + f"Invalid CRF BIOES path in row {row}: {predicted}; {prediction_error}" + ) + invalid_paths += 1 + false_negative += len(actual_spans) + continue - for pred_row, label_row in zip(predictions, labels): - predicted, actual = decode_row(pred_row, label_row) predicted_spans = extract_spans(predicted) - actual_spans = extract_spans(actual) true_positive += len(predicted_spans & actual_spans) false_positive += len(predicted_spans - actual_spans) false_negative += len(actual_spans - predicted_spans) @@ -352,29 +899,28 @@ def compute_metrics(eval_pred): "f1": f1, "precision": precision, "recall": recall, - "exact_match": exact / len(predictions) if len(predictions) else 0.0, + "exact_match": exact / len(predictions), "predicted_spans": true_positive + false_positive, "gold_spans": true_positive + false_negative, + "invalid_paths": invalid_paths, } def evaluate_isr(records, model, tokenizer, max_length): - """Return the fraction of predicted phrases ScanCode can locate.""" + """Return predicted-phrase locatability; this is not injection success.""" import torch from licensedcode.required_phrases import find_phrase_spans_in_text device = next(model.parameters()).device model.eval() - total = 0 - injectable = 0 - for record in records: - encoding, _, cut_phrase = align_labels( - record["tokens"], - record["bioes_labels"], - tokenizer, - max_length, - ) - if cut_phrase: + total = locatable = 0 + for item in records: + record = item.record if isinstance(item, ValidatedRecord) else item + try: + encoding, _truncated, _unused = align_labels( + record["tokens"], record["bioes_labels"], tokenizer, max_length + ) + except AlignmentError: continue inputs = { "input_ids": torch.tensor([encoding["input_ids"]], device=device), @@ -383,19 +929,23 @@ def evaluate_isr(records, model, tokenizer, max_length): } with torch.no_grad(): output = model(**inputs) - tags, _ = decode_row( + tags, _actual = decode_row( output["predictions"][0].tolist(), output["word_labels"][0].tolist(), ) + error = validate_bioes(tags) + if error: + if model.use_crf: + raise ValueError(f"Invalid CRF ISR path: {tags}; {error}") + continue for start, end in extract_spans(tags): if end >= len(record["tokens"]): continue phrase = " ".join(record["tokens"][start : end + 1]) total += 1 if find_phrase_spans_in_text(record["text"], phrase): - injectable += 1 - - return injectable / total if total else 0.0 + locatable += 1 + return locatable / total if total else 0.0 def sha256(path): @@ -416,47 +966,635 @@ def installed_version(package_name): def serializable_config(config): - """Return the training configuration with paths converted to strings.""" + """Return the complete training configuration with string paths.""" values = asdict(config) values["data_dir"] = str(values["data_dir"]) values["output_dir"] = str(values["output_dir"]) return values +def write_json_atomic(path, value): + """Durably replace one JSON file without exposing partial content.""" + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary_name = tempfile.mkstemp( + prefix=f".{path.name}.", suffix=".tmp", dir=path.parent + ) + temporary = Path(temporary_name) + try: + with os.fdopen(descriptor, "w", encoding="utf-8", newline="\n") as stream: + json.dump(value, stream, ensure_ascii=False, allow_nan=False, indent=2, sort_keys=True) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + try: + directory = os.open(path.parent, os.O_RDONLY) + except (AttributeError, OSError): + directory = None + if directory is not None: + try: + os.fsync(directory) + finally: + os.close(directory) + finally: + if temporary.exists(): + temporary.unlink() + + +def write_failure_manifest( + manifest_path, pre_run_manifest, phase, error, completed_checks, retained_artifacts +): + """Atomically record a failed phase without environment or traceback capture.""" + failure_manifest = { + **pre_run_manifest, + "state": "failure", + "failed_phase": phase, + "exception": {"type": type(error).__name__, "message": str(error)}, + "completed_checks": list(completed_checks), + "retained_artifacts": [str(path) for path in retained_artifacts if Path(path).exists()], + } + write_json_atomic(manifest_path, failure_manifest) + return failure_manifest + + +def _run_git(repo_dir, *arguments): + result = subprocess.run( + ["git", "-C", str(repo_dir), *arguments], + check=False, + capture_output=True, + text=True, + encoding="utf-8", + errors="strict", + ) + if result.returncode: + raise RuntimeError(f"Cannot collect source provenance with git {' '.join(arguments)}") + return result.stdout.rstrip("\n") + + +def _redacted_repository_identity(value): + if "://" not in value: + return value.split("@", 1)[-1] if "@" in value else value + parsed = urlsplit(value) + hostname = parsed.hostname or "" + if parsed.port: + hostname = f"{hostname}:{parsed.port}" + return urlunsplit((parsed.scheme, hostname, parsed.path, "", "")) + + +def collect_source_provenance(repo_dir=None): + """Return narrow repository identity without reading environment variables.""" + repo_dir = Path(repo_dir or Path(__file__).resolve().parents[2]) + root = Path(_run_git(repo_dir, "rev-parse", "--show-toplevel")).resolve() + branch = _run_git(root, "branch", "--show-current") or "DETACHED" + commit = _run_git(root, "rev-parse", "HEAD") + if not IMMUTABLE_REVISION.fullmatch(commit): + raise RuntimeError("Source commit is not an immutable full identity") + repository = _redacted_repository_identity( + _run_git(root, "config", "--get", "remote.origin.url") + ) + return { + "repository": repository, + "root": str(root), + "branch": branch, + "commit": commit, + "dirty": bool(_run_git(root, "status", "--porcelain")), + } + + +def _optimizer_provenance(name): + if name == "adamw": + from torch.optim import AdamW + + implementation = AdamW + package = "torch" + elif name == "adamw-8bit": + try: + from bitsandbytes.optim import AdamW8bit + except ImportError as error: + raise RuntimeError( + "adamw-8bit requires bitsandbytes; install the training-8bit extra" + ) from error + implementation = AdamW8bit + package = "bitsandbytes" + else: + raise ValueError(f"Unsupported optimizer: {name}") + return { + "configured": name, + "implementation": f"{implementation.__module__}.{implementation.__qualname__}", + "package": package, + "version": installed_version(package), + } + + +def collect_runtime_provenance(optimizer, precision): + """Return runtime API provenance without capturing environment variables.""" + import torch + import transformers + + cuda_available = torch.cuda.is_available() + return { + "python": platform.python_version(), + "implementation": platform.python_implementation(), + "platform": platform.platform(), + "package": installed_version("scancode-required-phrases"), + "scancode_toolkit": installed_version("scancode-toolkit"), + "torch": torch.__version__, + "transformers": transformers.__version__, + "pytorch_crf": installed_version("pytorch-crf"), + "cuda_available": cuda_available, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "device": torch.cuda.get_device_name(0) if cuda_available else "cpu", + "optimizer": _optimizer_provenance(optimizer), + "precision": precision, + } + + +def validate_state_structure(expected, actual, expected_name="expected", actual_name="actual"): + """Require ordered keys, shapes, dtypes, and finite tensor values.""" + import torch + + expected_keys = list(expected) + actual_keys = list(actual) + if expected_keys != actual_keys: + missing = [name for name in expected_keys if name not in actual] + unexpected = [name for name in actual_keys if name not in expected] + raise ValueError( + f"State key mismatch for {expected_name} and {actual_name}; " + f"missing={missing}, unexpected={unexpected}" + ) + for name in expected_keys: + left = expected[name].detach().cpu() + right = actual[name].detach().cpu() + if left.shape != right.shape: + raise ValueError( + f"Tensor {name!r} shape mismatch: {tuple(left.shape)} != {tuple(right.shape)}" + ) + if left.dtype != right.dtype: + raise ValueError(f"Tensor {name!r} dtype mismatch: {left.dtype} != {right.dtype}") + if not torch.isfinite(left).all() or not torch.isfinite(right).all(): + raise ValueError(f"Tensor {name!r} contains non-finite values") + + +def validate_state_dicts(expected, actual, expected_name="expected", actual_name="actual"): + """Require exact ordered state structure and tensor values.""" + import torch + + validate_state_structure(expected, actual, expected_name, actual_name) + for name in expected: + left = expected[name].detach().cpu() + right = actual[name].detach().cpu() + if not torch.equal(left, right): + raise ValueError(f"Tensor {name!r} values differ") + + +def _canonical_state_dict(state): + return {name: state[name] for name in sorted(state)} + + +def _load_state_file(model_path): + import torch + + model_path = Path(model_path) + if model_path.suffix == ".safetensors": + from safetensors.torch import load_file + return load_file(str(model_path)) + return torch.load(model_path, map_location="cpu", weights_only=True) + + def validate_saved_state(model, model_path): - """Validate the keys, shapes, and values in a saved safetensors model.""" + """Require one saved model state to equal the in-memory state exactly.""" + saved = _load_state_file(model_path) + validate_state_dicts( + _canonical_state_dict(model.state_dict()), + _canonical_state_dict(saved), + "in-memory", + str(model_path), + ) + + +def _artifact_config(config, resolved_revision): + values = serializable_config(config) + values.update( + { + "artifact_schema": ARTIFACT_SCHEMA, + "constraint_contract": CONSTRAINT_CONTRACT, + "labels": list(LABELS), + "requested_model_revision": config.model_revision, + "resolved_model_revision": resolved_revision, + "model_revision": resolved_revision, + } + ) + return values + + +def _validate_artifact_config(values): + if type(values) is not dict: + raise TypeError("train_config.json must contain an object") + config_fields = set(Config.__dataclass_fields__) + contract_fields = { + "artifact_schema", + "constraint_contract", + "labels", + "requested_model_revision", + "resolved_model_revision", + } + expected_fields = config_fields | contract_fields + if set(values) != expected_fields: + missing = sorted(expected_fields - set(values)) + unexpected = sorted(set(values) - expected_fields) + raise ValueError( + f"Artifact configuration fields differ; missing={missing}, unexpected={unexpected}" + ) + if values["artifact_schema"] != ARTIFACT_SCHEMA: + raise ValueError( + f"Unsupported artifact schema; expected {ARTIFACT_SCHEMA}. Old artifacts must be retrained." + ) + if type(values["labels"]) is not list or tuple(values["labels"]) != LABELS: + raise ValueError("Artifact label order does not match the supported LABELS") + if values["constraint_contract"] != CONSTRAINT_CONTRACT: + raise ValueError("Artifact constraint contract is unsupported") + revision = values["resolved_model_revision"] + if not isinstance(revision, str) or not IMMUTABLE_REVISION.fullmatch(revision): + raise ValueError("Artifact has no immutable resolved model revision") + if values["requested_model_revision"] != revision or values["model_revision"] != revision: + raise ValueError("Requested, resolved, and configured artifact revisions differ") + + config_values = {name: values[name] for name in config_fields} + for name in ("data_dir", "output_dir"): + if type(config_values[name]) is not str or not config_values[name]: + raise ValueError(f"Artifact {name} must be a non-empty path string") + config_values[name] = Path(config_values[name]) + validate_config(Config(**config_values), check_paths=False) + return values + + +def _load_local_model(model_dir, offline=True): + """Strictly reconstruct a supported model and tokenizer from local files only.""" import torch - from safetensors.torch import load_file + from transformers import AutoConfig + from transformers import AutoModel + from transformers import AutoTokenizer + + from scancode_required_phrases.model import PhraseTagger + + if offline is not True: + raise ValueError("Final_Model loading is local-only") + model_dir = Path(model_dir) + config_path = model_dir / "train_config.json" + try: + values = _validate_artifact_config(json.loads(config_path.read_text(encoding="utf-8"))) + except (OSError, json.JSONDecodeError) as error: + raise ValueError(f"Cannot read supported artifact configuration: {error}") from error + local_config = AutoConfig.from_pretrained(str(model_dir), local_files_only=True) + backbone = AutoModel.from_config(local_config) + tagger_config = SimpleNamespace(**values) + model = PhraseTagger(tagger_config, backbone=backbone) + state_path = model_dir / "model.safetensors" + if not state_path.is_file(): + raise ValueError(f"Final_Model is missing {state_path.name}") + state = _load_state_file(state_path) + validate_state_structure( + _canonical_state_dict(model.state_dict()), + _canonical_state_dict(state), + "constructed", + "saved", + ) + model.load_state_dict(state, strict=True) + validate_state_dicts( + _canonical_state_dict(state), + _canonical_state_dict(model.state_dict()), + "saved", + "loaded", + ) + tokenizer = AutoTokenizer.from_pretrained( + str(model_dir), use_fast=True, local_files_only=True + ) + if not tokenizer.is_fast: + raise ValueError("Final_Model tokenizer is not fast") + if any(not torch.isfinite(tensor).all() for tensor in model.state_dict().values()): + raise ValueError("Final_Model contains non-finite tensors") + return model.eval(), tokenizer + + +def load_final_model(model_dir, offline=True): + """Validate publication and load one supported Final_Model locally.""" + validate_publishable_model(model_dir) + return _load_local_model(model_dir, offline=offline) - saved = load_file(str(model_path)) - expected = model.state_dict() - if set(saved) != set(expected): - missing = sorted(set(expected) - set(saved)) - unexpected = sorted(set(saved) - set(expected)) - raise ValueError(f"Saved model key mismatch; missing={missing}, unexpected={unexpected}") - for name, tensor in saved.items(): - if tensor.shape != expected[name].shape: +def _all_file_hashes(directory): + directory = Path(directory) + return { + path.relative_to(directory).as_posix(): sha256(path) + for path in sorted(directory.rglob("*")) + if path.is_file() and path.name != "SUCCESS.json" + } + + +def validate_publishable_model(model_dir): + """Require a supported Success_Marker and exact Final_Model file hashes.""" + model_dir = Path(model_dir) + marker_path = model_dir / "SUCCESS.json" + try: + marker = json.loads(marker_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ValueError(f"Final_Model has no valid Success_Marker: {error}") from error + if marker.get("schema") != ARTIFACT_SCHEMA: + raise ValueError("Success_Marker uses an unsupported artifact schema") + if marker.get("constraint_contract") != CONSTRAINT_CONTRACT: + raise ValueError("Success_Marker uses an unsupported constraint contract") + selected = marker.get("selected_checkpoint") + if not isinstance(selected, str) or not selected: + raise ValueError("Success_Marker has no Selected_Checkpoint") + files = marker.get("files") + if type(files) is not dict or not files: + raise ValueError("Success_Marker has no required file inventory") + actual_files = _all_file_hashes(model_dir) + missing_required = sorted(REQUIRED_ARTIFACT_FILES - set(actual_files)) + if missing_required: + raise ValueError(f"Final_Model is missing required files: {missing_required}") + if not TOKENIZER_MODEL_FILES.intersection(actual_files): + raise ValueError("Final_Model is missing a complete fast-tokenizer model file") + if set(files) != set(actual_files): + raise ValueError("Final_Model file inventory does not match Success_Marker") + for name, digest in files.items(): + if not isinstance(digest, str) or digest != actual_files[name]: + raise ValueError(f"Final_Model hash mismatch for {name}") + manifest_path = model_dir / "run_manifest.json" + if marker.get("run_manifest_sha256") != sha256(manifest_path): + raise ValueError("Success_Marker run manifest hash does not match") + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + config = json.loads((model_dir / "train_config.json").read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ValueError(f"Final_Model metadata is malformed: {error}") from error + required_manifest_fields = { + "schema", + "state", + "config", + "contracts", + "dataset", + "model_identity", + "source", + "runtime", + "completed_checks", + "selected_checkpoint", + "best_validation_f1", + "validation_metrics", + "ordered_validation", + "test_metrics", + "artifact_files", + "log_history", + } + if set(manifest) != required_manifest_fields: + raise ValueError("Final_Model run manifest fields are incomplete or unsupported") + if manifest["schema"] != MANIFEST_SCHEMA or manifest["state"] != "success": + raise ValueError("Final_Model run manifest is not a supported success manifest") + if manifest["selected_checkpoint"] != selected: + raise ValueError("Success_Marker Selected_Checkpoint does not match the manifest") + if manifest["config"] != {name: config[name] for name in Config.__dataclass_fields__}: + raise ValueError("Final_Model manifest and artifact configurations differ") + expected_contracts = { + "artifact_schema": ARTIFACT_SCHEMA, + "constraint_contract": CONSTRAINT_CONTRACT, + "labels": list(LABELS), + } + if manifest["contracts"] != expected_contracts: + raise ValueError("Final_Model manifest contracts are unsupported") + dataset = manifest["dataset"] + if type(dataset) is not dict or not {"paths", "h0", "h1", "h2", "report"} <= set(dataset): + raise ValueError("Final_Model manifest dataset provenance is incomplete") + if type(manifest["source"]) is not dict or not { + "repository", "root", "branch", "commit", "dirty" + } <= set(manifest["source"]): + raise ValueError("Final_Model source provenance is incomplete") + if type(manifest["runtime"]) is not dict or not { + "python", "platform", "torch", "transformers", "optimizer", "precision" + } <= set(manifest["runtime"]): + raise ValueError("Final_Model runtime provenance is incomplete") + if not isinstance(manifest["completed_checks"], list) or not manifest["completed_checks"]: + raise ValueError("Final_Model completed checks are missing") + if not isinstance(manifest["best_validation_f1"], (int, float)) or not math.isfinite( + manifest["best_validation_f1"] + ): + raise ValueError("Final_Model best validation F1 is invalid") + for name in ("validation_metrics", "ordered_validation"): + if type(manifest[name]) is not dict: + raise ValueError(f"Final_Model {name} is invalid") + expected_components = { + name: digest for name, digest in actual_files.items() if name != "run_manifest.json" + } + if manifest["artifact_files"] != expected_components: + raise ValueError("Final_Model manifest artifact hashes do not match") + _validate_artifact_config(config) + return marker + + +def _checkpoint_state_path(checkpoint, output_dir): + checkpoint = Path(checkpoint).resolve() + output_root = Path(output_dir).resolve() + for filename in ("model.safetensors", "pytorch_model.bin"): + candidate = checkpoint / filename + if candidate.is_file(): + resolved = candidate.resolve() + if output_root not in resolved.parents or checkpoint not in resolved.parents: + raise ValueError("Selected_Checkpoint state resolves outside the run output") + return resolved + raise ValueError(f"Selected_Checkpoint has no model state: {checkpoint}") + + +def validate_selected_checkpoint(trainer, output_dir): + """Return the trainer-selected best-F1 checkpoint contained by output_dir.""" + selected = trainer.state.best_model_checkpoint + if not selected: + raise ValueError("Trainer state has no Selected_Checkpoint") + output_root = Path(output_dir).resolve() + selected_path = Path(selected).resolve() + if selected_path == output_root or output_root not in selected_path.parents: + raise ValueError("Selected_Checkpoint must be a checkpoint directory inside run output") + if not selected_path.is_dir(): + raise ValueError(f"Selected_Checkpoint does not exist: {selected_path}") + if trainer.state.best_metric is None or not math.isfinite(float(trainer.state.best_metric)): + raise ValueError("Trainer state has no finite best validation F1") + return selected_path + + +def stage_final_model(output_dir, model, tokenizer, artifact_config): + """Save one new self-contained local Final_Model staging directory.""" + from safetensors.torch import save_file + + output_dir = Path(output_dir) + stage = output_dir / "final-model.tmp" + destination = output_dir / "final-model" + if stage.exists() or destination.exists(): + raise ValueError("Final_Model staging and destination must both be absent") + stage.mkdir() + model.backbone.config.save_pretrained(str(stage)) + tokenizer.save_pretrained(str(stage)) + state = { + name: tensor.detach().cpu().contiguous() + for name, tensor in model.state_dict().items() + } + save_file(state, str(stage / "model.safetensors")) + write_json_atomic(stage / "train_config.json", artifact_config) + return stage + + +def _pad_metric_rows(predictions, labels): + if len(predictions) != len(labels) or not predictions: + raise ValueError("Ordered validation predictions and labels must be non-empty and equal") + width = max(len(row) for row in labels) + padded_predictions = [] + padded_labels = [] + for row, (prediction, label) in enumerate(zip(predictions, labels)): + if len(prediction) != len(label): + raise ValueError(f"Ordered validation row {row} lengths differ") + padded_predictions.append(list(prediction) + [IGNORE_INDEX] * (width - len(prediction))) + padded_labels.append(list(label) + [IGNORE_INDEX] * (width - len(label))) + return padded_predictions, padded_labels + + +def collect_ordered_validation_result(model, dataset, use_crf): + """Return exact discrete results and strict metrics in dataset order.""" + import torch + + try: + device = next(model.parameters()).device + except StopIteration: + device = torch.device("cpu") + model.eval() + predictions = [] + labels = [] + for example in dataset.examples: + inputs = { + name: torch.tensor([example[name]], device=device) + for name in ("input_ids", "attention_mask", "labels") + } + with torch.no_grad(): + output = model(**inputs) + if type(output) is not dict or "predictions" not in output or "word_labels" not in output: + raise ValueError("Ordered validation output must contain predictions and word_labels") + prediction_tensor = output["predictions"] + actual_tensor = output["word_labels"] + if not isinstance(prediction_tensor, torch.Tensor) or not isinstance( + actual_tensor, torch.Tensor + ): + raise TypeError("Ordered validation predictions and labels must be tensors") + if prediction_tensor.dim() != 2 or actual_tensor.dim() != 2: + raise ValueError("Ordered validation predictions and labels must have rank 2") + if prediction_tensor.shape != actual_tensor.shape or prediction_tensor.size(0) != 1: raise ValueError( - f"Saved tensor {name!r} has shape {tuple(tensor.shape)}, " - f"expected {tuple(expected[name].shape)}" + "Ordered validation predictions and labels must have equal single-row shapes" ) - if not torch.isfinite(tensor).all(): - raise ValueError(f"Saved tensor {name!r} contains non-finite values") + if prediction_tensor.dtype != torch.long or actual_tensor.dtype != torch.long: + raise TypeError("Ordered validation predictions and labels must use torch.long IDs") + prediction = prediction_tensor[0].detach().cpu().tolist() + actual = actual_tensor[0].detach().cpu().tolist() + predictions.append(prediction) + labels.append(actual) + metric_predictions, metric_labels = _pad_metric_rows(predictions, labels) + metrics = compute_metrics((metric_predictions, metric_labels), use_crf=use_crf) + return { + "prediction_ids": predictions, + "label_ids": labels, + "invalid_paths": metrics["invalid_paths"], + "metrics": metrics, + } + + +def compare_ordered_validation_results(selected, staged): + """Require exact ordered predictions, labels, invalid counts, and metrics.""" + for name in ("prediction_ids", "label_ids", "invalid_paths", "metrics"): + if selected.get(name) != staged.get(name): + raise ValueError(f"Ordered validation {name} differs after staged reload") + + +def _resolved_commit(component, name): + candidates = [ + getattr(component, "_commit_hash", None), + getattr(getattr(component, "config", None), "_commit_hash", None), + getattr(component, "init_kwargs", {}).get("_commit_hash") + if isinstance(getattr(component, "init_kwargs", None), dict) + else None, + ] + values = {value for value in candidates if value} + if len(values) != 1: + raise ValueError(f"{name} has no single resolved immutable revision") + value = values.pop() + if not isinstance(value, str) or not IMMUTABLE_REVISION.fullmatch(value): + raise ValueError(f"{name} resolved revision is not a full immutable commit") + return value + + +def resolve_tokenizer_revision(model_name, requested): + """Return the Hub revision for the tokenizer's requested model commit.""" + from huggingface_hub import model_info + + tokenizer_revision = model_info(model_name, revision=requested).sha + if not isinstance(tokenizer_revision, str) or not IMMUTABLE_REVISION.fullmatch( + tokenizer_revision + ): + raise ValueError("Tokenizer Hub revision is not a full immutable commit") + return tokenizer_revision + + +def resolve_model_identity(requested, tokenizer_revision, backbone_config): + """Require tokenizer and backbone identities to match the requested commit.""" + backbone_revision = _resolved_commit(backbone_config, "backbone") + if tokenizer_revision != requested or backbone_revision != requested: + raise ValueError( + "Requested, tokenizer, and backbone revisions are absent or inconsistent: " + f"{requested}, {tokenizer_revision}, {backbone_revision}" + ) + return requested + + +def promote_final_model( + stage, + destination, + marker, + manifest_path=None, + final_manifest=None, +): + """Atomically promote a stage, write SUCCESS last, and roll back failures.""" + stage = Path(stage) + destination = Path(destination) + if not stage.is_dir(): + raise ValueError(f"Final_Model stage does not exist: {stage}") + if destination.exists(): + raise ValueError(f"Final_Model destination already exists: {destination}") + if (manifest_path is None) != (final_manifest is None): + raise ValueError("manifest_path and final_manifest must be provided together") + os.replace(stage, destination) + try: + if manifest_path is not None: + write_json_atomic(manifest_path, final_manifest) + write_json_atomic(destination / "SUCCESS.json", marker) + validate_publishable_model(destination) + except Exception: + success_marker = destination / "SUCCESS.json" + if success_marker.exists(): + success_marker.unlink() + os.replace(destination, stage) + raise def validate_precision(precision): """Validate the selected training precision.""" import torch - if precision == "bf16" and not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()): + if precision == "bf16" and not ( + torch.cuda.is_available() and torch.cuda.is_bf16_supported() + ): raise ValueError("bf16 requires a CUDA device with BF16 support") def run_training(config): - """Train, validate, and save a required phrase model.""" + """Validate, train, verify, and transactionally publish a Final_Model.""" import torch - import transformers + from transformers import AutoConfig from transformers import AutoTokenizer from transformers import DataCollatorForTokenClassification from transformers import EarlyStoppingCallback @@ -466,13 +1604,14 @@ def run_training(config): from scancode_required_phrases.model import PhraseTrainer from scancode_required_phrases.model import build_optimizer - if config.with_isr and not config.evaluate_test: - raise ValueError("ISR evaluation requires --evaluate-test") - validate_config(config) validate_precision(config.precision) - prepare_output_dir(config.output_dir, config.resume) - set_seed(config.seed) + paths = { + "train": config.data_dir / "train.jsonl", + "validation": config.data_dir / "val.jsonl", + "test": config.data_dir / "test.jsonl", + } + records_by_split, raw_report = load_and_validate_splits(paths) tokenizer = AutoTokenizer.from_pretrained( config.model_name, @@ -481,161 +1620,272 @@ def run_training(config): ) if not tokenizer.is_fast: raise RuntimeError("Training requires a fast tokenizer with word IDs") - - paths = { - "train": config.data_dir / "train.jsonl", - "validation": config.data_dir / "val.jsonl", - "test": config.data_dir / "test.jsonl", - } - datasets = { - name: PhraseDataset(path, tokenizer, config.max_length, config.limit) - for name, path in paths.items() - } - validate_splits(datasets) - - for name, dataset in datasets.items(): - click.echo( - f"{name}: {len(dataset)} examples, {dataset.truncated} truncated, " - f"{dataset.cut_phrases} skipped with a cut phrase" - ) - - model = PhraseTagger(config) - collator = DataCollatorForTokenClassification( - tokenizer, - label_pad_token_id=IGNORE_INDEX, + tokenizer_revision = resolve_tokenizer_revision( + config.model_name, config.model_revision ) - - arguments = TrainingArguments( - output_dir=str(config.output_dir), - num_train_epochs=config.epochs, - per_device_train_batch_size=config.batch_size, - per_device_eval_batch_size=config.batch_size, - gradient_accumulation_steps=config.grad_accum, - learning_rate=config.base_lr, - weight_decay=config.weight_decay, - warmup_ratio=config.warmup_ratio, - lr_scheduler_type="cosine", - max_grad_norm=config.max_grad_norm, - eval_strategy="epoch", - save_strategy="epoch", - save_total_limit=2, - load_best_model_at_end=True, - metric_for_best_model="f1", - greater_is_better=True, - bf16=config.precision == "bf16", - fp16=False, - logging_steps=50, - report_to="none", - seed=config.seed, - data_seed=config.seed, - dataloader_num_workers=2, - save_safetensors=True, + backbone_config = AutoConfig.from_pretrained( + config.model_name, + revision=config.model_revision, ) - - trainer_kwargs = { - "model": model, - "args": arguments, - "train_dataset": datasets["train"], - "eval_dataset": datasets["validation"], - "data_collator": collator, - "compute_metrics": compute_metrics, - "optimizers": (build_optimizer(config, model), None), - "callbacks": [ - EarlyStoppingCallback( - early_stopping_patience=config.early_stopping_patience, - ) - ], - } - if "processing_class" in inspect.signature(PhraseTrainer.__init__).parameters: - trainer_kwargs["processing_class"] = tokenizer - else: - trainer_kwargs["tokenizer"] = tokenizer - - trainer = PhraseTrainer(**trainer_kwargs) - trainer.train(resume_from_checkpoint=config.resume or None) - - trainer.save_model(str(config.output_dir)) - tokenizer.save_pretrained(str(config.output_dir)) - - resolved_revision = getattr(model.backbone.config, "_commit_hash", None) - train_config = { - "model_name": config.model_name, - "model_revision": resolved_revision or config.model_revision, - "use_crf": config.use_crf, - "max_length": config.max_length, - "labels": LABELS, - } - (config.output_dir / "train_config.json").write_text( - json.dumps(train_config, indent=2), - encoding="utf-8", + resolved_revision = resolve_model_identity( + config.model_revision, tokenizer_revision, backbone_config ) - - model_path = config.output_dir / "model.safetensors" - validate_saved_state(model, model_path) - - validation_metrics = trainer.evaluate( - datasets["validation"], - metric_key_prefix="validation", + datasets, h2 = build_effective_datasets( + records_by_split, tokenizer, config.max_length, config.limit ) - test_metrics = None - if config.evaluate_test: - test_metrics = trainer.evaluate( - datasets["test"], - metric_key_prefix="test", - ) - if config.with_isr: - test_records = [ - validate_record(record, paths["test"], line_number) - for line_number, record in load_jsonl(paths["test"]) - ] - test_metrics["test_isr"] = evaluate_isr( - test_records, - model, - tokenizer, - config.max_length, - ) + validate_raw_split_hashes(paths, raw_report["h0"]) + dataset_report = { + **raw_report, + "h2": h2, + "effective": { + split: { + "accepted": len(dataset.effective_inventory), + "selected": len(dataset), + "truncations": dataset.truncations, + "rejections": dataset.rejections, + } + for split, dataset in datasets.items() + }, + } + prepare_output_dir(config.output_dir, config.resume) + source_provenance = collect_source_provenance() + runtime_provenance = collect_runtime_provenance(config.optimizer, config.precision) + report_path = config.output_dir / "dataset_report.json" + manifest_path = config.output_dir / "run_manifest.json" + write_json_atomic(report_path, dataset_report) manifest = { + "schema": MANIFEST_SCHEMA, + "state": "pre-run", "config": serializable_config(config), + "contracts": { + "artifact_schema": ARTIFACT_SCHEMA, + "constraint_contract": CONSTRAINT_CONTRACT, + "labels": list(LABELS), + }, "dataset": { - name: { - "path": str(path), - "sha256": sha256(path), - "examples": len(datasets[name]), - "truncated": datasets[name].truncated, - "cut_phrases": datasets[name].cut_phrases, - } - for name, path in paths.items() + "paths": {name: str(path) for name, path in paths.items()}, + "h0": raw_report["h0"], + "h1": raw_report["h1"], + "h2": h2, + "report": {"path": str(report_path), "sha256": sha256(report_path)}, }, - "model_revision": resolved_revision or config.model_revision, - "best_checkpoint": trainer.state.best_model_checkpoint, - "best_validation_f1": trainer.state.best_metric, - "validation_metrics": validation_metrics, - "test_metrics": test_metrics, - "log_history": trainer.state.log_history, - "versions": { - "python": platform.python_version(), - "torch": torch.__version__, - "transformers": transformers.__version__, - "scancode_toolkit": installed_version("scancode-toolkit"), - "pytorch_crf": installed_version("pytorch-crf"), + "model_identity": { + "name": config.model_name, + "requested_revision": config.model_revision, + "resolved_tokenizer_revision": resolved_revision, + "resolved_backbone_revision": resolved_revision, }, - "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None, + "source": source_provenance, + "runtime": runtime_provenance, + "completed_checks": [ + "configuration", + "raw-splits", + "identifier-uniqueness", + "content-duplicate-report", + "h0-h1", + "immutable-model-identity", + "effective-examples-h2", + ], } - (config.output_dir / "run_manifest.json").write_text( - json.dumps(manifest, indent=2), - encoding="utf-8", - ) + write_json_atomic(manifest_path, manifest) - click.echo(f"best checkpoint: {trainer.state.best_model_checkpoint}") + completed_checks = list(manifest["completed_checks"]) + phase = "model-construction" + stage = config.output_dir / "final-model.tmp" + destination = config.output_dir / "final-model" + try: + set_seed(config.seed) + model = PhraseTagger(config) + actual_revision = _resolved_commit(model.backbone.config, "constructed backbone") + if actual_revision != resolved_revision: + raise ValueError("Constructed backbone revision differs from the resolved revision") + completed_checks.append("model-construction") + collator = DataCollatorForTokenClassification( + tokenizer, + label_pad_token_id=IGNORE_INDEX, + ) + arguments = TrainingArguments( + output_dir=str(config.output_dir), + num_train_epochs=config.epochs, + per_device_train_batch_size=config.batch_size, + per_device_eval_batch_size=config.batch_size, + gradient_accumulation_steps=config.grad_accum, + learning_rate=config.base_lr, + weight_decay=config.weight_decay, + warmup_ratio=config.warmup_ratio, + lr_scheduler_type="cosine", + max_grad_norm=config.max_grad_norm, + eval_strategy="epoch", + save_strategy="epoch", + save_total_limit=2, + load_best_model_at_end=True, + metric_for_best_model="f1", + greater_is_better=True, + bf16=config.precision == "bf16", + fp16=False, + logging_steps=50, + report_to="none", + seed=config.seed, + data_seed=config.seed, + dataloader_num_workers=2, + save_safetensors=True, + ) + trainer_kwargs = { + "model": model, + "args": arguments, + "train_dataset": datasets["train"], + "eval_dataset": datasets["validation"], + "data_collator": collator, + "compute_metrics": partial(compute_metrics, use_crf=config.use_crf), + "optimizers": (build_optimizer(config, model), None), + "callbacks": [ + EarlyStoppingCallback( + early_stopping_patience=config.early_stopping_patience, + ) + ], + } + if "processing_class" in inspect.signature(PhraseTrainer.__init__).parameters: + trainer_kwargs["processing_class"] = tokenizer + else: + trainer_kwargs["tokenizer"] = tokenizer + + phase = "training" + trainer = PhraseTrainer(**trainer_kwargs) + trainer.train() + validate_raw_split_hashes(paths, raw_report["h0"]) + completed_checks.extend(["training", "post-training-dataset-hashes"]) + + phase = "selected-checkpoint" + selected_checkpoint = validate_selected_checkpoint(trainer, config.output_dir) + checkpoint_state = _load_state_file( + _checkpoint_state_path(selected_checkpoint, config.output_dir) + ) + validate_state_dicts( + _canonical_state_dict(checkpoint_state), + _canonical_state_dict(model.state_dict()), + "Selected_Checkpoint", + "selected in-memory model", + ) + completed_checks.extend(["selected-checkpoint", "checkpoint-in-memory-comparison"]) + selected_result = collect_ordered_validation_result( + model, datasets["validation"], config.use_crf + ) + selected_f1 = selected_result["metrics"]["f1"] + if float(trainer.state.best_metric) != selected_f1: + raise ValueError( + "Selected_Checkpoint best metric does not equal strict validation F1: " + f"{trainer.state.best_metric} != {selected_f1}" + ) + trainer.remove_callback(EarlyStoppingCallback) + validation_metrics = trainer.evaluate( + datasets["validation"], metric_key_prefix="validation" + ) + test_metrics = None + if config.evaluate_test: + test_metrics = trainer.evaluate(datasets["test"], metric_key_prefix="test") + if config.with_isr: + test_metrics["test_isr"] = evaluate_isr( + records_by_split["test"], model, tokenizer, config.max_length + ) + test_metrics["test_isr_scope"] = ( + "predicted-phrase locatability only; injection gates and rule mutation " + "are outside this metric" + ) + + phase = "staging" + artifact_config = _artifact_config(config, resolved_revision) + stage = stage_final_model( + config.output_dir, model, tokenizer, artifact_config + ) + staged_state = _load_state_file(stage / "model.safetensors") + validate_state_dicts( + _canonical_state_dict(model.state_dict()), + _canonical_state_dict(staged_state), + "selected in-memory model", + "staged model", + ) + completed_checks.extend(["staging", "in-memory-staged-comparison"]) + + phase = "offline-reload" + staged_model, _staged_tokenizer = _load_local_model(stage, offline=True) + validate_state_dicts( + _canonical_state_dict(model.state_dict()), + _canonical_state_dict(staged_model.state_dict()), + "selected in-memory model", + "offline staged model", + ) + staged_result = collect_ordered_validation_result( + staged_model, datasets["validation"], config.use_crf + ) + compare_ordered_validation_results(selected_result, staged_result) + completed_checks.extend(["offline-reload", "ordered-validation-comparison"]) + + phase = "final-manifest" + component_hashes = _all_file_hashes(stage) + completed_checks.append("artifact-hashes") + success_manifest = { + **manifest, + "state": "success", + "completed_checks": completed_checks + ["final-manifest"], + "selected_checkpoint": str(selected_checkpoint), + "best_validation_f1": trainer.state.best_metric, + "validation_metrics": validation_metrics, + "ordered_validation": selected_result, + "test_metrics": test_metrics, + "artifact_files": component_hashes, + "log_history": trainer.state.log_history, + } + write_json_atomic(stage / "run_manifest.json", success_manifest) + artifact_hashes = _all_file_hashes(stage) + marker = { + "schema": ARTIFACT_SCHEMA, + "constraint_contract": CONSTRAINT_CONTRACT, + "selected_checkpoint": str(selected_checkpoint), + "run_manifest_sha256": sha256(stage / "run_manifest.json"), + "files": artifact_hashes, + } + + phase = "promotion" + promote_final_model( + stage, + destination, + marker, + manifest_path=manifest_path, + final_manifest=success_manifest, + ) + phase = "success-marker" + except Exception as error: + if destination.exists(): + if stage.exists(): + raise RuntimeError( + "Both failed Final_Model destination and staging directory exist" + ) from error + success_marker = destination / "SUCCESS.json" + if success_marker.exists(): + success_marker.unlink() + os.replace(destination, stage) + retained_artifacts = [config.output_dir, report_path, stage] + retained_artifacts.extend(sorted(config.output_dir.glob("checkpoint-*"))) + write_failure_manifest( + manifest_path, + manifest, + phase, + error, + completed_checks, + retained_artifacts, + ) + raise + + click.echo(f"best checkpoint: {selected_checkpoint}") click.echo(f"best validation F1: {trainer.state.best_metric}") click.echo(f"validation: {validation_metrics}") if test_metrics is not None: click.echo(f"test: {test_metrics}") - return { "validation": validation_metrics, "test": test_metrics, + "final_model": destination, } @@ -650,75 +1900,51 @@ def run_training(config): "--output-dir", default="model-output", type=click.Path(file_okay=False, path_type=Path), - help="Directory for checkpoints and the final model.", + help="New or empty directory for this run.", ) @click.option("--model-name", default=MODEL_NAME, help="Base model to fine-tune.") -@click.option("--model-revision", default=None, help="Optional model revision to pin.") @click.option( - "--max-length", - default=MAX_LENGTH, - type=click.IntRange(min=3, max=MAX_LENGTH), - show_default=True, + "--model-revision", + required=True, + help="Full immutable model and tokenizer commit revision.", ) -@click.option("--epochs", default=8, type=click.IntRange(min=1), show_default=True) -@click.option("--batch-size", default=1, type=int, show_default=True) -@click.option("--grad-accum", default=16, type=int, show_default=True) -@click.option("--base-lr", default=2e-5, type=float, show_default=True) -@click.option("--head-lr", default=1e-4, type=float, show_default=True) -@click.option("--aux-ce-weight", default=0.3, type=float, show_default=True) @click.option( - "--optimizer", - type=click.Choice(["adamw", "adamw-8bit"]), - default="adamw", - show_default=True, + "--max-length", default=MAX_LENGTH, + type=click.IntRange(min=3, max=MAX_LENGTH), show_default=True, ) +@click.option("--epochs", default=8, type=click.IntRange(min=1), show_default=True) +@click.option("--batch-size", default=1, type=click.IntRange(min=1), show_default=True) +@click.option("--grad-accum", default=16, type=click.IntRange(min=1), show_default=True) +@click.option("--base-lr", default=2e-5, type=click.FloatRange(min=0, min_open=True), show_default=True) +@click.option("--head-lr", default=1e-4, type=click.FloatRange(min=0, min_open=True), show_default=True) +@click.option("--aux-ce-weight", default=0.3, type=click.FloatRange(min=0), show_default=True) @click.option( - "--precision", - type=click.Choice(["fp32", "bf16"]), - default="fp32", - show_default=True, + "--optimizer", type=click.Choice(["adamw", "adamw-8bit"]), + default="adamw", show_default=True, ) -@click.option("--no-crf", is_flag=True, default=False, help="Train without the CRF head.") @click.option( - "--evaluate-test", - is_flag=True, - default=False, - help="Evaluate the test split after training.", + "--precision", type=click.Choice(["fp32", "bf16"]), + default="fp32", show_default=True, ) +@click.option("--no-crf", is_flag=True, default=False, help="Train without the CRF head.") +@click.option("--evaluate-test", is_flag=True, help="Evaluate test data after selection.") @click.option( - "--with-isr", - is_flag=True, - default=False, - help="Report injection success rate with final test evaluation.", + "--with-isr", is_flag=True, + help="Report predicted-phrase locatability during final test evaluation.", ) -@click.option("--limit", default=0, type=int, help="Limit examples per split; zero uses all.") -@click.option("--resume", is_flag=True, default=False, help="Resume from the latest checkpoint.") -@click.option("--seed", default=42, type=int, show_default=True) +@click.option("--limit", default=0, type=click.IntRange(min=0), help="Limit effective examples per split.") +@click.option("--resume", is_flag=True, hidden=True, help="Unsupported.") +@click.option("--seed", default=42, type=click.IntRange(min=0), show_default=True) def main( - data_dir, - output_dir, - model_name, - model_revision, - max_length, - epochs, - batch_size, - grad_accum, - base_lr, - head_lr, - aux_ce_weight, - optimizer, - precision, - no_crf, - evaluate_test, - with_isr, - limit, - resume, - seed, + data_dir, output_dir, model_name, model_revision, max_length, epochs, + batch_size, grad_accum, base_lr, head_lr, aux_ce_weight, optimizer, + precision, no_crf, evaluate_test, with_isr, limit, resume, seed, ): - """Train the required phrase tagger from a BIOES dataset.""" + """Train the required phrase tagger from a positive BIOES dataset.""" if with_isr and not evaluate_test: raise click.UsageError("--with-isr requires --evaluate-test") - + if resume: + raise click.UsageError("--resume is unsupported for final hardened training") config = Config( data_dir=data_dir, output_dir=output_dir, diff --git a/tests/test_export.py b/tests/test_export.py index 174ebd4..59d75ba 100644 --- a/tests/test_export.py +++ b/tests/test_export.py @@ -3,13 +3,161 @@ # Copyright (c) nexB Inc. and others. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import builtins +import json +from pathlib import Path +import sys +import types + import pytest +from click.testing import CliRunner np = pytest.importorskip("numpy") +from scancode_required_phrases import export as export_module +from scancode_required_phrases import training +from scancode_required_phrases.export import check_viterbi_matches_crf +from scancode_required_phrases.export import export_crf_matrices +from scancode_required_phrases.export import export_onnx_emissions +from scancode_required_phrases.export import main +from scancode_required_phrases.export import _load_publishable_artifact +from scancode_required_phrases.export import OnnxDependencyError +from scancode_required_phrases.export import PARITY_CASES from scancode_required_phrases.export import sha256 from scancode_required_phrases.export import viterbi_decode +LABELS = ("O", "B-REQ", "I-REQ", "E-REQ", "S-REQ") +ARTIFACT_SCHEMA = "test-artifact-v2" +CONSTRAINT_CONTRACT = "bioes-hard-v1" + + +def validate_bioes(labels): + starts = {"O", "B-REQ", "S-REQ"} + ends = {"O", "E-REQ", "S-REQ"} + transitions = { + "O": {"O", "B-REQ", "S-REQ"}, + "B-REQ": {"I-REQ", "E-REQ"}, + "I-REQ": {"I-REQ", "E-REQ"}, + "E-REQ": {"O", "B-REQ", "S-REQ"}, + "S-REQ": {"O", "B-REQ", "S-REQ"}, + } + if labels[0] not in starts: + return f"starts with {labels[0]}" + for previous, current in zip(labels, labels[1:]): + if current not in transitions[previous]: + return f"contains invalid transition {previous} -> {current}" + if labels[-1] not in ends: + return f"ends with {labels[-1]}" + return None + + +def make_masks(torch): + start = torch.tensor([True, True, False, False, True]) + transition = torch.tensor( + [ + [True, True, False, False, True], + [False, False, True, True, False], + [False, False, True, True, False], + [True, True, False, False, True], + [True, True, False, False, True], + ] + ) + end = torch.tensor([True, False, False, True, True]) + return start, transition, end + + +class FakeConstrainedCRF: + def __init__(self, torch): + self.torch = torch + self.start_transitions = torch.nn.Parameter(torch.zeros(5)) + self.transitions = torch.nn.Parameter(torch.zeros((5, 5))) + self.end_transitions = torch.nn.Parameter(torch.zeros(5)) + self.start_mask, self.transition_mask, self.end_mask = make_masks(torch) + + @property + def effective_start_transitions(self): + return self.start_transitions.masked_fill(~self.start_mask, -self.torch.inf) + + @property + def effective_transitions(self): + return self.transitions.masked_fill(~self.transition_mask, -self.torch.inf) + + @property + def effective_end_transitions(self): + return self.end_transitions.masked_fill(~self.end_mask, -self.torch.inf) + + def decode(self, emissions, mask): + start = self.effective_start_transitions + transitions = self.effective_transitions + end = self.effective_end_transitions + paths = [] + for row in range(emissions.shape[0]): + length = int(mask[row].sum().item()) + score = start + emissions[row, 0] + backpointers = [] + for step in range(1, length): + candidates = score[:, None] + transitions + score, sources = candidates.max(dim=0) + score = score + emissions[row, step] + backpointers.append(sources) + best = int((score + end).argmax()) + path = [best] + for sources in reversed(backpointers): + best = int(sources[best]) + path.append(best) + paths.append(list(reversed(path))) + return paths + + +class FakeTagger: + def __init__(self, torch): + self.use_crf = True + self.num_labels = len(LABELS) + self.crf = FakeConstrainedCRF(torch) + + def eval(self): + return self + + +def install_artifact_helpers(monkeypatch, model_dir, tagger): + from scancode_required_phrases import training + + config = { + "artifact_schema": ARTIFACT_SCHEMA, + "constraint_contract": CONSTRAINT_CONTRACT, + "labels": list(LABELS), + "use_crf": True, + } + marker = { + "schema": ARTIFACT_SCHEMA, + "files": {"model.safetensors": "validated-test-hash"}, + } + (model_dir / "train_config.json").write_text(json.dumps(config), encoding="utf-8") + (model_dir / "SUCCESS.json").write_text(json.dumps(marker), encoding="utf-8") + validation_calls = [] + + def validate_publishable_model(path): + validation_calls.append(path) + + def load_final_model(path, offline): + assert path == model_dir + assert offline is True + return tagger, object() + + def write_json_atomic(path, value): + path.write_text(json.dumps(value, sort_keys=True) + "\n", encoding="utf-8") + + monkeypatch.setattr(training, "ARTIFACT_SCHEMA", ARTIFACT_SCHEMA, raising=False) + monkeypatch.setattr(training, "CONSTRAINT_CONTRACT", CONSTRAINT_CONTRACT, raising=False) + monkeypatch.setattr(training, "LABELS", LABELS) + monkeypatch.setattr(training, "validate_bioes", validate_bioes) + monkeypatch.setattr( + training, "validate_publishable_model", validate_publishable_model, raising=False + ) + monkeypatch.setattr(training, "load_final_model", load_final_model, raising=False) + monkeypatch.setattr(training, "write_json_atomic", write_json_atomic, raising=False) + return config, marker, validation_calls + def test_viterbi_with_zero_transitions_is_argmax(): emissions = np.array( @@ -22,7 +170,7 @@ def test_viterbi_with_zero_transitions_is_argmax(): transitions = np.zeros((5, 5)) edges = np.zeros(5) - assert viterbi_decode(emissions, edges, transitions, edges) == [1, 0, 4] + assert viterbi_decode(emissions, edges, transitions, edges, 5) == [1, 0, 4] def test_viterbi_obeys_transition_scores(): @@ -35,6 +183,330 @@ def test_viterbi_obeys_transition_scores(): assert path[0] == path[1] +def test_viterbi_ties_are_deterministic_and_forbidden_edges_cannot_win(): + emissions = np.zeros((2, 2), dtype=np.float32) + transitions = np.array([[0.0, -np.inf], [0.0, 0.0]], dtype=np.float32) + edges = np.zeros(2, dtype=np.float32) + + assert viterbi_decode(emissions, edges, transitions, edges) == [0, 0] + + +@pytest.mark.parametrize( + ("argument", "value", "error"), + [ + ("emissions", np.empty((0, 2)), "at least one token"), + ("emissions", np.zeros((1, 2, 1)), "emissions must have shape"), + ("emissions", [[0.0, 1.0]], "NumPy array"), + ("emissions", np.zeros((1, 2), dtype=np.int64), "floating-point"), + ("emissions", np.array([[np.nan, 0.0]]), "NaN"), + ("start", np.zeros(3), "shape"), + ("transitions", np.zeros((2, 3)), "shape"), + ("end", np.zeros(3), "shape"), + ], +) +def test_viterbi_rejects_malformed_inputs(argument, value, error): + emissions = np.zeros((2, 2)) + start = np.zeros(2) + transitions = np.zeros((2, 2)) + end = np.zeros(2) + values = { + "emissions": emissions, + "start": start, + "transitions": transitions, + "end": end, + } + values[argument] = value + + with pytest.raises((TypeError, ValueError), match=error): + viterbi_decode( + values["emissions"], + values["start"], + values["transitions"], + values["end"], + num_labels=2, + ) + + +def test_viterbi_rejects_wrong_explicit_label_count(): + emissions = np.zeros((1, 2)) + edges = np.zeros(2) + + with pytest.raises(ValueError, match="expected 3"): + viterbi_decode(emissions, edges, np.zeros((2, 2)), edges, num_labels=3) + + +def test_deterministic_adversarial_numpy_pytorch_parity(monkeypatch): + torch = pytest.importorskip("torch") + tagger = FakeTagger(torch) + with torch.no_grad(): + tagger.crf.start_transitions.copy_(torch.tensor([0.7, -0.2, 0.4, 0.8, 1.1])) + tagger.crf.end_transitions.copy_(torch.tensor([-0.3, 0.6, 0.2, -0.9, 0.5])) + monkeypatch.setattr( + "scancode_required_phrases.export._training_interfaces", + lambda: ( + ARTIFACT_SCHEMA, + CONSTRAINT_CONTRACT, + LABELS, + None, + validate_bioes, + None, + None, + ), + ) + + start, transitions, end = check_viterbi_matches_crf(tagger, len(LABELS)) + + assert start.shape == (5,) + assert transitions.shape == (5, 5) + assert end.shape == (5,) + + +def test_crf_export_is_publishable_effective_and_manifested(tmp_path, monkeypatch): + torch = pytest.importorskip("torch") + model_dir = tmp_path / "final-model" + output_dir = tmp_path / "export" + model_dir.mkdir() + tagger = FakeTagger(torch) + _config, marker, validation_calls = install_artifact_helpers(monkeypatch, model_dir, tagger) + + matrix_path, manifest_path = export_crf_matrices(model_dir, output_dir) + + assert validation_calls == [model_dir] + with np.load(matrix_path) as matrices: + assert np.isneginf(matrices["start"][[2, 3]]).all() + assert np.isneginf(matrices["transitions"][0, [2, 3]]).all() + assert np.isneginf(matrices["end"][[1, 2]]).all() + assert np.isfinite(matrices["start"][[0, 1, 4]]).all() + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + assert manifest["artifact_schema"] == ARTIFACT_SCHEMA + assert manifest["constraint_contract"] == CONSTRAINT_CONTRACT + assert manifest["labels"] == list(LABELS) + assert manifest["source_final_model"] == marker + assert manifest["crf_transitions"] == sha256(matrix_path) + assert manifest["parity"] == {"passed": True, "cases": list(PARITY_CASES)} + assert "onnx_model" not in manifest + + +@pytest.mark.parametrize( + ("field", "value", "error"), + [ + ("artifact_schema", "old", "artifact schema"), + ("labels", list(reversed(LABELS)), "label order"), + ("constraint_contract", "old", "constraint contract"), + ("use_crf", False, "with a CRF"), + ], +) +def test_crf_export_rejects_unsupported_artifact(tmp_path, monkeypatch, field, value, error): + torch = pytest.importorskip("torch") + model_dir = tmp_path / "final-model" + model_dir.mkdir() + tagger = FakeTagger(torch) + config, _marker, _calls = install_artifact_helpers(monkeypatch, model_dir, tagger) + config[field] = value + (model_dir / "train_config.json").write_text(json.dumps(config), encoding="utf-8") + + with pytest.raises(ValueError, match=error): + export_crf_matrices(model_dir, tmp_path / "export") + + +def test_generic_publishable_load_allows_non_crf_onnx_artifacts(tmp_path, monkeypatch): + torch = pytest.importorskip("torch") + model_dir = tmp_path / "final-model" + model_dir.mkdir() + + class NonCrfTagger: + use_crf = False + + def eval(self): + return self + + tagger = NonCrfTagger() + config, _marker, _calls = install_artifact_helpers(monkeypatch, model_dir, tagger) + config["use_crf"] = False + (model_dir / "train_config.json").write_text(json.dumps(config), encoding="utf-8") + + loaded, _tokenizer, loaded_config, _marker = _load_publishable_artifact(model_dir) + assert loaded is tagger + assert loaded_config["use_crf"] is False + with pytest.raises(ValueError, match="with a CRF"): + _load_publishable_artifact(model_dir, require_crf=True) + + +def test_crf_export_rejects_output_inside_final_model(tmp_path, monkeypatch): + torch = pytest.importorskip("torch") + model_dir = tmp_path / "final-model" + model_dir.mkdir() + install_artifact_helpers(monkeypatch, model_dir, FakeTagger(torch)) + + with pytest.raises(ValueError, match="outside the Final_Model"): + export_crf_matrices(model_dir, model_dir / "export") + + assert not (model_dir / "export").exists() + + +def test_crf_export_rejects_a_nonempty_output_directory(tmp_path, monkeypatch): + torch = pytest.importorskip("torch") + model_dir = tmp_path / "final-model" + output_dir = tmp_path / "export" + model_dir.mkdir() + output_dir.mkdir() + (output_dir / "model.onnx").write_bytes(b"stale") + install_artifact_helpers(monkeypatch, model_dir, FakeTagger(torch)) + + with pytest.raises(ValueError, match="not empty"): + export_crf_matrices(model_dir, output_dir) + + assert (output_dir / "model.onnx").read_bytes() == b"stale" + assert not (output_dir / "crf_transitions.npz").exists() + + +def test_failed_crf_export_removes_staging_and_never_publishes(tmp_path, monkeypatch): + torch = pytest.importorskip("torch") + model_dir = tmp_path / "final-model" + output_dir = tmp_path / "export" + model_dir.mkdir() + install_artifact_helpers(monkeypatch, model_dir, FakeTagger(torch)) + + def fail_manifest(path, value): + raise OSError("injected export manifest failure") + + monkeypatch.setattr(training, "write_json_atomic", fail_manifest) + with pytest.raises(OSError, match="injected export manifest failure"): + export_crf_matrices(model_dir, output_dir) + + assert not output_dir.exists() + assert not output_dir.with_name("export.tmp").exists() + + +def test_crf_export_does_not_import_onnx_and_onnx_error_is_actionable(tmp_path, monkeypatch): + torch = pytest.importorskip("torch") + model_dir = tmp_path / "final-model" + model_dir.mkdir() + install_artifact_helpers(monkeypatch, model_dir, FakeTagger(torch)) + original_import = builtins.__import__ + + def without_onnx(name, *args, **kwargs): + if name in {"onnx", "onnxruntime"}: + raise ImportError(f"blocked {name}") + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", without_onnx) + + matrix_path, manifest_path = export_crf_matrices(model_dir, tmp_path / "crf") + assert matrix_path.exists() + assert manifest_path.exists() + with pytest.raises(OnnxDependencyError, match=r"\[training,onnx\]"): + export_onnx_emissions(model_dir, tmp_path / "onnx") + + +def test_onnx_export_validates_shape_and_publishes_transactionally(tmp_path, monkeypatch): + torch = pytest.importorskip("torch") + + class Tokenizer: + def __call__(self, text, return_tensors): + assert return_tensors == "pt" + return { + "input_ids": torch.zeros((1, 3), dtype=torch.long), + "attention_mask": torch.ones((1, 3), dtype=torch.long), + } + + class Emissions: + def __call__(self, input_ids, attention_mask): + return torch.zeros((1, 3, len(LABELS)), dtype=torch.float32) + + runtime = types.ModuleType("onnxruntime") + runtime.output = np.zeros((1, 1, len(LABELS)), dtype=np.float32) + + class Session: + def __init__(self, path, providers): + self.path = path + self.providers = providers + + def run(self, names, feeds): + return [runtime.output] + + runtime.InferenceSession = Session + monkeypatch.setitem(sys.modules, "onnx", types.ModuleType("onnx")) + monkeypatch.setitem(sys.modules, "onnxruntime", runtime) + monkeypatch.setattr( + export_module, + "_load_publishable_artifact", + lambda model_dir: ( + object(), + Tokenizer(), + {"resolved_model_revision": "a" * 40}, + {"schema": ARTIFACT_SCHEMA, "files": {}}, + ), + ) + monkeypatch.setattr(export_module, "build_emissions_module", lambda tagger: Emissions()) + monkeypatch.setattr( + torch.onnx, + "export", + lambda module, inputs, path, **kwargs: Path(path).write_bytes(b"onnx"), + ) + model_dir = tmp_path / "final-model" + model_dir.mkdir() + + failed_output = tmp_path / "failed-onnx" + with pytest.raises(ValueError, match="shape"): + export_onnx_emissions(model_dir, failed_output) + assert not failed_output.exists() + assert not failed_output.with_name("failed-onnx.tmp").exists() + + runtime.output = np.zeros((1, 3, len(LABELS)), dtype=np.float32) + output_dir = tmp_path / "onnx" + onnx_path, manifest_path = export_onnx_emissions(model_dir, output_dir, opset=17) + assert onnx_path.is_file() + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + assert manifest["onnx_model"] == sha256(onnx_path) + assert manifest["opset"] == 17 + assert not output_dir.with_name("onnx.tmp").exists() + + +def test_cli_exposes_crf_and_onnx_as_separate_operations(tmp_path, monkeypatch): + model_dir = tmp_path / "final-model" + model_dir.mkdir() + calls = [] + + def export_crf(model, output): + calls.append(("crf", model, output)) + return output / "crf_transitions.npz", output / "manifest.json" + + def export_onnx(model, output, opset): + calls.append(("onnx", model, output, opset)) + return output / "model.onnx", output / "manifest.json" + + monkeypatch.setattr("scancode_required_phrases.export.export_crf_matrices", export_crf) + monkeypatch.setattr("scancode_required_phrases.export.export_onnx_emissions", export_onnx) + runner = CliRunner() + + crf_output = tmp_path / "crf" + result = runner.invoke( + main, + ["--model-dir", str(model_dir), "--output-dir", str(crf_output)], + ) + assert result.exit_code == 0 + assert calls == [("crf", model_dir, crf_output)] + + calls.clear() + onnx_output = tmp_path / "onnx" + result = runner.invoke( + main, + [ + "--model-dir", + str(model_dir), + "--output-dir", + str(onnx_output), + "--operation", + "onnx", + "--opset", + "17", + ], + ) + assert result.exit_code == 0 + assert calls == [("onnx", model_dir, onnx_output, 17)] + + def test_sha256_is_stable(tmp_path): path = tmp_path / "model.onnx" path.write_bytes(b"model") diff --git a/tests/test_model.py b/tests/test_model.py index b7ae65f..1d64105 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -15,8 +15,11 @@ pytest.importorskip("transformers") from scancode_required_phrases import model as model_module +from scancode_required_phrases.model import build_constraint_masks from scancode_required_phrases.model import build_optimizer +from scancode_required_phrases.model import ConstrainedCRF from scancode_required_phrases.model import PhraseTagger +from scancode_required_phrases.model import PhraseTrainer from scancode_required_phrases.training import LABEL2ID from scancode_required_phrases.training import LABELS @@ -32,12 +35,10 @@ def __init__(self): self.encoder = torch.nn.Module() self.encoder.layer = torch.nn.ModuleList([torch.nn.Linear(4, 4)]) self.embeddings = torch.nn.Linear(4, 4) + self.gradient_checkpointing_kwargs = None - def gradient_checkpointing_enable(self): - pass - - def enable_input_require_grads(self): - pass + def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): + self.gradient_checkpointing_kwargs = gradient_checkpointing_kwargs @pytest.fixture @@ -78,9 +79,45 @@ def from_pretrained(*args, **kwargs): return FakeBackbone() monkeypatch.setattr(model_module.AutoModel, "from_pretrained", from_pretrained) - PhraseTagger(config) + tagger = PhraseTagger(config) assert calls == [(("fake-model",), {"revision": "revision"})] + assert tagger.backbone.gradient_checkpointing_kwargs == {"use_reentrant": False} + + +def test_trainer_accepts_gathered_finite_losses_and_rejects_non_finite_losses(): + trainer = PhraseTrainer.__new__(PhraseTrainer) + + class Model: + def __init__(self, loss): + self.loss = loss + + def __call__(self, **inputs): + return {"loss": self.loss} + + loss = torch.tensor([1.0, 2.0], requires_grad=True) + assert torch.equal(trainer.compute_loss(Model(loss), {}), loss) + + with pytest.raises(FloatingPointError, match="non-finite loss"): + trainer.compute_loss(Model(torch.tensor([1.0, torch.inf])), {}) + + +def test_prediction_step_averages_a_gathered_loss(monkeypatch): + trainer = PhraseTrainer.__new__(PhraseTrainer) + monkeypatch.setattr(trainer, "_prepare_inputs", lambda inputs: inputs) + + class Model: + def __call__(self, **inputs): + return { + "loss": torch.tensor([2.0, 4.0]), + "predictions": torch.tensor([[0], [0]]), + "word_labels": torch.tensor([[0], [0]]), + } + + loss, _, _ = trainer.prediction_step(Model(), {}, prediction_loss_only=False) + + assert loss.ndim == 0 + assert loss.item() == 3.0 def test_build_optimizer_uses_explicit_adamw(monkeypatch, config): @@ -100,13 +137,12 @@ def test_build_optimizer_uses_explicit_adamw(monkeypatch, config): def make_crf_tagger(): - from torchcrf import CRF - tagger = PhraseTagger.__new__(PhraseTagger) torch.nn.Module.__init__(tagger) tagger.use_crf = True + tagger.aux_ce_weight = 0 tagger.num_labels = len(LABELS) - tagger.crf = CRF(len(LABELS), batch_first=True) + tagger.crf = ConstrainedCRF(len(LABELS), batch_first=True) with torch.no_grad(): for parameter in tagger.crf.parameters(): parameter.zero_() @@ -131,9 +167,192 @@ def test_predict_words_uses_first_subwords(): assert tags == [LABEL2ID["B-REQ"], LABEL2ID["E-REQ"]] -def test_predict_words_does_not_run_for_an_empty_sequence(): +def test_predict_words_rejects_an_empty_sequence(): tagger = make_crf_tagger() tagger.emissions = lambda *args: pytest.fail("emissions should not be computed") input_ids = torch.zeros((1, 2), dtype=torch.long) - assert tagger.predict_words(input_ids, input_ids, [None, None]) == [] + with pytest.raises(ValueError, match="must not be empty"): + tagger.predict_words(input_ids, input_ids, [None, None]) + + +def test_forward_and_predict_words_share_constrained_decode(): + tagger = make_crf_tagger() + tagger.eval() + emissions = torch.zeros((1, 1, len(LABELS))) + emissions[0, 0, LABEL2ID["I-REQ"]] = 1000.0 + tagger.emissions = lambda input_ids, attention_mask: emissions + input_ids = torch.zeros((1, 1), dtype=torch.long) + labels = torch.tensor([[LABEL2ID["O"]]]) + + result = tagger(input_ids, input_ids, labels=labels) + predicted = tagger.predict_words(input_ids, input_ids, [0]) + + assert result["predictions"].tolist() == [[LABEL2ID["O"]]] + assert predicted == [LABEL2ID["O"]] + assert torch.isfinite(result["loss"]) + + +def test_constraint_masks_match_exact_bioes_contract(): + start_mask, transition_mask, end_mask = build_constraint_masks() + + assert start_mask.tolist() == [True, True, False, False, True] + assert end_mask.tolist() == [True, False, False, True, True] + assert transition_mask.tolist() == [ + [True, True, False, False, True], + [False, False, True, True, False], + [False, False, True, True, False], + [True, True, False, False, True], + [True, True, False, False, True], + ] + + +def test_constraint_masks_are_not_saved_and_learned_scores_stay_finite(): + crf = ConstrainedCRF(len(LABELS), batch_first=True) + + assert set(crf.state_dict()) == { + "start_transitions", + "transitions", + "end_transitions", + } + assert all(parameter.requires_grad for parameter in crf.parameters()) + assert all(torch.isfinite(parameter).all() for parameter in crf.parameters()) + assert torch.isneginf(crf.effective_start_transitions[~crf.start_mask]).all() + assert torch.isneginf(crf.effective_transitions[~crf.transition_mask]).all() + assert torch.isneginf(crf.effective_end_transitions[~crf.end_mask]).all() + + +def test_constrained_decode_cannot_select_an_invalid_single_tag_path(): + crf = ConstrainedCRF(len(LABELS), batch_first=True) + with torch.no_grad(): + for parameter in crf.parameters(): + parameter.zero_() + emissions = torch.zeros((1, 1, len(LABELS))) + emissions[0, 0, LABEL2ID["I-REQ"]] = 1000.0 + emissions[0, 0, LABEL2ID["E-REQ"]] = 900.0 + + assert crf.decode(emissions) == [[LABEL2ID["O"]]] + + +def test_constrained_loss_rejects_an_invalid_gold_path(): + crf = ConstrainedCRF(len(LABELS), batch_first=True) + emissions = torch.zeros((1, 1, len(LABELS))) + tags = torch.tensor([[LABEL2ID["I-REQ"]]]) + + with pytest.raises(ValueError, match="Invalid BIOES gold path in row 0"): + crf(emissions, tags) + + +def test_constrained_crf_padding_is_ignored_without_non_finite_loss(): + crf = ConstrainedCRF(len(LABELS), batch_first=True) + with torch.no_grad(): + for parameter in crf.parameters(): + parameter.zero_() + active_emissions = torch.tensor( + [ + [ + [0.0, 4.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 4.0, 0.0], + ] + ] + ) + active_tags = torch.tensor([[LABEL2ID["B-REQ"], LABEL2ID["E-REQ"]]]) + active_mask = torch.ones((1, 2), dtype=torch.bool) + + padded_emissions = torch.cat( + [active_emissions, torch.full((1, 1, len(LABELS)), 1000.0)], + dim=1, + ) + padded_tags = torch.tensor( + [[LABEL2ID["B-REQ"], LABEL2ID["E-REQ"], len(LABELS) + 10]] + ) + padded_mask = torch.tensor([[True, True, False]]) + + active_likelihood = crf( + active_emissions, + active_tags, + mask=active_mask, + reduction="none", + ) + padded_likelihood = crf( + padded_emissions, + padded_tags, + mask=padded_mask, + reduction="none", + ) + + assert torch.isfinite(padded_likelihood).all() + assert torch.equal(active_likelihood, padded_likelihood) + assert crf.decode(active_emissions, active_mask) == crf.decode( + padded_emissions, + padded_mask, + ) + + +def test_constrained_crf_rejects_invalid_inputs(): + crf = ConstrainedCRF(len(LABELS), batch_first=True) + emissions = torch.zeros((1, 2, len(LABELS))) + tags = torch.tensor([[LABEL2ID["B-REQ"], LABEL2ID["E-REQ"]]]) + + with pytest.raises(ValueError, match="non-empty"): + crf.decode(torch.zeros((1, 0, len(LABELS)))) + with pytest.raises(TypeError, match="boolean"): + crf.decode(emissions, torch.ones((1, 2), dtype=torch.long)) + three_emissions = torch.zeros((1, 3, len(LABELS))) + with pytest.raises(ValueError, match="left aligned"): + crf.decode(three_emissions, torch.tensor([[True, False, True]])) + with pytest.raises(ValueError, match="mask shape"): + crf.decode(emissions, torch.ones((1, 1), dtype=torch.bool)) + with pytest.raises(ValueError, match="valid label IDs"): + crf(emissions, torch.tensor([[LABEL2ID["B-REQ"], len(LABELS)]])) + with pytest.raises(FloatingPointError, match="finite scores"): + invalid_emissions = emissions.clone() + invalid_emissions[0, 0, 0] = torch.nan + crf.decode(invalid_emissions) + + with torch.no_grad(): + crf.transitions[0, 0] = torch.inf + with pytest.raises(FloatingPointError, match="transitions"): + crf(emissions, tags) + + +def test_gather_words_uses_o_for_padded_crf_tags(): + tagger = make_crf_tagger() + emissions = torch.zeros((2, 3, len(LABELS))) + labels = torch.tensor( + [ + [LABEL2ID["S-REQ"], -100, -100], + [LABEL2ID["B-REQ"], LABEL2ID["E-REQ"], -100], + ] + ) + + _, crf_tags, _, mask = tagger.gather_words(emissions, labels) + + assert crf_tags.tolist() == [ + [LABEL2ID["S-REQ"], LABEL2ID["O"]], + [LABEL2ID["B-REQ"], LABEL2ID["E-REQ"]], + ] + assert mask.tolist() == [[True, False], [True, True]] + + +def test_gather_words_rejects_an_empty_word_sequence(): + tagger = make_crf_tagger() + emissions = torch.zeros((1, 2, len(LABELS))) + labels = torch.full((1, 2), -100) + + with pytest.raises(ValueError, match="must not be empty"): + tagger.gather_words(emissions, labels) + + +def test_phrase_tagger_accepts_an_already_constructed_backbone(monkeypatch, config): + monkeypatch.setattr( + model_module.AutoModel, + "from_pretrained", + lambda *args, **kwargs: pytest.fail("local construction must not load a model"), + ) + backbone = FakeBackbone() + + tagger = PhraseTagger(config, backbone=backbone) + + assert tagger.backbone is backbone + assert isinstance(tagger.crf, ConstrainedCRF) diff --git a/tests/test_training.py b/tests/test_training.py index ceb9801..a9f1dee 100644 --- a/tests/test_training.py +++ b/tests/test_training.py @@ -4,34 +4,57 @@ # SPDX-License-Identifier: Apache-2.0 import json +from pathlib import Path +import types + from click.testing import CliRunner import pytest +from scancode_required_phrases import training from scancode_required_phrases.training import align_labels +from scancode_required_phrases.training import AlignmentError +from scancode_required_phrases.training import ARTIFACT_SCHEMA +from scancode_required_phrases.training import build_effective_datasets +from scancode_required_phrases.training import compare_ordered_validation_results from scancode_required_phrases.training import compute_metrics from scancode_required_phrases.training import Config +from scancode_required_phrases.training import CONSTRAINT_CONTRACT from scancode_required_phrases.training import decode_row from scancode_required_phrases.training import extract_spans from scancode_required_phrases.training import first_subword_positions from scancode_required_phrases.training import IGNORE_INDEX from scancode_required_phrases.training import LABEL2ID +from scancode_required_phrases.training import LABELS +from scancode_required_phrases.training import load_and_validate_splits from scancode_required_phrases.training import main -from scancode_required_phrases.training import PhraseDataset from scancode_required_phrases.training import prepare_output_dir +from scancode_required_phrases.training import promote_final_model +from scancode_required_phrases.training import resolve_model_identity +from scancode_required_phrases.training import resolve_tokenizer_revision from scancode_required_phrases.training import serializable_config from scancode_required_phrases.training import sha256 from scancode_required_phrases.training import validate_bioes from scancode_required_phrases.training import validate_config +from scancode_required_phrases.training import validate_publishable_model from scancode_required_phrases.training import validate_record from scancode_required_phrases.training import validate_saved_state -from scancode_required_phrases.training import validate_splits +from scancode_required_phrases.training import validate_raw_split_hashes +from scancode_required_phrases.training import validate_state_dicts +from scancode_required_phrases.training import validate_state_structure +from scancode_required_phrases.training import write_failure_manifest +from scancode_required_phrases.training import write_json_atomic + +REVISION = "a" * 40 class FakeEncoding(dict): def __init__(self, word_ids): super().__init__() self._word_ids = word_ids - self["input_ids"] = [0] * len(word_ids) + self["input_ids"] = [ + (0 if index == 0 else 1) if word_id is None else 10 + word_id + for index, word_id in enumerate(word_ids) + ] self["attention_mask"] = [1] * len(word_ids) def word_ids(self): @@ -39,9 +62,21 @@ def word_ids(self): class FakeTokenizer: - def __call__(self, tokens, max_length=512, **kwargs): - word_ids = [None, *range(len(tokens)), None] - return FakeEncoding(word_ids[:max_length]) + is_fast = True + all_special_ids = [0, 1] + vocab_size = 100 + + def __init__(self, full_ids=None, retained_ids=None): + self.full_ids = full_ids + self.retained_ids = retained_ids + + def __call__(self, tokens, truncation=False, max_length=512, **kwargs): + default = [None, *range(len(tokens)), None] + ids = self.full_ids if not truncation else self.retained_ids + ids = list(ids if ids is not None else default) + if truncation and self.retained_ids is None: + ids = ids[:max_length] + return FakeEncoding(ids) def make_record(identifier="mit_1.RULE", tokens=None, labels=None): @@ -50,8 +85,8 @@ def make_record(identifier="mit_1.RULE", tokens=None, labels=None): "license_expression": "mit", "rule_type": "is_license_notice", "text": "MIT License terms apply", - "tokens": tokens or ["MIT", "License", "terms", "apply"], - "bioes_labels": labels or ["B-REQ", "E-REQ", "O", "O"], + "tokens": tokens if tokens is not None else ["MIT", "License", "terms", "apply"], + "bioes_labels": labels if labels is not None else ["B-REQ", "E-REQ", "O", "O"], } @@ -62,6 +97,39 @@ def write_jsonl(path, records): ) +def make_split_paths(tmp_path, records=None): + records = records or { + "train": [make_record("train.RULE")], + "validation": [make_record("val.RULE")], + "test": [make_record("test.RULE")], + } + paths = {} + for split, filename in (("train", "train.jsonl"), ("validation", "val.jsonl"), ("test", "test.jsonl")): + path = tmp_path / filename + write_jsonl(path, records[split]) + paths[split] = path + return paths + + +def make_config(tmp_path, **changes): + data_dir = tmp_path / "data" + data_dir.mkdir(parents=True) + make_split_paths(data_dir) + values = { + "data_dir": data_dir, + "output_dir": tmp_path / "model", + "model_revision": REVISION, + } + values.update(changes) + return Config(**values) + + +def test_contract_constants_are_exact(): + assert LABELS == ("O", "B-REQ", "I-REQ", "E-REQ", "S-REQ") + assert ARTIFACT_SCHEMA == "scancode-required-phrases-model-v2" + assert CONSTRAINT_CONTRACT == "bioes-hard-v1" + + @pytest.mark.parametrize( "labels", [ @@ -77,54 +145,98 @@ def test_validate_bioes_accepts_valid_sequences(labels): @pytest.mark.parametrize( "labels", - [ - [], - ["I-REQ"], - ["E-REQ"], - ["B-REQ", "O"], - ["O", "I-REQ"], - ["B-REQ", "I-REQ"], - ], + [[], ["I-REQ"], ["E-REQ"], ["B-REQ", "O"], ["O", "I-REQ"], ["B-REQ", "I-REQ"]], ) def test_validate_bioes_rejects_invalid_sequences(labels): assert validate_bioes(labels) -def test_validate_record_accepts_a_complete_record(): +def test_validate_record_requires_exact_non_empty_types_and_positive_labels(): record = make_record() assert validate_record(record, "train.jsonl", 1) is record + for field in ("identifier", "license_expression", "rule_type", "text"): + invalid = make_record() + invalid[field] = 1 + with pytest.raises(TypeError, match=field): + validate_record(invalid, "train.jsonl", 2) -@pytest.mark.parametrize( - "field_name", - ["identifier", "license_expression", "rule_type", "text", "tokens", "bioes_labels"], -) -def test_validate_record_rejects_a_missing_field(field_name): + invalid = make_record(tokens=["MIT", 1]) + invalid["bioes_labels"] = ["B-REQ", "E-REQ"] + with pytest.raises(ValueError, match="token 1"): + validate_record(invalid, "train.jsonl", 3) + + with pytest.raises(ValueError, match="no required phrase"): + validate_record(make_record(labels=["O", "O", "O", "O"]), "train.jsonl", 4) + + +@pytest.mark.parametrize("field", training.RECORD_FIELDS) +def test_validate_record_rejects_missing_fields(field): record = make_record() - del record[field_name] - with pytest.raises(ValueError, match=field_name): - validate_record(record, "train.jsonl", 4) + del record[field] + with pytest.raises(ValueError, match=field): + validate_record(record, "train.jsonl", 9) -def test_validate_record_rejects_mismatched_tokens_and_labels(): - record = make_record(labels=["S-REQ"]) - with pytest.raises(ValueError, match="tokens and"): - validate_record(record, "train.jsonl", 2) +def test_all_raw_splits_are_validated_and_duplicate_ids_report_both_locations(tmp_path): + records = { + "train": [make_record("same.RULE")], + "validation": [make_record("same.RULE")], + "test": [make_record("test.RULE")], + } + paths = make_split_paths(tmp_path, records) + with pytest.raises(ValueError, match=r"same\.RULE.*train split.*validation split"): + load_and_validate_splits(paths) -def test_validate_record_rejects_unknown_labels(): - record = make_record(labels=["B-REQ", "BAD", "O", "O"]) - with pytest.raises(ValueError, match="unknown labels"): - validate_record(record, "train.jsonl", 2) +def test_malformed_late_raw_line_is_not_hidden_by_limit(tmp_path): + paths = make_split_paths(tmp_path) + with paths["test"].open("a", encoding="utf-8") as stream: + stream.write("{bad json}\n") + with pytest.raises(ValueError, match=r"test\.jsonl line 2: malformed JSON"): + load_and_validate_splits(paths) -def test_align_labels_keeps_only_first_subwords(): - tokenizer = lambda tokens, **kwargs: FakeEncoding([None, 0, 1, 1, None]) - encoding, truncated, cut_phrase = align_labels( - ["MIT", "License"], - ["B-REQ", "E-REQ"], - tokenizer, - 512, +def test_raw_hashes_and_content_duplicate_reports_are_complete(tmp_path): + duplicate = make_record("duplicate.RULE") + records = { + "train": [make_record("train.RULE")], + "validation": [duplicate], + "test": [make_record("test.RULE")], + } + paths = make_split_paths(tmp_path, records) + loaded, report = load_and_validate_splits(paths) + + assert [record.identifier for record in loaded["validation"]] == ["duplicate.RULE"] + assert report["h0"]["train"]["serializer"] == "raw-bytes-v1" + assert report["h1"]["train"]["serializer"].startswith("validated-records") + assert report["h0"]["train"]["sha256"] == sha256(paths["train"]) + assert len(report["duplicates"]["tokens"]) == 1 + assert len(report["duplicates"]["tokens"][0]["records"]) == 3 + assert len(report["duplicates"]["labels"]) == 1 + + +def test_h0_changes_with_raw_bytes_and_h1_tracks_validated_content(tmp_path): + paths = make_split_paths(tmp_path) + _records, first = load_and_validate_splits(paths) + validate_raw_split_hashes(paths, first["h0"]) + + paths["train"].write_text(paths["train"].read_text() + "\n", encoding="utf-8") + with pytest.raises(ValueError, match="train split changed during the run"): + validate_raw_split_hashes(paths, first["h0"]) + + _records, second = load_and_validate_splits(paths) + assert first["h0"]["train"]["sha256"] != second["h0"]["train"]["sha256"] + assert first["h1"]["train"]["sha256"] == second["h1"]["train"]["sha256"] + + +def test_align_labels_uses_first_subwords_and_full_coverage(): + tokenizer = FakeTokenizer( + full_ids=[None, 0, 1, 1, None], + retained_ids=[None, 0, 1, 1, None], + ) + encoding, truncated, cut = align_labels( + ["MIT", "License"], ["B-REQ", "E-REQ"], tokenizer, 8 ) assert encoding["labels"] == [ IGNORE_INDEX, @@ -133,75 +245,97 @@ def test_align_labels_keeps_only_first_subwords(): IGNORE_INDEX, IGNORE_INDEX, ] - assert not truncated - assert not cut_phrase + assert not truncated and not cut -def test_align_labels_detects_a_phrase_cut_by_truncation(): - encoding, truncated, cut_phrase = align_labels( - ["prefix", "GNU", "General", "Public", "License"], - ["O", "B-REQ", "I-REQ", "I-REQ", "E-REQ"], - FakeTokenizer(), - 4, +@pytest.mark.parametrize( + ( + "full_ids", "retained_ids", "reason" + ), + [ + ([None, 0, 1, 1, None], [None, 0, 1, None], "omitted-non-o"), + ([None, 0, 2, None], [None, 0, 2, None], "noncontiguous-coverage"), + ([None, 0, None], [None, 0, None], "zero-coverage"), + ], +) +def test_align_labels_rejects_partial_gap_and_zero_coverage(full_ids, retained_ids, reason): + tokens = ["a", "b"] if 2 not in full_ids else ["a", "b", "c"] + labels = ["B-REQ", "E-REQ"] if len(tokens) == 2 else ["B-REQ", "I-REQ", "E-REQ"] + with pytest.raises(AlignmentError) as caught: + align_labels(tokens, labels, FakeTokenizer(full_ids, retained_ids), 8) + assert caught.value.reason == reason + + +def test_align_labels_accepts_only_omitted_o_words(): + tokenizer = FakeTokenizer( + full_ids=[None, 0, 1, 2, 3, None], + retained_ids=[None, 0, 1, None], ) - assert encoding["labels"][-1] == LABEL2ID["I-REQ"] - assert truncated - assert cut_phrase - - -def test_align_labels_allows_safe_truncation(): - _, truncated, cut_phrase = align_labels( + encoding, truncated, _cut = align_labels( ["MIT", "License", "terms", "apply"], ["B-REQ", "E-REQ", "O", "O"], - FakeTokenizer(), + tokenizer, 4, ) assert truncated - assert not cut_phrase - - -def test_phrase_dataset_skips_a_cut_phrase(tmp_path): - path = tmp_path / "train.jsonl" - write_jsonl( - path, - [ - make_record(identifier="safe.RULE"), - make_record( - identifier="cut.RULE", - tokens=["prefix", "GNU", "General", "Public", "License"], - labels=["O", "B-REQ", "I-REQ", "I-REQ", "E-REQ"], - ), - ], + assert encoding["labels"] == [IGNORE_INDEX, 1, 3, IGNORE_INDEX] + + with pytest.raises(AlignmentError) as caught: + align_labels( + ["prefix", "GNU", "License"], + ["O", "B-REQ", "E-REQ"], + FakeTokenizer([None, 0, 1, 2, None], [None, 0, None]), + 3, + ) + assert caught.value.reason == "omitted-non-o" + assert "positions [1, 2]" in str(caught.value) + + +def test_align_labels_retokenizes_a_partial_o_boundary_to_a_complete_prefix(): + class PartialBoundaryTokenizer(FakeTokenizer): + def __call__(self, tokens, truncation=False, max_length=512, **kwargs): + if len(tokens) == 2: + return FakeEncoding([None, 0, 1, None]) + if truncation: + return FakeEncoding([None, 0, 1, 2, None]) + return FakeEncoding([None, 0, 1, 2, 2, 3, None]) + + encoding, truncated, _cut = align_labels( + ["MIT", "License", "longword", "tail"], + ["B-REQ", "E-REQ", "O", "O"], + PartialBoundaryTokenizer(), + 5, ) - dataset = PhraseDataset(path, FakeTokenizer(), max_length=4) - - assert dataset.identifiers == ["safe.RULE"] - assert dataset.truncated == 2 - assert dataset.cut_phrases == 1 - - -def test_validate_splits_with_real_dataset_objects(tmp_path): - paths = {} - for name, identifier in (("train", "a.RULE"), ("validation", "b.RULE")): - path = tmp_path / f"{name}.jsonl" - write_jsonl(path, [make_record(identifier=identifier)]) - paths[name] = PhraseDataset(path, FakeTokenizer(), 512) - validate_splits(paths) + assert truncated + assert encoding.word_ids() == [None, 0, 1, None] + assert encoding["labels"] == [ + IGNORE_INDEX, + LABEL2ID["B-REQ"], + LABEL2ID["E-REQ"], + IGNORE_INDEX, + ] -def test_validate_splits_rejects_duplicate_identifiers(tmp_path): - datasets = {} - for name in ("train", "validation"): - path = tmp_path / f"{name}.jsonl" - write_jsonl(path, [make_record(identifier="same.RULE")]) - datasets[name] = PhraseDataset(path, FakeTokenizer(), 512) - with pytest.raises(ValueError, match="both train and validation"): - validate_splits(datasets) +def test_effective_limit_is_applied_after_all_alignment_accounting(tmp_path): + records = { + "train": [make_record("first.RULE"), make_record("second.RULE")], + "validation": [make_record("val.RULE")], + "test": [make_record("test.RULE")], + } + paths = make_split_paths(tmp_path, records) + loaded, _report = load_and_validate_splits(paths) + tokenizer = FakeTokenizer() + datasets, h2 = build_effective_datasets(loaded, tokenizer, 512, limit=1) + assert datasets["train"].identifiers == ["first.RULE"] + assert len(datasets["train"].effective_inventory) == 2 + assert h2["train"]["effective_count"] == 2 + assert h2["train"]["selected_count"] == 1 + assert h2["train"]["serializer"].startswith("effective-examples") @pytest.mark.parametrize( - "tags, expected", + ("tags", "expected"), [ (["O", "B-REQ", "I-REQ", "E-REQ", "O"], {(1, 3)}), (["O", "S-REQ", "O"], {(1, 1)}), @@ -209,124 +343,882 @@ def test_validate_splits_rejects_duplicate_identifiers(tmp_path): (["O", "O"], set()), ], ) -def test_extract_spans(tags, expected): +def test_extract_spans_requires_valid_sequences(tags, expected): assert extract_spans(tags) == expected + with pytest.raises(ValueError, match="invalid BIOES"): + extract_spans(["I-REQ"]) -def test_decode_row_drops_ignored_positions(): +def test_decode_row_is_strict_and_drops_only_ignored_gold_positions(): predicted, actual = decode_row( - [LABEL2ID["B-REQ"], 0, LABEL2ID["E-REQ"]], + [LABEL2ID["B-REQ"], 99, LABEL2ID["E-REQ"]], [LABEL2ID["B-REQ"], IGNORE_INDEX, LABEL2ID["E-REQ"]], ) assert predicted == ["B-REQ", "E-REQ"] assert actual == ["B-REQ", "E-REQ"] + with pytest.raises(ValueError, match="lengths differ"): + decode_row([0], [0, 0]) + with pytest.raises(ValueError, match="Unknown prediction ID"): + decode_row([99], [0]) -def test_compute_metrics_scores_exact_spans(): +def test_compute_metrics_scores_exact_spans_and_reports_invalid_count(): predictions = [[LABEL2ID["B-REQ"], LABEL2ID["E-REQ"], LABEL2ID["O"]]] labels = [[LABEL2ID["B-REQ"], LABEL2ID["E-REQ"], LABEL2ID["O"]]] - scores = compute_metrics((predictions, labels)) - assert scores["f1"] == 1.0 - assert scores["exact_match"] == 1.0 - assert scores["predicted_spans"] == scores["gold_spans"] == 1 + scores = compute_metrics((predictions, labels), use_crf=True) + assert scores == { + "f1": 1.0, + "precision": 1.0, + "recall": 1.0, + "exact_match": 1.0, + "predicted_spans": 1, + "gold_spans": 1, + "invalid_paths": 0, + } + + +def test_non_crf_invalid_path_is_not_repaired_and_counts_all_gold_false_negative(): + predictions = [[LABEL2ID["I-REQ"], LABEL2ID["E-REQ"]]] + labels = [[LABEL2ID["B-REQ"], LABEL2ID["E-REQ"]]] + scores = compute_metrics((predictions, labels), use_crf=False) + assert scores["invalid_paths"] == 1 + assert scores["predicted_spans"] == 0 + assert scores["gold_spans"] == 1 + assert scores["recall"] == 0 + assert scores["exact_match"] == 0 + + with pytest.raises(ValueError, match="Invalid CRF BIOES path in row 0"): + compute_metrics((predictions, labels), use_crf=True) + + +def test_metrics_reject_shapes_ids_invalid_gold_and_crf_padding(): + with pytest.raises(ValueError, match="rank 2"): + compute_metrics(([0], [0])) + with pytest.raises(ValueError, match="batch sizes"): + compute_metrics(([[0]], [[0], [0]])) + with pytest.raises(ValueError, match="row lengths"): + compute_metrics(([[0, 0]], [[0]])) + with pytest.raises(ValueError, match="Unknown prediction ID"): + compute_metrics(([[9]], [[0]])) + with pytest.raises(ValueError, match="Unknown gold label ID"): + compute_metrics(([[0]], [[9]])) + with pytest.raises(ValueError, match="Invalid gold BIOES path"): + compute_metrics(([[0]], [[LABEL2ID["I-REQ"]]])) + with pytest.raises(ValueError, match="padding must be left aligned"): + compute_metrics( + ( + [[LABEL2ID["B-REQ"], IGNORE_INDEX, LABEL2ID["E-REQ"]]], + [[LABEL2ID["B-REQ"], IGNORE_INDEX, LABEL2ID["E-REQ"]]], + ), + use_crf=True, + ) + with pytest.raises(ValueError, match="prediction padding"): + compute_metrics( + ([[LABEL2ID["O"], LABEL2ID["O"]]], [[LABEL2ID["O"], IGNORE_INDEX]]), + use_crf=True, + ) def test_first_subword_positions_skips_specials_and_continuations(): assert first_subword_positions([None, 0, 1, 1, 2, None]) == [1, 2, 4] +def test_config_validation_is_exhaustive_and_requires_immutable_revision(tmp_path): + validate_config(make_config(tmp_path)) + + invalid = make_config(tmp_path / "bad-revision", model_revision="main") + with pytest.raises(ValueError, match="full immutable"): + validate_config(invalid) + + +@pytest.mark.parametrize( + ("field", "value", "error"), + [ + ("epochs", 0, "between"), + ("batch_size", True, "integer"), + ("base_lr", float("nan"), "finite"), + ("head_lr", 0, "greater than"), + ("layer_decay", 1.1, "at most"), + ("warmup_ratio", -1, "at least"), + ("max_grad_norm", 0, "greater than"), + ("optimizer", "other", "Unsupported optimizer"), + ("precision", "fp16", "Unsupported precision"), + ("limit", -1, "between"), + ("label_weights", [1.0], "exactly 5"), + ("resume", True, "unsupported"), + ], +) +def test_config_rejects_invalid_fields(tmp_path, field, value, error): + config = make_config(tmp_path) + setattr(config, field, value) + with pytest.raises((TypeError, ValueError), match=error): + validate_config(config) + + +def test_config_rejects_non_empty_output(tmp_path): + config = make_config(tmp_path) + config.output_dir.mkdir() + (config.output_dir / "anything").write_text("x") + with pytest.raises(ValueError, match="not empty"): + validate_config(config) + + +def test_prepare_output_dir_is_empty_only_and_never_resumes(tmp_path): + output = tmp_path / "model" + prepare_output_dir(output) + assert output.is_dir() + with pytest.raises(ValueError, match="unsupported"): + prepare_output_dir(output, resume=True) + (output / "checkpoint-1").mkdir() + with pytest.raises(ValueError, match="not empty"): + prepare_output_dir(output) + + +def test_atomic_json_is_complete_and_leaves_no_temporary_file(tmp_path): + path = tmp_path / "manifest.json" + write_json_atomic(path, {"state": "pre-run", "value": 1}) + write_json_atomic(path, {"state": "success", "value": 2}) + assert json.loads(path.read_text()) == {"state": "success", "value": 2} + assert list(tmp_path.iterdir()) == [path] + + +def test_serializable_config_contains_every_field_and_string_paths(tmp_path): + config = make_config(tmp_path) + values = serializable_config(config) + assert set(values) == set(config.__dataclass_fields__) + assert values["data_dir"] == str(config.data_dir) + assert values["output_dir"] == str(config.output_dir) + + def test_sha256_is_stable(tmp_path): path = tmp_path / "data.jsonl" path.write_bytes(b"required phrase\n") assert sha256(path) == "792616e2062f96efb6ae2f69e8637b834e2db74354eb4f51e78eda329038cc70" -def test_serializable_config_converts_paths(tmp_path): - config = Config(data_dir=tmp_path / "data", output_dir=tmp_path / "model") - values = serializable_config(config) - assert values["data_dir"] == str(tmp_path / "data") - assert values["output_dir"] == str(tmp_path / "model") +def test_exact_state_comparison_and_saved_state(tmp_path): + torch = pytest.importorskip("torch") + safetensors = pytest.importorskip("safetensors.torch") + model = torch.nn.Linear(2, 1) + state = model.state_dict() + reversed_state = dict(reversed(list(state.items()))) + with pytest.raises(ValueError, match="key mismatch"): + validate_state_dicts(state, reversed_state) + + changed = dict(state) + changed["weight"] = changed["weight"].clone() + changed["weight"][0, 0] += 1 + validate_state_structure(state, changed) + with pytest.raises(ValueError, match="values differ"): + validate_state_dicts(state, changed) + changed = dict(state) + changed["unexpected"] = torch.ones(1) + with pytest.raises(ValueError, match="key mismatch"): + validate_state_dicts(state, changed) + model_path = tmp_path / "model.safetensors" + safetensors.save_file(state, str(model_path)) + validate_saved_state(model, model_path) -def test_validate_config_accepts_the_default_training_settings(tmp_path): - config = Config(data_dir=tmp_path / "data", output_dir=tmp_path / "model") - validate_config(config) +def test_ordered_validation_comparison_is_exact(): + result = { + "prediction_ids": [[0, 4]], + "label_ids": [[0, 4]], + "invalid_paths": 0, + "metrics": {"f1": 1.0}, + } + compare_ordered_validation_results(result, dict(result)) + for field, value in ( + ("prediction_ids", [[4, 0]]), + ("label_ids", [[4, 0]]), + ("invalid_paths", 1), + ("metrics", {"f1": 0.0}), + ): + changed = dict(result) + changed[field] = value + with pytest.raises(ValueError, match=field): + compare_ordered_validation_results(result, changed) + + +def write_complete_artifact(directory, tmp_path, selected="checkpoint-1"): + directory.mkdir(parents=True) + files = { + "config.json": "{}\n", + "model.safetensors": "weights", + "special_tokens_map.json": "{}\n", + "tokenizer_config.json": "{}\n", + "tokenizer.json": "{}\n", + } + for name, content in files.items(): + (directory / name).write_text(content, encoding="utf-8") + config = make_config(tmp_path / "artifact-config") + artifact_config = training._artifact_config(config, REVISION) + write_json_atomic(directory / "train_config.json", artifact_config) + manifest = { + "schema": training.MANIFEST_SCHEMA, + "state": "success", + "config": serializable_config(config), + "contracts": { + "artifact_schema": ARTIFACT_SCHEMA, + "constraint_contract": CONSTRAINT_CONTRACT, + "labels": list(LABELS), + }, + "dataset": {"paths": {}, "h0": {}, "h1": {}, "h2": {}, "report": {}}, + "model_identity": { + "name": config.model_name, + "requested_revision": REVISION, + "resolved_tokenizer_revision": REVISION, + "resolved_backbone_revision": REVISION, + }, + "source": { + "repository": "example/repository", + "root": "repository", + "branch": "test", + "commit": REVISION, + "dirty": False, + }, + "runtime": { + "python": "3", + "platform": "test", + "torch": "test", + "transformers": "test", + "optimizer": {"configured": "adamw"}, + "precision": "fp32", + }, + "completed_checks": ["test"], + "selected_checkpoint": selected, + "best_validation_f1": 1.0, + "validation_metrics": {"f1": 1.0}, + "ordered_validation": {"metrics": {"f1": 1.0}}, + "test_metrics": None, + "artifact_files": training._all_file_hashes(directory), + "log_history": [], + } + write_json_atomic(directory / "run_manifest.json", manifest) + return { + "schema": ARTIFACT_SCHEMA, + "constraint_contract": CONSTRAINT_CONTRACT, + "selected_checkpoint": selected, + "run_manifest_sha256": sha256(directory / "run_manifest.json"), + "files": training._all_file_hashes(directory), + } -def test_validate_config_rejects_invalid_settings(tmp_path): - config = Config( - data_dir=tmp_path / "data", - output_dir=tmp_path / "model", - aux_ce_weight=-1, - ) - with pytest.raises(ValueError, match="cannot be negative"): - validate_config(config) +def make_publishable_model(tmp_path): + model_dir = tmp_path / "final-model" + marker = write_complete_artifact(model_dir, tmp_path) + write_json_atomic(model_dir / "SUCCESS.json", marker) + return model_dir -def test_prepare_output_dir_refuses_to_overwrite_a_run(tmp_path): - output_dir = tmp_path / "model" - output_dir.mkdir() - (output_dir / "model.safetensors").write_bytes(b"model") - with pytest.raises(ValueError, match="not empty"): - prepare_output_dir(output_dir, resume=False) +def test_publishability_requires_marker_exact_inventory_and_hashes(tmp_path): + model_dir = make_publishable_model(tmp_path) + assert validate_publishable_model(model_dir)["schema"] == ARTIFACT_SCHEMA + (model_dir / "model.safetensors").write_bytes(b"tampered") + with pytest.raises(ValueError, match="hash mismatch"): + validate_publishable_model(model_dir) -def test_prepare_output_dir_requires_a_checkpoint_to_resume(tmp_path): - with pytest.raises(ValueError, match="No checkpoint"): - prepare_output_dir(tmp_path / "model", resume=True) +def test_publishability_rejects_missing_marker_and_unrecorded_file(tmp_path): + empty = tmp_path / "empty" + empty.mkdir() + with pytest.raises(ValueError, match="Success_Marker"): + validate_publishable_model(empty) -def test_validate_saved_state_accepts_matching_tensors(tmp_path): - torch = pytest.importorskip("torch") - safetensors = pytest.importorskip("safetensors.torch") + model_dir = make_publishable_model(tmp_path / "nested") + (model_dir / "extra").write_text("not recorded") + with pytest.raises(ValueError, match="inventory"): + validate_publishable_model(model_dir) - model = torch.nn.Linear(2, 1) - path = tmp_path / "model.safetensors" - safetensors.save_file(model.state_dict(), str(path)) - validate_saved_state(model, path) +def test_cli_requires_revision_and_describes_isr_as_locatability(tmp_path): + runner = CliRunner() + result = runner.invoke(main, ["--data-dir", str(tmp_path), "--with-isr"]) + assert result.exit_code == 2 + assert "Missing option '--model-revision'" in result.output + help_result = runner.invoke(main, ["--help"]) + assert "predicted-phrase locatability" in help_result.output + assert "Resume" not in help_result.output -def test_validate_saved_state_rejects_unexpected_keys(tmp_path): - torch = pytest.importorskip("torch") - safetensors = pytest.importorskip("safetensors.torch") - model = torch.nn.Linear(2, 1) - state = dict(model.state_dict()) - state["unexpected"] = torch.ones(1) - path = tmp_path / "model.safetensors" - safetensors.save_file(state, str(path)) - with pytest.raises(ValueError, match="key mismatch"): - validate_saved_state(model, path) +def test_cli_rejects_resume_and_isr_without_test(tmp_path): + arguments = [ + "--data-dir", str(tmp_path), "--model-revision", REVISION, + ] + result = CliRunner().invoke(main, arguments + ["--with-isr"]) + assert result.exit_code == 2 + assert "--with-isr requires --evaluate-test" in result.output + result = CliRunner().invoke(main, arguments + ["--resume"]) + assert result.exit_code == 2 + assert "--resume is unsupported" in result.output -def test_validate_saved_state_rejects_wrong_shapes(tmp_path): - torch = pytest.importorskip("torch") - safetensors = pytest.importorskip("safetensors.torch") +def test_tokenizer_revision_uses_the_requested_hub_commit(monkeypatch): + calls = [] - model = torch.nn.Linear(2, 1) - state = dict(model.state_dict()) - state["weight"] = torch.ones((1, 3)) - path = tmp_path / "model.safetensors" - safetensors.save_file(state, str(path)) - with pytest.raises(ValueError, match="has shape"): - validate_saved_state(model, path) + def get_model_info(model_name, revision): + calls.append((model_name, revision)) + return types.SimpleNamespace(sha=REVISION) + + hub = types.ModuleType("huggingface_hub") + hub.model_info = get_model_info + monkeypatch.setitem(__import__("sys").modules, "huggingface_hub", hub) + + assert resolve_tokenizer_revision("model", REVISION) == REVISION + assert calls == [("model", REVISION)] + + hub.model_info = lambda *args, **kwargs: types.SimpleNamespace(sha="main") + with pytest.raises(ValueError, match="full immutable commit"): + resolve_tokenizer_revision("model", REVISION) -def test_validate_saved_state_rejects_non_finite_values(tmp_path): +def test_model_identity_requires_exact_tokenizer_and_backbone_commits(): + backbone = type("BackboneConfig", (), {"_commit_hash": REVISION})() + assert resolve_model_identity(REVISION, REVISION, backbone) == REVISION + + backbone._commit_hash = "b" * 40 + with pytest.raises(ValueError, match="inconsistent"): + resolve_model_identity(REVISION, REVISION, backbone) + with pytest.raises(ValueError, match="inconsistent"): + resolve_model_identity(REVISION, "b" * 40, backbone) + + +def test_source_provenance_uses_narrow_git_values_and_no_environment(monkeypatch, tmp_path): + values = { + ("rev-parse", "--show-toplevel"): str(tmp_path), + ("branch", "--show-current"): "gsoc/training-pipeline", + ("rev-parse", "HEAD"): REVISION, + ("config", "--get", "remote.origin.url"): "https://secret@example.com/project.git", + ("status", "--porcelain"): " M allowed.py", + } + monkeypatch.setattr(training, "_run_git", lambda _root, *args: values[args]) + monkeypatch.setenv("SECRET_TOKEN", "must-not-appear") + provenance = training.collect_source_provenance(tmp_path) + serialized = json.dumps(provenance) + assert provenance["repository"] == "https://example.com/project.git" + assert provenance["branch"] == "gsoc/training-pipeline" + assert provenance["commit"] == REVISION + assert provenance["dirty"] is True + assert "SECRET_TOKEN" not in serialized + assert "must-not-appear" not in serialized + + +def test_promotion_writes_marker_last_and_is_publishable(tmp_path): + stage = tmp_path / "final-model.tmp" + marker = write_complete_artifact(stage, tmp_path) + destination = tmp_path / "final-model" + promote_final_model(stage, destination, marker) + assert not stage.exists() + assert validate_publishable_model(destination) == marker + + +def test_promotion_marker_failure_rolls_back_and_leaves_no_final_model( + monkeypatch, tmp_path +): + stage = tmp_path / "final-model.tmp" + stage.mkdir() + (stage / "model.safetensors").write_bytes(b"weights") + destination = tmp_path / "final-model" + + def fail_marker(path, value): + raise OSError("injected marker failure") + + monkeypatch.setattr(training, "write_json_atomic", fail_marker) + with pytest.raises(OSError, match="injected marker failure"): + promote_final_model(stage, destination, {}) + assert stage.is_dir() + assert not destination.exists() + assert not (stage / "SUCCESS.json").exists() + + +class RejectingLateTokenizer(FakeTokenizer): + def __call__(self, tokens, truncation=False, max_length=512, **kwargs): + if tokens[0] == "bad": + return FakeEncoding([None, 0, None]) + return super().__call__(tokens, truncation, max_length, **kwargs) + + +def test_limit_does_not_hide_late_alignment_rejections(tmp_path): + records = { + "train": [ + make_record("selected.RULE"), + make_record( + "rejected.RULE", + tokens=["bad", "record"], + labels=["B-REQ", "E-REQ"], + ), + ], + "validation": [make_record("val.RULE")], + "test": [make_record("test.RULE")], + } + loaded, _report = load_and_validate_splits(make_split_paths(tmp_path, records)) + datasets, h2 = build_effective_datasets( + loaded, RejectingLateTokenizer(), max_length=512, limit=1 + ) + assert datasets["train"].identifiers == ["selected.RULE"] + assert datasets["train"].rejections[0]["identifier"] == "rejected.RULE" + assert datasets["train"].rejections[0]["reason"] == "zero-coverage" + assert h2["train"]["rejected_count"] == 1 + + +def test_split_with_no_effective_example_is_rejected(tmp_path): + bad = make_record( + "bad.RULE", tokens=["bad", "record"], labels=["B-REQ", "E-REQ"] + ) + records = { + "train": [bad], + "validation": [make_record("val.RULE")], + "test": [make_record("test.RULE")], + } + loaded, _report = load_and_validate_splits(make_split_paths(tmp_path, records)) + with pytest.raises(ValueError, match="train split has no selected Effective_Example"): + build_effective_datasets(loaded, RejectingLateTokenizer(), 512) + + +def test_failure_manifest_is_atomic_scoped_and_reports_retained_paths(tmp_path): + manifest_path = tmp_path / "run_manifest.json" + retained = tmp_path / "dataset_report.json" + retained.write_text("{}\n") + value = write_failure_manifest( + manifest_path, + {"schema": "run-v2", "state": "pre-run"}, + "offline-reload", + ValueError("state mismatch"), + ["configuration", "training"], + [retained, tmp_path / "absent"], + ) + assert json.loads(manifest_path.read_text()) == value + assert value["state"] == "failure" + assert value["failed_phase"] == "offline-reload" + assert value["exception"] == {"type": "ValueError", "message": "state mismatch"} + assert value["completed_checks"] == ["configuration", "training"] + assert value["retained_artifacts"] == [str(retained)] + assert "environment" not in value + + +def test_duplicate_check_waits_until_every_raw_record_is_validated(tmp_path): + records = { + "train": [make_record("duplicate.RULE")], + "validation": [make_record("duplicate.RULE")], + "test": [make_record("test.RULE")], + } + paths = make_split_paths(tmp_path, records) + with paths["test"].open("a", encoding="utf-8") as stream: + stream.write("{malformed}\n") + with pytest.raises(ValueError, match=r"test\.jsonl line 2: malformed JSON"): + load_and_validate_splits(paths) + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("epochs", 1_001), + ("batch_size", 1_025), + ("base_lr", 1.1), + ("weight_decay", 1.1), + ("aux_ce_weight", 101), + ("label_weights", [1.0, 2.0, 1.5, 1.5, 1_000_001]), + ], +) +def test_config_numeric_contracts_are_upper_bounded(tmp_path, field, value): + config = make_config(tmp_path) + setattr(config, field, value) + with pytest.raises(ValueError): + validate_config(config) + + +@pytest.mark.parametrize( + ("full_ids", "retained_ids", "reason"), + [ + ([0, 1, None], [0, 1, None], "missing-special-token"), + ([None, 0, 1], [None, 0, 1], "missing-special-token"), + ([None, 0, None, 1, None], [None, 0, None, 1, None], "coverage-gap"), + ], +) +def test_alignment_requires_boundary_specials_and_no_internal_specials( + full_ids, retained_ids, reason +): + with pytest.raises(AlignmentError) as caught: + align_labels( + ["MIT", "License"], + ["B-REQ", "E-REQ"], + FakeTokenizer(full_ids, retained_ids), + 8, + ) + assert caught.value.reason == reason + + +def test_alignment_rejects_word_id_and_model_input_shape_mismatch(): + class BadShapeTokenizer(FakeTokenizer): + def __call__(self, tokens, **kwargs): + encoding = FakeEncoding([None, 0, 1, None]) + encoding["attention_mask"].pop() + return encoding + + with pytest.raises(AlignmentError) as caught: + align_labels( + ["MIT", "License"], ["B-REQ", "E-REQ"], BadShapeTokenizer(), 8 + ) + assert caught.value.reason == "shape-mismatch" + + +def test_publishability_requires_complete_reloadable_inventory(tmp_path): + model_dir = make_publishable_model(tmp_path) + (model_dir / "config.json").unlink() + marker = json.loads((model_dir / "SUCCESS.json").read_text()) + marker["files"].pop("config.json") + write_json_atomic(model_dir / "SUCCESS.json", marker) + with pytest.raises(ValueError, match="missing required files"): + validate_publishable_model(model_dir) + + +def test_publishability_rejects_unknown_artifact_configuration_fields(tmp_path): + model_dir = make_publishable_model(tmp_path) + config_path = model_dir / "train_config.json" + config = json.loads(config_path.read_text()) + config["unknown"] = True + write_json_atomic(config_path, config) + marker = json.loads((model_dir / "SUCCESS.json").read_text()) + marker["files"]["train_config.json"] = sha256(config_path) + write_json_atomic(model_dir / "SUCCESS.json", marker) + with pytest.raises(ValueError, match="artifact hashes|configuration fields differ"): + validate_publishable_model(model_dir) + + +def test_local_loader_loads_saved_values_after_structural_validation( + monkeypatch, tmp_path +): torch = pytest.importorskip("torch") safetensors = pytest.importorskip("safetensors.torch") + transformers = pytest.importorskip("transformers") + from scancode_required_phrases import model as model_module + + model_dir = tmp_path / "model" + model_dir.mkdir() + config = make_config(tmp_path) + write_json_atomic( + model_dir / "train_config.json", + training._artifact_config(config, REVISION), + ) + safetensors.save_file( + {"weight": torch.tensor([3.0])}, + str(model_dir / "model.safetensors"), + ) - model = torch.nn.Linear(2, 1) - state = dict(model.state_dict()) - state["weight"] = torch.full_like(state["weight"], float("nan")) - path = tmp_path / "model.safetensors" - safetensors.save_file(state, str(path)) - with pytest.raises(ValueError, match="non-finite"): - validate_saved_state(model, path) + class LocalTagger(torch.nn.Module): + def __init__(self, config, backbone=None): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor([0.0])) + class LocalTokenizer: + is_fast = True -def test_cli_requires_test_evaluation_for_isr(tmp_path): - result = CliRunner().invoke(main, ["--data-dir", str(tmp_path), "--with-isr"]) - assert result.exit_code == 2 - assert "--with-isr requires --evaluate-test" in result.output + monkeypatch.setattr(model_module, "PhraseTagger", LocalTagger) + monkeypatch.setattr( + transformers.AutoConfig, + "from_pretrained", + lambda *args, **kwargs: object(), + ) + monkeypatch.setattr( + transformers.AutoModel, + "from_config", + lambda config: object(), + ) + monkeypatch.setattr( + transformers.AutoTokenizer, + "from_pretrained", + lambda *args, **kwargs: LocalTokenizer(), + ) + + loaded, tokenizer = training._load_local_model(model_dir) + + assert loaded.weight.item() == 3.0 + assert tokenizer.is_fast + + +def test_final_model_loader_rejects_unpublished_stage_before_local_loading( + monkeypatch, tmp_path +): + stage = tmp_path / "final-model.tmp" + stage.mkdir() + monkeypatch.setattr( + training, + "_load_local_model", + lambda *args, **kwargs: pytest.fail("unpublished artifact must not load"), + ) + with pytest.raises(ValueError, match="Success_Marker"): + training.load_final_model(stage) + + +def test_repository_identity_removes_all_url_credentials(): + assert training._redacted_repository_identity( + "https://user:password@example.com/repo.git?token=secret#credential" + ) == "https://example.com/repo.git" + assert training._redacted_repository_identity( + "git@example.com:aboutcode/repo.git" + ) == "example.com:aboutcode/repo.git" + + +def test_selected_checkpoint_must_be_a_child_directory(tmp_path): + state = type("State", (), {"best_model_checkpoint": str(tmp_path), "best_metric": 1.0})() + trainer = type("Trainer", (), {"state": state})() + with pytest.raises(ValueError, match="inside run output"): + training.validate_selected_checkpoint(trainer, tmp_path) + + outside = tmp_path.parent / "outside-checkpoint" + outside.mkdir(exist_ok=True) + state.best_model_checkpoint = str(outside) + with pytest.raises(ValueError, match="inside run output"): + training.validate_selected_checkpoint(trainer, tmp_path) + + +def test_run_training_validates_every_raw_line_before_tokenizer_loading( + monkeypatch, tmp_path +): + config = make_config(tmp_path) + with (config.data_dir / "test.jsonl").open("a", encoding="utf-8") as stream: + stream.write("{malformed}\n") + calls = [] + + class AutoTokenizer: + @staticmethod + def from_pretrained(*args, **kwargs): + calls.append((args, kwargs)) + pytest.fail("tokenizer must not load before complete raw validation") + + fake_transformers = types.ModuleType("transformers") + fake_transformers.AutoConfig = object + fake_transformers.AutoTokenizer = AutoTokenizer + fake_transformers.DataCollatorForTokenClassification = object + fake_transformers.EarlyStoppingCallback = object + fake_transformers.TrainingArguments = object + fake_model = types.ModuleType("scancode_required_phrases.model") + fake_model.PhraseTagger = object + fake_model.PhraseTrainer = object + fake_model.build_optimizer = object + modules = __import__("sys").modules + monkeypatch.setitem(modules, "transformers", fake_transformers) + monkeypatch.setitem(modules, "torch", types.ModuleType("torch")) + monkeypatch.setitem(modules, "scancode_required_phrases.model", fake_model) + monkeypatch.setattr(training, "validate_precision", lambda precision: None) + + with pytest.raises(ValueError, match=r"test\.jsonl line 2: malformed JSON"): + training.run_training(config) + assert calls == [] + assert not config.output_dir.exists() + + +def test_run_training_records_offline_reload_failure_without_publication( + monkeypatch, tmp_path +): + config = make_config(tmp_path) + raw_report = { + "h0": {}, + "h1": {}, + "duplicates": {"tokens": [], "labels": []}, + "raw_counts": {"train": 1, "validation": 1, "test": 1}, + } + records = {name: [] for name in ("train", "validation", "test")} + + class Dataset: + examples = [{"input_ids": [1], "attention_mask": [1], "labels": [0]}] + effective_inventory = examples + rejections = [] + truncations = [] + truncated = 0 + cut_phrases = 0 + + def __len__(self): + return 1 + + datasets = {name: Dataset() for name in records} + h2 = {name: {"sha256": name} for name in records} + + class Tokenizer: + is_fast = True + init_kwargs = {"_commit_hash": REVISION} + + class AutoTokenizer: + @staticmethod + def from_pretrained(*args, **kwargs): + return Tokenizer() + + class AutoConfig: + @staticmethod + def from_pretrained(*args, **kwargs): + return types.SimpleNamespace(_commit_hash=REVISION) + + class TrainingArguments: + def __init__(self, output_dir, **kwargs): + self.output_dir = output_dir + + class EarlyStoppingCallback: + def __init__(self, **kwargs): + pass + + class PhraseTagger: + def __init__(self, config): + self.backbone = types.SimpleNamespace( + config=types.SimpleNamespace(_commit_hash=REVISION) + ) + + def state_dict(self): + return {"weight": object()} + + class PhraseTrainer: + def __init__(self, model, args, **kwargs): + self.model = model + self.args = args + self.state = types.SimpleNamespace( + best_model_checkpoint=None, + best_metric=1.0, + log_history=[], + ) + + def train(self): + checkpoint = Path(self.args.output_dir) / "checkpoint-1" + checkpoint.mkdir() + (checkpoint / "model.safetensors").write_bytes(b"checkpoint") + self.state.best_model_checkpoint = str(checkpoint) + + def remove_callback(self, callback): + pass + + def evaluate(self, dataset, metric_key_prefix): + return {f"{metric_key_prefix}_f1": 1.0} + + fake_transformers = types.ModuleType("transformers") + fake_transformers.AutoConfig = AutoConfig + fake_transformers.AutoTokenizer = AutoTokenizer + fake_transformers.DataCollatorForTokenClassification = lambda *args, **kwargs: object() + fake_transformers.EarlyStoppingCallback = EarlyStoppingCallback + fake_transformers.TrainingArguments = TrainingArguments + fake_model = types.ModuleType("scancode_required_phrases.model") + fake_model.PhraseTagger = PhraseTagger + fake_model.PhraseTrainer = PhraseTrainer + fake_model.build_optimizer = lambda config, model: None + modules = __import__("sys").modules + monkeypatch.setitem(modules, "torch", types.ModuleType("torch")) + monkeypatch.setitem(modules, "transformers", fake_transformers) + monkeypatch.setitem(modules, "scancode_required_phrases.model", fake_model) + monkeypatch.setattr(training, "validate_precision", lambda precision: None) + monkeypatch.setattr( + training, + "resolve_tokenizer_revision", + lambda model_name, requested: REVISION, + ) + monkeypatch.setattr(training, "validate_raw_split_hashes", lambda paths, hashes: None) + monkeypatch.setattr(training, "set_seed", lambda seed: None) + monkeypatch.setattr( + training, + "load_and_validate_splits", + lambda paths: (records, raw_report), + ) + monkeypatch.setattr( + training, + "build_effective_datasets", + lambda records, tokenizer, max_length, limit: (datasets, h2), + ) + monkeypatch.setattr( + training, + "collect_source_provenance", + lambda: { + "repository": "example/repository", + "root": "repository", + "branch": "test", + "commit": REVISION, + "dirty": False, + }, + ) + monkeypatch.setattr( + training, + "collect_runtime_provenance", + lambda optimizer, precision: {"optimizer": optimizer, "precision": precision}, + ) + monkeypatch.setattr(training, "_load_state_file", lambda path: {"weight": object()}) + monkeypatch.setattr(training, "validate_state_dicts", lambda *args, **kwargs: None) + monkeypatch.setattr( + training, + "collect_ordered_validation_result", + lambda model, dataset, use_crf: { + "prediction_ids": [[0]], + "label_ids": [[0]], + "invalid_paths": 0, + "metrics": {"f1": 1.0}, + }, + ) + + def stage_final_model(output_dir, model, tokenizer, artifact_config): + stage = Path(output_dir) / "final-model.tmp" + stage.mkdir() + (stage / "model.safetensors").write_bytes(b"staged") + return stage + + monkeypatch.setattr(training, "stage_final_model", stage_final_model) + monkeypatch.setattr( + training, + "_load_local_model", + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("offline failure")), + ) + + with pytest.raises(RuntimeError, match="offline failure"): + training.run_training(config) + + manifest = json.loads((config.output_dir / "run_manifest.json").read_text()) + assert manifest["state"] == "failure" + assert manifest["failed_phase"] == "offline-reload" + assert str(config.output_dir / "checkpoint-1") in manifest["retained_artifacts"] + assert str(config.output_dir / "final-model.tmp") in manifest["retained_artifacts"] + assert not (config.output_dir / "final-model").exists() + assert not (config.output_dir / "final-model.tmp" / "SUCCESS.json").exists() + + +def test_alignment_rejects_invalid_model_ids_and_inactive_masks(): + class InvalidModelInputTokenizer(FakeTokenizer): + def __init__(self, input_ids=None, attention_mask=None): + super().__init__() + self.input_ids = input_ids + self.attention_mask = attention_mask + + def __call__(self, tokens, **kwargs): + encoding = super().__call__(tokens, **kwargs) + if self.input_ids is not None: + encoding["input_ids"] = self.input_ids + if self.attention_mask is not None: + encoding["attention_mask"] = self.attention_mask + return encoding + + with pytest.raises(AlignmentError) as caught: + align_labels( + ["MIT", "License"], + ["B-REQ", "E-REQ"], + InvalidModelInputTokenizer(input_ids=[0, 100, 11, 1]), + 8, + ) + assert caught.value.reason == "invalid-input-id" + + with pytest.raises(AlignmentError) as caught: + align_labels( + ["MIT", "License"], + ["B-REQ", "E-REQ"], + InvalidModelInputTokenizer(attention_mask=[0, 0, 0, 0]), + 8, + ) + assert caught.value.reason == "invalid-attention-mask" + + +def test_publishability_rejects_rehashed_incomplete_success_manifest(tmp_path): + model_dir = make_publishable_model(tmp_path) + manifest_path = model_dir / "run_manifest.json" + manifest = json.loads(manifest_path.read_text()) + del manifest["source"] + write_json_atomic(manifest_path, manifest) + marker_path = model_dir / "SUCCESS.json" + marker = json.loads(marker_path.read_text()) + marker["run_manifest_sha256"] = sha256(manifest_path) + marker["files"]["run_manifest.json"] = sha256(manifest_path) + write_json_atomic(marker_path, marker) + with pytest.raises(ValueError, match="manifest fields"): + validate_publishable_model(model_dir) From 959e3dfce817eeeaafbeaec402884c313bc5569e Mon Sep 17 00:00:00 2001 From: Kaushik Date: Mon, 7 Sep 2026 12:25:34 +0000 Subject: [PATCH 3/5] Harden model training and inference Signed-off-by: Kaushik --- src/scancode_required_phrases/inference.py | 139 +++++++++++++++++++++ tests/test_inference.py | 108 ++++++++++++++++ 2 files changed, 247 insertions(+) create mode 100644 src/scancode_required_phrases/inference.py create mode 100644 tests/test_inference.py diff --git a/src/scancode_required_phrases/inference.py b/src/scancode_required_phrases/inference.py new file mode 100644 index 0000000..6d758fd --- /dev/null +++ b/src/scancode_required_phrases/inference.py @@ -0,0 +1,139 @@ +# Copyright (c) nexB Inc. and others. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Read-only required-phrase prediction from a validated final model.""" + +import json +import unicodedata +from dataclasses import dataclass +from pathlib import Path + +from licensedcode.tokenize import required_phrase_splitter + +from scancode_required_phrases.training import extract_spans +from scancode_required_phrases.training import first_subword_positions +from scancode_required_phrases.training import ID2LABEL +from scancode_required_phrases.training import load_final_model + + +@dataclass(frozen=True) +class PhrasePrediction: + """One predicted required-phrase candidate.""" + + text: str + start_word: int + end_word: int + confidence: float + + +@dataclass(frozen=True) +class PredictionResult: + """Read-only predictions and tokenization details for one rule text.""" + + words: tuple[str, ...] + phrases: tuple[PhrasePrediction, ...] + truncated: bool + + +def words_from_text(text): + """Tokenize rule text exactly as the training dataset does.""" + text = text.replace("\r\n", "\n").replace("\r", "\n") + return required_phrase_splitter(unicodedata.normalize("NFKC", text)) + + +def span_confidence(crf, word_emissions, tags, mask, free, span): + """Return the CRF probability mass agreeing with one decoded span.""" + start, end = span + pinned = word_emissions.clone() + floor = float(word_emissions.min()) - 10000.0 + + for position in range(start, end + 1): + label = int(tags[0, position]) + keep = float(pinned[0, position, label]) + pinned[0, position] = floor + pinned[0, position, label] = keep + + constrained = crf(pinned, tags, mask=mask, reduction="none") + confidence = float((free - constrained).detach().exp()) + return min(max(confidence, 0.0), 1.0) + + +class RequiredPhrasePredictor: + """Run a validated hardened model without changing ScanCode rules.""" + + def __init__(self, model, tokenizer, max_length): + self.model = model + self.tokenizer = tokenizer + self.max_length = max_length + + @classmethod + def from_model_dir(cls, model_dir): + """Validate and load a local final-model directory.""" + model_dir = Path(model_dir) + model, tokenizer = load_final_model(model_dir) + config = json.loads((model_dir / "train_config.json").read_text(encoding="utf-8")) + return cls(model=model, tokenizer=tokenizer, max_length=config["max_length"]) + + def predict(self, text): + """Return candidate spans for text without mutating any rule or file.""" + import torch + + words = words_from_text(text) + if not words: + return PredictionResult(words=(), phrases=(), truncated=False) + + encoding = self.tokenizer( + words, + is_split_into_words=True, + truncation=True, + max_length=self.max_length, + return_tensors="pt", + ) + positions = first_subword_positions(encoding.word_ids()) + if not positions: + return PredictionResult(words=tuple(words), phrases=(), truncated=False) + + device = next(self.model.parameters()).device + input_ids = encoding["input_ids"].to(device) + attention_mask = encoding["attention_mask"].to(device) + + with torch.inference_mode(): + emissions = self.model.emissions(input_ids, attention_mask) + word_emissions = emissions[:, positions].float() + mask = torch.ones( + word_emissions.shape[:2], + dtype=torch.bool, + device=word_emissions.device, + ) + decoded = self.model.crf.decode(word_emissions, mask=mask)[0] + tags = torch.tensor([decoded], device=word_emissions.device) + free = self.model.crf(word_emissions, tags, mask=mask, reduction="none") + + labels = [ID2LABEL[int(label)] for label in decoded] + truncated = len(labels) < len(words) + predictions = [] + for start, end in extract_spans(labels): + if truncated and end == len(labels) - 1: + continue + predictions.append( + PhrasePrediction( + text=" ".join(words[start : end + 1]), + start_word=start, + end_word=end, + confidence=span_confidence( + self.model.crf, + word_emissions, + tags, + mask, + free, + (start, end), + ), + ) + ) + + predictions.sort(key=lambda item: (item.start_word, item.end_word)) + return PredictionResult( + words=tuple(words), + phrases=tuple(predictions), + truncated=truncated, + ) diff --git a/tests/test_inference.py b/tests/test_inference.py new file mode 100644 index 0000000..299abf5 --- /dev/null +++ b/tests/test_inference.py @@ -0,0 +1,108 @@ +# Copyright (c) nexB Inc. and others. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from scancode_required_phrases.inference import RequiredPhrasePredictor +from scancode_required_phrases.inference import words_from_text +from scancode_required_phrases.model import ConstrainedCRF +from scancode_required_phrases.training import LABEL2ID +from scancode_required_phrases.training import LABELS + + +class FakeEncoding(dict): + + def __init__(self, word_ids): + super().__init__( + input_ids=torch.zeros((1, len(word_ids)), dtype=torch.long), + attention_mask=torch.ones((1, len(word_ids)), dtype=torch.long), + ) + self._word_ids = word_ids + + def word_ids(self): + return self._word_ids + + +class FakeTokenizer: + + def __call__(self, words, max_length=512, **kwargs): + word_ids = [None] + list(range(len(words))) + [None] + return FakeEncoding(word_ids[:max_length]) + + +class StubTagger(torch.nn.Module): + + def __init__(self, tagged): + super().__init__() + self.anchor = torch.nn.Parameter(torch.zeros(1)) + self.tagged = tagged + self.crf = ConstrainedCRF(len(LABELS), batch_first=True) + with torch.no_grad(): + for parameter in self.crf.parameters(): + parameter.zero_() + + def emissions(self, input_ids, attention_mask): + scores = torch.zeros((1, input_ids.shape[1], len(LABELS))) + scores[:, :, LABEL2ID["O"]] = 4.0 + for word, label in self.tagged.items(): + scores[0, word + 1] = 0.0 + scores[0, word + 1, LABEL2ID[label]] = 9.0 + return scores + + +def test_words_match_dataset_tokenization(): + assert words_from_text("Apache-2.0 License") == ["Apache", "2", "0", "License"] + assert words_from_text("a\ufb01x\r\nnotice") == ["afix", "notice"] + + +def test_predicts_bioes_phrase_without_mutation(): + text = "Granted under the MIT License to everyone" + predictor = RequiredPhrasePredictor( + model=StubTagger({3: "B-REQ", 4: "E-REQ"}), + tokenizer=FakeTokenizer(), + max_length=512, + ) + + result = predictor.predict(text) + + assert text == "Granted under the MIT License to everyone" + assert [phrase.text for phrase in result.phrases] == ["MIT License"] + assert result.phrases[0].start_word == 3 + assert result.phrases[0].end_word == 4 + assert 0.0 <= result.phrases[0].confidence <= 1.0 + assert not result.truncated + + +def test_predicts_single_word_phrase(): + predictor = RequiredPhrasePredictor( + model=StubTagger({2: "S-REQ"}), + tokenizer=FakeTokenizer(), + max_length=512, + ) + + result = predictor.predict("one two MIT four") + + assert [phrase.text for phrase in result.phrases] == ["MIT"] + + +def test_drops_phrase_cut_by_truncation(): + predictor = RequiredPhrasePredictor( + model=StubTagger({2: "B-REQ", 3: "I-REQ"}), + tokenizer=FakeTokenizer(), + max_length=5, + ) + + result = predictor.predict("one two three four five six") + + assert result.truncated + assert result.phrases == () + + +def test_empty_text_does_not_run_model(): + predictor = RequiredPhrasePredictor( + model=StubTagger({}), + tokenizer=FakeTokenizer(), + max_length=512, + ) + + assert predictor.predict("").phrases == () From dc89f9a5d04928f36e9a9b3c7fca3ca6f071bb75 Mon Sep 17 00:00:00 2001 From: Kaushik Date: Mon, 7 Sep 2026 12:45:27 +0000 Subject: [PATCH 4/5] Document and test read-only inference Signed-off-by: Kaushik --- README.rst | 15 +++++++++++++++ azure-pipelines.yml | 3 ++- docs/source/training.rst | 17 +++++++++++++++++ tests/test_inference.py | 32 +++++++++++++++++++++++++++++++- 4 files changed, 65 insertions(+), 2 deletions(-) diff --git a/README.rst b/README.rst index 2e09892..a4c7050 100644 --- a/README.rst +++ b/README.rst @@ -51,6 +51,21 @@ Test-set evaluation is opt-in and should only be used for a final selected run. The completed model is written to ``model-output/final-model`` only after local reload and validation succeed. See ``docs/source/training.rst`` for details. +Run read-only prediction +======================== + +Load a validated final model and return candidate required phrases without +changing a ScanCode rule or file: + +.. code-block:: python + + from scancode_required_phrases.inference import RequiredPhrasePredictor + + predictor = RequiredPhrasePredictor.from_model_dir("model-output/final-model") + result = predictor.predict("Permission is hereby granted ...") + +Predictions require human review before they are added to license rules. + Development =========== diff --git a/azure-pipelines.yml b/azure-pipelines.yml index 1619850..bfea243 100644 --- a/azure-pipelines.yml +++ b/azure-pipelines.yml @@ -27,7 +27,8 @@ jobs: USE_TF=0 venv/bin/pytest -q \ tests/test_training.py \ tests/test_model.py \ - tests/test_export.py + tests/test_export.py \ + tests/test_inference.py displayName: Run training unit tests - template: etc/ci/azure-posix.yml diff --git a/docs/source/training.rst b/docs/source/training.rst index dd41f8f..2dc92df 100644 --- a/docs/source/training.rst +++ b/docs/source/training.rst @@ -62,6 +62,23 @@ ISR retains its existing name but measures predicted-phrase locatability only. It does not exercise injection gates or rule mutation and is not evidence of injection success. +Read-only prediction +-------------------- + +``RequiredPhrasePredictor`` loads only a final model that passes the publication +checks. Its ``predict()`` method uses the same ScanCode tokenization as the +training dataset and returns phrase text, word offsets, confidence, and whether +the input was truncated. It does not change rules or write files. + +.. code-block:: python + + from scancode_required_phrases.inference import RequiredPhrasePredictor + + predictor = RequiredPhrasePredictor.from_model_dir("model-output/final-model") + result = predictor.predict("Permission is hereby granted ...") + +Treat every prediction as a candidate requiring human review. + Export ------ diff --git a/tests/test_inference.py b/tests/test_inference.py index 299abf5..30cb1ed 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -1,8 +1,18 @@ # Copyright (c) nexB Inc. and others. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import torch +import json +import os +import pytest + +os.environ.setdefault("USE_TF", "0") + +torch = pytest.importorskip("torch") +pytest.importorskip("torchcrf") +pytest.importorskip("transformers") + +from scancode_required_phrases import inference from scancode_required_phrases.inference import RequiredPhrasePredictor from scancode_required_phrases.inference import words_from_text from scancode_required_phrases.model import ConstrainedCRF @@ -55,6 +65,26 @@ def test_words_match_dataset_tokenization(): assert words_from_text("a\ufb01x\r\nnotice") == ["afix", "notice"] +def test_loads_predictor_from_validated_model(tmp_path, monkeypatch): + model = StubTagger({}) + tokenizer = FakeTokenizer() + (tmp_path / "train_config.json").write_text( + json.dumps({"max_length": 384}), + encoding="utf-8", + ) + monkeypatch.setattr( + inference, + "load_final_model", + lambda model_dir: (model, tokenizer), + ) + + predictor = RequiredPhrasePredictor.from_model_dir(tmp_path) + + assert predictor.model is model + assert predictor.tokenizer is tokenizer + assert predictor.max_length == 384 + + def test_predicts_bioes_phrase_without_mutation(): text = "Granted under the MIT License to everyone" predictor = RequiredPhrasePredictor( From b2d9a3df6d5da5b11e4bdc182629150d7a314724 Mon Sep 17 00:00:00 2001 From: Kaushik Date: Mon, 7 Sep 2026 13:12:58 +0000 Subject: [PATCH 5/5] Reduce final model loading memory Signed-off-by: Kaushik --- src/scancode_required_phrases/training.py | 29 +++++++++++++++++++---- tests/test_training.py | 13 +++++++++- 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/src/scancode_required_phrases/training.py b/src/scancode_required_phrases/training.py index 0ac3ecb..b84bb42 100644 --- a/src/scancode_required_phrases/training.py +++ b/src/scancode_required_phrases/training.py @@ -1125,7 +1125,7 @@ def validate_state_structure(expected, actual, expected_name="expected", actual_ f"missing={missing}, unexpected={unexpected}" ) for name in expected_keys: - left = expected[name].detach().cpu() + left = expected[name].detach() right = actual[name].detach().cpu() if left.shape != right.shape: raise ValueError( @@ -1133,7 +1133,9 @@ def validate_state_structure(expected, actual, expected_name="expected", actual_ ) if left.dtype != right.dtype: raise ValueError(f"Tensor {name!r} dtype mismatch: {left.dtype} != {right.dtype}") - if not torch.isfinite(left).all() or not torch.isfinite(right).all(): + if left.device.type != "meta" and not torch.isfinite(left.cpu()).all(): + raise ValueError(f"Tensor {name!r} contains non-finite values") + if not torch.isfinite(right).all(): raise ValueError(f"Tensor {name!r} contains non-finite values") @@ -1248,7 +1250,11 @@ def _load_local_model(model_dir, offline=True): except (OSError, json.JSONDecodeError) as error: raise ValueError(f"Cannot read supported artifact configuration: {error}") from error local_config = AutoConfig.from_pretrained(str(model_dir), local_files_only=True) - backbone = AutoModel.from_config(local_config) + # Construct the large backbone without allocating random parameters. The + # saved tensors are assigned below, avoiding a second full FP32 model during + # inference startup on memory-constrained hosts. + with torch.device("meta"): + backbone = AutoModel.from_config(local_config) tagger_config = SimpleNamespace(**values) model = PhraseTagger(tagger_config, backbone=backbone) state_path = model_dir / "model.safetensors" @@ -1261,7 +1267,22 @@ def _load_local_model(model_dir, offline=True): "constructed", "saved", ) - model.load_state_dict(state, strict=True) + model.load_state_dict(state, strict=True, assign=True) + for module in model.modules(): + position_ids = getattr(module, "position_ids", None) + if position_ids is not None and position_ids.device.type == "meta": + module.register_buffer( + "position_ids", + torch.arange(position_ids.shape[-1]).expand(position_ids.shape), + persistent=False, + ) + meta_tensors = [ + name + for name, tensor in (*model.named_parameters(), *model.named_buffers()) + if tensor.device.type == "meta" + ] + if meta_tensors: + raise ValueError(f"Final_Model did not materialize tensors: {meta_tensors}") validate_state_dicts( _canonical_state_dict(state), _canonical_state_dict(model.state_dict()), diff --git a/tests/test_training.py b/tests/test_training.py index a9f1dee..7f7a687 100644 --- a/tests/test_training.py +++ b/tests/test_training.py @@ -907,9 +907,19 @@ def test_local_loader_loads_saved_values_after_structural_validation( str(model_dir / "model.safetensors"), ) + class LocalBackbone(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer( + "position_ids", + torch.arange(4).expand((1, 4)), + persistent=False, + ) + class LocalTagger(torch.nn.Module): def __init__(self, config, backbone=None): super().__init__() + self.backbone = backbone self.weight = torch.nn.Parameter(torch.tensor([0.0])) class LocalTokenizer: @@ -924,7 +934,7 @@ class LocalTokenizer: monkeypatch.setattr( transformers.AutoModel, "from_config", - lambda config: object(), + lambda config: LocalBackbone(), ) monkeypatch.setattr( transformers.AutoTokenizer, @@ -935,6 +945,7 @@ class LocalTokenizer: loaded, tokenizer = training._load_local_model(model_dir) assert loaded.weight.item() == 3.0 + assert loaded.backbone.position_ids.device.type == "cpu" assert tokenizer.is_fast