diff --git a/etc/requirements-ml.txt b/etc/requirements-ml.txt new file mode 100644 index 0000000000..319fddb28e --- /dev/null +++ b/etc/requirements-ml.txt @@ -0,0 +1,19 @@ +# Dependencies for training and exporting the required phrase model. +# Install with: pip install -r etc/requirements-ml.txt + +torch>=2.0 +transformers==4.57.3 +huggingface-hub==0.36.2 +sentencepiece>=0.2 +protobuf>=3.20 +accelerate>=0.33 +pytorch-crf==0.7.2 +safetensors>=0.4 + +# Used explicitly for DeBERTa-large training on a 16 GB GPU. +bitsandbytes>=0.43 + +# ONNX export and CPU verification. +numpy>=1.24 +onnx>=1.16 +onnxruntime>=1.18 diff --git a/etc/scripts/dataset_pipeline/export_onnx.py b/etc/scripts/dataset_pipeline/export_onnx.py new file mode 100644 index 0000000000..0fbd405397 --- /dev/null +++ b/etc/scripts/dataset_pipeline/export_onnx.py @@ -0,0 +1,595 @@ +# -*- coding: utf-8 -*- +# +# Copyright (c) nexB Inc. and others. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Export a completed required phrase model for CPU inference.""" + +import hashlib +import json +import os +from pathlib import Path +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", +) + + +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): + 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) + + 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): + 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 _training_interfaces(): + """Import final-artifact interfaces without importing optional ONNX code.""" + from train_model import ARTIFACT_SCHEMA + from train_model import CONSTRAINT_CONTRACT + from train_model import LABELS + from train_model import load_final_model + from train_model import validate_bioes + from train_model import validate_publishable_model + from train_model import write_json_atomic + + return ( + ARTIFACT_SCHEMA, + CONSTRAINT_CONTRACT, + tuple(LABELS), + load_final_model, + validate_bioes, + validate_publishable_model, + write_json_atomic, + ) + + +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) + + 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 + + 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,), + ), + ) + 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 + + 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,)), + ) + + +def _exact_one_word_tie(start, end): + """Return emissions that produce an exact tie under exported edge scores.""" + import numpy as np + + 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 check_viterbi_matches_crf(tagger, num_tags): + """Verify exact NumPy/PyTorch constrained paths on deterministic cases.""" + import numpy as np + import torch + + ( + _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 = { + "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 + + 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), + ) + _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 + + +def build_emissions_module(tagger): + """Wrap only emissions computation for ONNX export.""" + import torch.nn as nn + + 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() +@click.option( + "--model-dir", + required=True, + type=click.Path(exists=True, file_okay=False, path_type=Path), + help="Completed local Final_Model directory.", +) +@click.option( + "--output-dir", + required=True, + type=click.Path(file_okay=False, path_type=Path), + 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, operation, opset): + """Export constrained matrices or optional ONNX emissions.""" + try: + 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__": + main() diff --git a/etc/scripts/dataset_pipeline/phrase_model.py b/etc/scripts/dataset_pipeline/phrase_model.py new file mode 100644 index 0000000000..29d2d04d5c --- /dev/null +++ b/etc/scripts/dataset_pipeline/phrase_model.py @@ -0,0 +1,445 @@ +# -*- 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 train_model import first_subword_positions +from train_model import IGNORE_INDEX +from train_model 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, backbone=None): + super().__init__() + self.use_crf = config.use_crf + self.aux_ce_weight = config.aux_ce_weight + self.num_labels = len(LABELS) + + 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) + self.dropout = nn.Dropout(dropout) + self.classifier = nn.Linear(hidden_size, self.num_labels) + + if self.use_crf: + self.crf = ConstrainedCRF(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) + 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)) + 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) + + 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: + raise ValueError("Inference word sequence must not be empty") + + 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).all(): + 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.mean().detach() + if prediction_loss_only: + return loss, None, None + return loss, outputs["predictions"], outputs["word_labels"] diff --git a/etc/scripts/dataset_pipeline/test_export_onnx.py b/etc/scripts/dataset_pipeline/test_export_onnx.py new file mode 100644 index 0000000000..414d4ceb9e --- /dev/null +++ b/etc/scripts/dataset_pipeline/test_export_onnx.py @@ -0,0 +1,514 @@ +# -*- coding: utf-8 -*- +# +# 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") + +import export_onnx as export_module +import train_model as training +from export_onnx import check_viterbi_matches_crf +from export_onnx import export_crf_matrices +from export_onnx import export_onnx_emissions +from export_onnx import main +from export_onnx import _load_publishable_artifact +from export_onnx import OnnxDependencyError +from export_onnx import PARITY_CASES +from export_onnx import sha256 +from export_onnx 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): + import train_model as 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( + [ + [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, 5) == [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_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( + "export_onnx._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("export_onnx.export_crf_matrices", export_crf) + monkeypatch.setattr("export_onnx.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") + + assert sha256(path) == "9372c470eeadd5ecd9c3c74c2b3cb633f8e2f2fad799250a0f70d652b6b825e4" diff --git a/etc/scripts/dataset_pipeline/test_phrase_model.py b/etc/scripts/dataset_pipeline/test_phrase_model.py new file mode 100644 index 0000000000..075cfd63a1 --- /dev/null +++ b/etc/scripts/dataset_pipeline/test_phrase_model.py @@ -0,0 +1,358 @@ +# -*- 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") + +import phrase_model as model_module +from phrase_model import build_constraint_masks +from phrase_model import build_optimizer +from phrase_model import ConstrainedCRF +from phrase_model import PhraseTagger +from phrase_model import PhraseTrainer +from train_model import LABEL2ID +from train_model 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) + self.gradient_checkpointing_kwargs = None + + def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): + self.gradient_checkpointing_kwargs = gradient_checkpointing_kwargs + + +@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) + 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): + 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(): + 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 = ConstrainedCRF(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_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) + + 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/etc/scripts/dataset_pipeline/test_train_model.py b/etc/scripts/dataset_pipeline/test_train_model.py new file mode 100644 index 0000000000..15e279ba27 --- /dev/null +++ b/etc/scripts/dataset_pipeline/test_train_model.py @@ -0,0 +1,1224 @@ +# -*- coding: utf-8 -*- +# +# Copyright (c) nexB Inc. and others. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +from pathlib import Path +import types + +from click.testing import CliRunner +import pytest + +import train_model as training +from train_model import align_labels +from train_model import AlignmentError +from train_model import ARTIFACT_SCHEMA +from train_model import build_effective_datasets +from train_model import compare_ordered_validation_results +from train_model import compute_metrics +from train_model import Config +from train_model import CONSTRAINT_CONTRACT +from train_model import decode_row +from train_model import extract_spans +from train_model import first_subword_positions +from train_model import IGNORE_INDEX +from train_model import LABEL2ID +from train_model import LABELS +from train_model import load_and_validate_splits +from train_model import main +from train_model import prepare_output_dir +from train_model import promote_final_model +from train_model import resolve_model_identity +from train_model import resolve_tokenizer_revision +from train_model import serializable_config +from train_model import sha256 +from train_model import validate_bioes +from train_model import validate_config +from train_model import validate_publishable_model +from train_model import validate_record +from train_model import validate_saved_state +from train_model import validate_raw_split_hashes +from train_model import validate_state_dicts +from train_model import validate_state_structure +from train_model import write_failure_manifest +from train_model 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 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): + return self._word_ids + + +class FakeTokenizer: + 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): + return { + "identifier": identifier, + "license_expression": "mit", + "rule_type": "is_license_notice", + "text": "MIT License terms apply", + "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"], + } + + +def write_jsonl(path, records): + path.write_text( + "".join(json.dumps(record) + "\n" for record in records), + encoding="utf-8", + ) + + +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", + [ + ["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_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) + + 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] + with pytest.raises(ValueError, match=field): + validate_record(record, "train.jsonl", 9) + + +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_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_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, + LABEL2ID["B-REQ"], + LABEL2ID["E-REQ"], + IGNORE_INDEX, + IGNORE_INDEX, + ] + assert not truncated and not cut + + +@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], + ) + encoding, truncated, _cut = align_labels( + ["MIT", "License", "terms", "apply"], + ["B-REQ", "E-REQ", "O", "O"], + tokenizer, + 4, + ) + assert truncated + 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, + ) + + 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_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"), + [ + (["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_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_is_strict_and_drops_only_ignored_gold_positions(): + predicted, actual = decode_row( + [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_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), 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_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_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 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_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_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) + + 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) + + +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_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_tokenizer_revision_uses_the_requested_hub_commit(monkeypatch): + calls = [] + + 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_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") + import phrase_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"), + ) + + 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 + + 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("phrase_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, "phrase_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("phrase_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, "phrase_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) diff --git a/etc/scripts/dataset_pipeline/train_model.py b/etc/scripts/dataset_pipeline/train_model.py new file mode 100644 index 0000000000..d7ed4d63e2 --- /dev/null +++ b/etc/scripts/dataset_pipeline/train_model.py @@ -0,0 +1,1978 @@ +# -*- 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 +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") +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", + "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]) + + +@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}") + + 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) + 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) + + +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 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}" + 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]}" + return None + + +def validate_record(record, path, line_number): + """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"] + 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" + ) + 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): + """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, + 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") + + 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 + return encoding, bool(omitted_positions), False + + +def first_subword_positions(word_ids): + """Return positions that start each contiguous 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: + """Hold selected effective examples for one validated split.""" + + 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 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 + self.truncations.append( + {"identifier": validated.identifier, **inventory["location"]} + ) + + 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( + {name: inventory[name] for name in ("input_ids", "attention_mask", "labels")} + ) + + def __len__(self): + return len(self.examples) + + def __getitem__(self, index): + return self.examples[index] + + +def validate_splits(datasets): + """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 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 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.add((index, index)) + elif tag == "B-REQ": + start = index + elif tag == "E-REQ": + spans.add((start, index)) + start = None + 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, 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 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 + 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 _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 = 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 + + predicted_spans = extract_spans(predicted) + 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), + "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 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 = 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), + "attention_mask": torch.tensor([encoding["attention_mask"]], device=device), + "labels": torch.tensor([encoding["labels"]], device=device), + } + with torch.no_grad(): + output = model(**inputs) + 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): + locatable += 1 + return locatable / 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 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): + """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 transformers import AutoConfig + from transformers import AutoModel + from transformers import AutoTokenizer + + from phrase_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) + + +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( + "Ordered validation predictions and labels must have equal single-row shapes" + ) + 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() + ): + raise ValueError("bf16 requires a CUDA device with BF16 support") + + +def run_training(config): + """Validate, train, verify, and transactionally publish a Final_Model.""" + import torch + from transformers import AutoConfig + from transformers import AutoTokenizer + from transformers import DataCollatorForTokenClassification + from transformers import EarlyStoppingCallback + from transformers import TrainingArguments + + from phrase_model import PhraseTagger + from phrase_model import PhraseTrainer + from phrase_model import build_optimizer + + validate_config(config) + validate_precision(config.precision) + 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, + revision=config.model_revision, + use_fast=True, + ) + if not tokenizer.is_fast: + raise RuntimeError("Training requires a fast tokenizer with word IDs") + tokenizer_revision = resolve_tokenizer_revision( + config.model_name, config.model_revision + ) + backbone_config = AutoConfig.from_pretrained( + config.model_name, + revision=config.model_revision, + ) + resolved_revision = resolve_model_identity( + config.model_revision, tokenizer_revision, backbone_config + ) + datasets, h2 = build_effective_datasets( + records_by_split, tokenizer, config.max_length, config.limit + ) + 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": { + "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_identity": { + "name": config.model_name, + "requested_revision": config.model_revision, + "resolved_tokenizer_revision": resolved_revision, + "resolved_backbone_revision": resolved_revision, + }, + "source": source_provenance, + "runtime": runtime_provenance, + "completed_checks": [ + "configuration", + "raw-splits", + "identifier-uniqueness", + "content-duplicate-report", + "h0-h1", + "immutable-model-identity", + "effective-examples-h2", + ], + } + write_json_atomic(manifest_path, manifest) + + 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, + } + + +@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="New or empty directory for this run.", +) +@click.option("--model-name", default=MODEL_NAME, help="Base model to fine-tune.") +@click.option( + "--model-revision", + required=True, + help="Full immutable model and tokenizer commit revision.", +) +@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=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( + "--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, help="Evaluate test data after selection.") +@click.option( + "--with-isr", is_flag=True, + help="Report predicted-phrase locatability during final test evaluation.", +) +@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, +): + """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, + 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()