diff --git a/etc/requirements-ml.txt b/etc/requirements-ml.txt new file mode 100644 index 0000000000..af9d7c8945 --- /dev/null +++ b/etc/requirements-ml.txt @@ -0,0 +1,21 @@ +# Dependencies for training the required phrase NER model +# These are kept apart from scancode runtime requirements, they are only +# needed to train and export the model, not to run a scan +# +# install with: pip install -r etc/requirements-ml.txt + +torch>=2.0 +transformers>=4.44 +sentencepiece>=0.2 +protobuf>=3.20 +accelerate>=0.33 + +# CRF layer on top of the token classifier +pytorch-crf>=0.7.2 + +# 8 bit optimizer to fit deberta-large on a 16gb gpu, optional at runtime +bitsandbytes>=0.43 + +# ONNX export for CPU inference +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..12cd54a887 --- /dev/null +++ b/etc/scripts/dataset_pipeline/export_onnx.py @@ -0,0 +1,191 @@ +# exports a trained tagger to ONNX so it can run on CPU without torch +# +# the CRF layer cannot be put into ONNX, so we export only the backbone and +# the classifier (the emission scores) and save the learned CRF transition +# matrix beside it. at inference a small numpy viterbi decoder reproduces the +# exact CRF decoding from those transitions +import hashlib +import json +from pathlib import Path +from types import SimpleNamespace + +import click +import numpy as np + + +def viterbi_decode(emissions, start_transitions, transitions, end_transitions): + """Best valid tag path for one sequence, matches pytorch-crf decoding + + emissions is (seq_len, num_tags), transitions[i, j] scores tag i to tag j + """ + seq_len = emissions.shape[0] + score = start_transitions + emissions[0] + backpointers = [] + for step in range(1, seq_len): + candidates = score[:, None] + transitions + best_source = candidates.argmax(axis=0) + score = candidates.max(axis=0) + emissions[step] + backpointers.append(best_source) + + score = score + end_transitions + best = int(score.argmax()) + path = [best] + for sources in reversed(backpointers): + best = int(sources[best]) + path.append(best) + path.reverse() + return path + + +def sha256(path): + """Hex sha256 of a file""" + digest = hashlib.sha256() + digest.update(Path(path).read_bytes()) + return digest.hexdigest() + + +def build_emissions_module(tagger): + """Wrap the backbone and classifier as a plain emissions only model""" + import torch.nn as nn + + class EmissionsModule(nn.Module): + def __init__(self): + super().__init__() + self.backbone = tagger.backbone + self.classifier = tagger.classifier + + def forward(self, input_ids, attention_mask): + hidden = self.backbone( + input_ids=input_ids, + attention_mask=attention_mask, + ).last_hidden_state + return self.classifier(hidden) + + return EmissionsModule().eval() + + +def load_tagger(model_dir, train_config): + """Rebuild the tagger and load the finetuned weights""" + import torch + from phrase_model import PhraseTagger + + cfg = SimpleNamespace( + model_name=train_config['model_name'], + use_crf=train_config['use_crf'], + aux_ce_weight=0.0, + label_weights=[1.0] * len(train_config['labels']), + ) + tagger = PhraseTagger(cfg) + + safetensors_file = Path(model_dir) / 'model.safetensors' + if safetensors_file.exists(): + from safetensors.torch import load_file + state = load_file(str(safetensors_file)) + else: + state = torch.load(Path(model_dir) / 'pytorch_model.bin', map_location='cpu') + # class_weights is a training only buffer, it may be absent here + state = {k: v for k, v in state.items() if k != 'class_weights'} + tagger.load_state_dict(state) + return tagger.eval() + + +def check_viterbi_matches_crf(tagger, num_tags): + """Random emissions must decode the same in numpy and in pytorch-crf""" + import torch + + start = tagger.crf.start_transitions.detach().cpu().numpy() + transitions = tagger.crf.transitions.detach().cpu().numpy() + end = tagger.crf.end_transitions.detach().cpu().numpy() + + emissions = torch.randn(3, 14, num_tags) + mask = torch.ones(3, 14, dtype=torch.bool) + crf_paths = tagger.crf.decode(emissions, mask=mask) + for row in range(emissions.size(0)): + numpy_path = viterbi_decode(emissions[row].numpy(), start, transitions, end) + if numpy_path != crf_paths[row]: + raise AssertionError('numpy viterbi disagrees with pytorch-crf decoding') + + return start, transitions, end + + +def export(model_dir, output_dir, opset): + """Export to ONNX, dump CRF transitions and write a checksum manifest""" + import torch + from transformers import AutoTokenizer + + model_dir = Path(model_dir) + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + train_config = json.loads((model_dir / 'train_config.json').read_text()) + labels = train_config['labels'] + use_crf = train_config['use_crf'] + + tagger = load_tagger(model_dir, train_config) + emissions_module = build_emissions_module(tagger) + tokenizer = AutoTokenizer.from_pretrained(str(model_dir)) + + sample = tokenizer( + 'Licensed under the Apache License Version 2.0', + return_tensors='pt', + ) + inputs = (sample['input_ids'], sample['attention_mask']) + + onnx_path = output_dir / 'model.onnx' + torch.onnx.export( + emissions_module, + inputs, + str(onnx_path), + input_names=['input_ids', 'attention_mask'], + output_names=['emissions'], + dynamic_axes={ + 'input_ids': {0: 'batch', 1: 'sequence'}, + 'attention_mask': {0: 'batch', 1: 'sequence'}, + 'emissions': {0: 'batch', 1: 'sequence'}, + }, + opset_version=opset, + do_constant_folding=True, + ) + click.echo(f'wrote {onnx_path}') + + manifest = {'labels': labels, 'onnx_model': sha256(onnx_path)} + + if use_crf: + start, transitions, end = check_viterbi_matches_crf(tagger, len(labels)) + transitions_path = output_dir / 'crf_transitions.npz' + np.savez(transitions_path, start=start, transitions=transitions, end=end) + manifest['crf_transitions'] = sha256(transitions_path) + click.echo(f'wrote {transitions_path} (viterbi matches pytorch-crf)') + + # the onnx graph must produce the same emissions as pytorch + import onnxruntime + + session = onnxruntime.InferenceSession(str(onnx_path), providers=['CPUExecutionProvider']) + feeds = { + 'input_ids': sample['input_ids'].numpy(), + 'attention_mask': sample['attention_mask'].numpy(), + } + onnx_emissions = session.run(['emissions'], feeds)[0] + torch_emissions = emissions_module(*inputs).detach().numpy() + if not np.allclose(onnx_emissions, torch_emissions, atol=1e-3): + raise AssertionError('onnx emissions differ from pytorch emissions') + click.echo('onnx emissions match pytorch within tolerance') + + manifest_path = output_dir / 'manifest.json' + manifest_path.write_text(json.dumps(manifest, indent=2)) + click.echo(f'wrote {manifest_path}') + + +@click.command() +@click.option('--model-dir', required=True, type=click.Path(exists=True), + help='Directory with the trained model and train_config.json') +@click.option('--output-dir', default=None, + help='Where to write the ONNX files, defaults to the model dir') +@click.option('--opset', default=14, type=int, help='ONNX opset version') +def main(model_dir, output_dir, opset): + """Export a trained required phrase tagger to ONNX""" + export(model_dir, output_dir or model_dir, opset) + + +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..f1e0c12014 --- /dev/null +++ b/etc/scripts/dataset_pipeline/phrase_model.py @@ -0,0 +1,188 @@ +# the DeBERTa token classifier used to tag required phrases +# kept in its own module so train_model.py stays importable without torch +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 IGNORE_INDEX +from train_model import LABELS + + +class PhraseTagger(nn.Module): + """DeBERTa backbone with a token classifier and an optional CRF head + + The CRF scores the first subword of each word, so tagging happens at word + level and it decodes the best valid label sequence per rule + """ + + def __init__(self, config): + super().__init__() + self.use_crf = config.use_crf + self.aux_ce_weight = config.aux_ce_weight + self.num_labels = len(LABELS) + + # newer transformers loads the checkpoint in its stored fp16, force fp32 + # so the fresh classifier head matches and deberta stays numerically stable + self.backbone = AutoModel.from_pretrained(config.model_name).float() + self.backbone.gradient_checkpointing_enable() + # without this the checkpointed backbone can get no gradient and only the + # head trains, while the loss still looks fine + self.backbone.enable_input_require_grads() + + hidden_size = self.backbone.config.hidden_size + dropout = getattr(self.backbone.config, 'hidden_dropout_prob', 0.1) + self.dropout = nn.Dropout(dropout) + self.classifier = nn.Linear(hidden_size, self.num_labels) + + if self.use_crf: + self.crf = CRF(self.num_labels, batch_first=True) + + if self.aux_ce_weight > 0: + self.register_buffer( + 'class_weights', + torch.tensor(config.label_weights, dtype=torch.float), + ) + else: + self.class_weights = None + + def emissions(self, input_ids, attention_mask): + """Per subword label scores from the backbone""" + 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): + """Plain weighted cross entropy over subwords, ignores -100""" + 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 the first subword of each word into a dense left aligned sequence + + crf_tags pad with 0 so the CRF has valid indices, eval_tags pad with + IGNORE_INDEX so the metrics skip them + """ + batch, _, num_labels = emissions.shape + is_word = labels.ne(IGNORE_INDEX) + lengths = is_word.sum(dim=1) + width = int(lengths.max().item()) + + word_emissions = emissions.new_zeros((batch, width, num_labels)) + crf_tags = labels.new_zeros((batch, width)) + eval_tags = labels.new_full((batch, width), IGNORE_INDEX) + mask = torch.zeros((batch, width), dtype=torch.bool, device=emissions.device) + + for row in range(batch): + positions = is_word[row].nonzero(as_tuple=True)[0] + count = positions.numel() + word_emissions[row, :count] = emissions[row, positions] + tags = labels[row, positions] + crf_tags[row, :count] = tags + eval_tags[row, :count] = tags + mask[row, :count] = True + + return word_emissions, crf_tags, eval_tags, mask + + def forward(self, input_ids, attention_mask, labels=None): + emissions = self.emissions(input_ids, attention_mask) + result = {} + + if not self.use_crf: + if labels is not None: + result['loss'] = self.token_cross_entropy(emissions, labels) + result['word_labels'] = labels + if not self.training: + result['predictions'] = emissions.argmax(dim=-1) + return result + + if labels is None: + raise ValueError('CRF head needs labels to locate words, use the ONNX export for inference') + + word_emissions, crf_tags, eval_tags, mask = self.gather_words(emissions, labels) + # the CRF math is more stable in fp32 under mixed precision + 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 pad_decoded(self, decoded, width, device): + """Turn the variable length CRF paths into a padded tensor""" + preds = torch.full((len(decoded), width), IGNORE_INDEX, dtype=torch.long, device=device) + for row, path in enumerate(decoded): + if path: + preds[row, :len(path)] = torch.tensor(path, dtype=torch.long, device=device) + return preds + + +def build_optimizer(config, model): + """AdamW with layer wise learning rate decay + + The head learns fastest and the embeddings slowest so the pretrained lower + layers keep more of what they know + """ + 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)) + # embeddings and relative position embeddings sit at the bottom + return config.base_lr * (config.layer_decay ** (num_layers + 1)) + + groups = [] + for name, param in model.named_parameters(): + if not param.requires_grad: + continue + decay = 0.0 if any(nd in name for nd in no_decay) else config.weight_decay + groups.append({'params': [param], 'lr': rate_for(name), 'weight_decay': decay}) + + # the 8 bit optimizer keeps deberta-large inside a 16gb gpu, fall back to + # plain AdamW where bitsandbytes is not installed + try: + from bitsandbytes.optim import AdamW8bit + return AdamW8bit(groups, lr=config.base_lr, eps=config.adam_epsilon, betas=(0.9, 0.999)) + except ImportError: + return AdamW(groups, lr=config.base_lr, eps=config.adam_epsilon, betas=(0.9, 0.999)) + + +class PhraseTrainer(Trainer): + """Trainer that reads loss and predictions from our model output dict""" + + def compute_loss(self, model, inputs, return_outputs=False, **kwargs): + outputs = model(**inputs) + loss = outputs['loss'] + return (loss, outputs) if return_outputs else loss + + def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None): + inputs = self._prepare_inputs(inputs) + with torch.no_grad(): + outputs = model(**inputs) + loss = outputs.get('loss') + if loss is not None: + loss = loss.detach() + if prediction_loss_only: + return (loss, None, None) + return (loss, outputs['predictions'], outputs['word_labels']) diff --git a/etc/scripts/dataset_pipeline/test_train_model.py b/etc/scripts/dataset_pipeline/test_train_model.py new file mode 100644 index 0000000000..1b0919de9c --- /dev/null +++ b/etc/scripts/dataset_pipeline/test_train_model.py @@ -0,0 +1,132 @@ +# tests for train_model.py and export_onnx.py +# these cover the pure logic, the heavy model parts are exercised on a gpu +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent)) +from train_model import align_labels +from train_model import decode_row +from train_model import extract_spans +from train_model import IGNORE_INDEX +from train_model import LABEL2ID + + +class FakeEncoding(dict): + """Mimics a tokenizer encoding with a fixed word_ids mapping""" + + def __init__(self, word_ids): + super().__init__() + self._word_ids = word_ids + self['input_ids'] = [0] * len(word_ids) + self['attention_mask'] = [1] * len(word_ids) + + def word_ids(self): + return self._word_ids + + +class FakeTokenizer: + def __init__(self, word_ids): + self._word_ids = word_ids + + def __call__(self, tokens, **kwargs): + return FakeEncoding(self._word_ids) + + +class TestExtractSpans: + + def test_single_phrase(self): + assert extract_spans(['O', 'B-REQ', 'I-REQ', 'E-REQ', 'O']) == {(1, 3)} + + def test_single_token_phrase(self): + assert extract_spans(['O', 'S-REQ', 'O']) == {(1, 1)} + + def test_two_phrases(self): + assert extract_spans(['S-REQ', 'O', 'B-REQ', 'E-REQ']) == {(0, 0), (2, 3)} + + def test_no_phrase(self): + assert extract_spans(['O', 'O', 'O']) == set() + + def test_handles_inside_without_begin(self): + # raw softmax output can be malformed, an I with no B still yields a span + assert extract_spans(['O', 'I-REQ', 'E-REQ']) == {(1, 2)} + + def test_handles_unterminated_begin(self): + assert extract_spans(['O', 'B-REQ', 'I-REQ']) == {(1, 2)} + + +class TestDecodeRow: + + def test_drops_ignored_positions(self): + preds = [1, 0, 2, 0] + labels = [1, IGNORE_INDEX, 2, IGNORE_INDEX] + pred_tags, true_tags = decode_row(preds, labels) + assert true_tags == ['B-REQ', 'I-REQ'] + assert pred_tags == ['B-REQ', 'I-REQ'] + + def test_unknown_prediction_falls_back_to_o(self): + pred_tags, true_tags = decode_row([99], [0]) + assert true_tags == ['O'] + assert pred_tags == ['O'] + + +class TestAlignLabels: + + def test_first_subword_keeps_label(self): + # CLS Apache Lic ##ense SEP, Apache is one subword, License is two + tokenizer = FakeTokenizer([None, 0, 1, 1, None]) + encoding = align_labels(['Apache', 'License'], ['B-REQ', 'E-REQ'], tokenizer, 512) + assert encoding['labels'] == [ + IGNORE_INDEX, LABEL2ID['B-REQ'], LABEL2ID['E-REQ'], IGNORE_INDEX, IGNORE_INDEX, + ] + + def test_all_outside(self): + tokenizer = FakeTokenizer([None, 0, 1, None]) + encoding = align_labels(['the', 'license'], ['O', 'O'], tokenizer, 512) + assert encoding['labels'] == [IGNORE_INDEX, 0, 0, IGNORE_INDEX] + + +def test_viterbi_with_zero_transitions_is_argmax(): + import numpy as np + from export_onnx import viterbi_decode + + 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]]) + zeros_t = np.zeros((5, 5)) + zeros_v = np.zeros(5) + path = viterbi_decode(emissions, zeros_v, zeros_t, zeros_v) + assert path == [1, 0, 4] + + +def test_viterbi_avoids_forbidden_transition(): + import numpy as np + from export_onnx import viterbi_decode + + # emissions want tag 1 then tag 0, but 1 -> 0 is heavily penalised + emissions = np.array([[0.0, 1.0], [1.0, 0.0]]) + transitions = np.array([[0.0, 0.0], [-100.0, 0.0]]) + zeros_v = np.zeros(2) + path = viterbi_decode(emissions, zeros_v, transitions, zeros_v) + assert path[0] == path[1] + + +def test_compute_metrics_perfect_and_exact(): + from train_model import compute_metrics + + # two rules, word level predictions and labels line up perfectly + predictions = [[1, 3, 0], [4, IGNORE_INDEX, IGNORE_INDEX]] + labels = [[1, 3, 0], [4, IGNORE_INDEX, IGNORE_INDEX]] + scores = compute_metrics((predictions, labels)) + assert scores['f1'] == 1.0 + assert scores['exact_match'] == 1.0 + + +def test_compute_metrics_counts_a_miss(): + from train_model import compute_metrics + + # predict O everywhere, the gold has one phrase, so f1 is 0 + predictions = [[0, 0, 0]] + labels = [[1, 3, 0]] + scores = compute_metrics((predictions, labels)) + assert scores['f1'] == 0.0 + assert scores['exact_match'] == 0.0 diff --git a/etc/scripts/dataset_pipeline/train_model.py b/etc/scripts/dataset_pipeline/train_model.py new file mode 100644 index 0000000000..aa973b7951 --- /dev/null +++ b/etc/scripts/dataset_pipeline/train_model.py @@ -0,0 +1,400 @@ +# finetunes a DeBERTa token classifier to predict required phrase spans +# reads the BIOES JSONL that build_dataset.py produces +# torch and transformers are imported lazily so the tests can run without them +import inspect +import json +import random +from dataclasses import dataclass +from dataclasses import field +from pathlib import Path + +import click + + +# BIOES scheme, single source of truth shared by the model and the exporter +LABELS = ['O', 'B-REQ', 'I-REQ', 'E-REQ', 'S-REQ'] +LABEL2ID = {label: i for i, label in enumerate(LABELS)} +ID2LABEL = {i: label for i, label in enumerate(LABELS)} + +# label id ignored by the loss for padding and non first subwords +IGNORE_INDEX = -100 + +MODEL_NAME = 'microsoft/deberta-v3-large' +MAX_LENGTH = 512 + + +@dataclass +class Config: + """Holds the settings for one training run""" + data_dir: Path + output_dir: Path + model_name: str = MODEL_NAME + max_length: int = MAX_LENGTH + + epochs: int = 8 + batch_size: int = 1 + grad_accum: int = 16 + base_lr: float = 2e-5 + # the classifier and crf start from scratch so they learn at a higher rate + 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 + + # cap examples per split for a quick smoke test, 0 uses everything + limit: int = 0 + # pick up from the last saved checkpoint in output_dir if a run died + resume: bool = False + + # turn off to train a plain softmax tagger for ablation + use_crf: bool = True + # light weighted cross entropy on top of the CRF loss to protect the rare + # boundary labels against the O heavy imbalance, set 0 for pure CRF + aux_ce_weight: float = 0.3 + # report injection success rate too, needs scancode importable + with_isr: bool = False + + seed: int = 42 + + # used only when aux_ce_weight is set, one weight per BIOES label + label_weights: list = field(default_factory=lambda: [1.0, 2.0, 1.5, 1.5, 2.0]) + + +def set_seed(seed): + """Seed python, numpy and torch for reproducible runs""" + import numpy as np + import torch + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def load_jsonl(path): + """Yield parsed records from a JSONL file""" + with open(path, encoding='utf-8') as f: + for line in f: + line = line.strip() + if line: + yield json.loads(line) + + +def align_labels(tokens, word_labels, tokenizer, max_length): + """Tokenize words into subwords and put each label on its first subword + + Continuation subwords and special tokens get IGNORE_INDEX so the loss + skips them + """ + encoding = tokenizer( + tokens, + is_split_into_words=True, + truncation=True, + max_length=max_length, + ) + + label_ids = [] + previous_word = None + for word_id in encoding.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 + + +class PhraseDataset: + """Reads a BIOES JSONL split and encodes each rule for the model""" + + def __init__(self, path, tokenizer, max_length, limit=0): + self.examples = [] + self.truncated = 0 + for record in load_jsonl(path): + if limit and len(self.examples) >= limit: + break + tokens = record['tokens'] + labels = record['bioes_labels'] + if not tokens: + continue + encoding = align_labels(tokens, labels, tokenizer, max_length) + kept = sum(1 for lid in encoding['labels'] if lid != IGNORE_INDEX) + if kept < len(tokens): + self.truncated += 1 + self.examples.append({ + 'input_ids': encoding['input_ids'], + 'attention_mask': encoding['attention_mask'], + 'labels': encoding['labels'], + }) + + def __len__(self): + return len(self.examples) + + def __getitem__(self, index): + return self.examples[index] + + +def extract_spans(tags): + """Return the set of (start, end) word spans in a BIOES tag sequence + + Tolerates malformed sequences so it also works on raw softmax output + """ + spans = [] + start = None + for i, tag in enumerate(tags): + if tag == 'S-REQ': + spans.append((i, i)) + start = None + elif tag == 'B-REQ': + if start is not None: + spans.append((start, i - 1)) + start = i + elif tag == 'I-REQ': + if start is None: + start = i + elif tag == 'E-REQ': + if start is None: + start = i + spans.append((start, i)) + start = None + else: + if start is not None: + spans.append((start, i - 1)) + start = None + if start is not None: + spans.append((start, len(tags) - 1)) + return set(spans) + + +def decode_row(pred_row, label_row): + """Drop the ignored positions and map ids back to BIOES tags""" + pred_tags = [] + true_tags = [] + for pred, label in zip(pred_row, label_row): + if int(label) == IGNORE_INDEX: + continue + true_tags.append(ID2LABEL[int(label)]) + pred_tags.append(ID2LABEL.get(int(pred), 'O')) + return pred_tags, true_tags + + +def compute_metrics(eval_pred): + """Strict entity level micro F1 plus a rule level exact match + + With a single phrase type this is just the span set overlap, so we score + it directly and skip the seqeval dependency + """ + predictions, labels = eval_pred + tp = fp = fn = 0 + exact = 0 + for pred_row, label_row in zip(predictions, labels): + predicted, actual = decode_row(pred_row, label_row) + pred_spans = extract_spans(predicted) + true_spans = extract_spans(actual) + tp += len(pred_spans & true_spans) + fp += len(pred_spans - true_spans) + fn += len(true_spans - pred_spans) + if pred_spans == true_spans: + exact += 1 + + precision = tp / (tp + fp) if tp + fp else 0.0 + recall = tp / (tp + fn) if tp + fn else 0.0 + f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0 + return { + 'f1': f1, + 'precision': precision, + 'recall': recall, + 'exact_match': exact / len(predictions) if len(predictions) else 0.0, + } + + +def evaluate_isr(records, model, tokenizer, max_length): + """Fraction of predicted phrases scancode can still locate in the rule text + + A phrase that cannot be located is not injectable, so this is the metric + that actually matters + """ + import torch + from licensedcode.required_phrases import find_phrase_spans_in_text + + device = next(model.parameters()).device + model.eval() + total = 0 + injectable = 0 + for record in records: + tokens = record['tokens'] + text = record['text'] + if not tokens: + continue + encoding = align_labels(tokens, record['bioes_labels'], tokenizer, max_length) + inputs = { + 'input_ids': torch.tensor([encoding['input_ids']], device=device), + 'attention_mask': torch.tensor([encoding['attention_mask']], device=device), + 'labels': torch.tensor([encoding['labels']], device=device), + } + with torch.no_grad(): + output = model(**inputs) + tags, _ = decode_row(output['predictions'][0].tolist(), output['word_labels'][0].tolist()) + for start, end in extract_spans(tags): + if end >= len(tokens): + continue + phrase = ' '.join(tokens[start:end + 1]) + total += 1 + if find_phrase_spans_in_text(text, phrase): + injectable += 1 + + return injectable / total if total else 0.0 + + +def run_training(config): + """Train, evaluate and save the model, callable from a notebook or main()""" + import torch + 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 + + config.output_dir.mkdir(parents=True, exist_ok=True) + set_seed(config.seed) + + head = 'CRF' if config.use_crf else 'softmax' + click.echo(f'training {config.model_name} with a {head} head') + + tokenizer = AutoTokenizer.from_pretrained(config.model_name, use_fast=True) + if not tokenizer.is_fast: + raise RuntimeError('need a fast tokenizer for word_ids, got a slow one') + train_ds = PhraseDataset(config.data_dir / 'train.jsonl', tokenizer, config.max_length, config.limit) + val_ds = PhraseDataset(config.data_dir / 'val.jsonl', tokenizer, config.max_length, config.limit) + test_ds = PhraseDataset(config.data_dir / 'test.jsonl', tokenizer, config.max_length, config.limit) + + click.echo(f'examples train: {len(train_ds)} val: {len(val_ds)} test: {len(test_ds)}') + if train_ds.truncated: + click.echo(f'note: {train_ds.truncated} train rules were longer than {config.max_length} subwords and got truncated') + + model = PhraseTagger(config) + collator = DataCollatorForTokenClassification(tokenizer, label_pad_token_id=IGNORE_INDEX) + + # deberta breaks in fp16, so use bf16 only where the gpu supports it else fp32 + use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported() + + args = 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=use_bf16, + fp16=False, + logging_steps=50, + report_to='none', + seed=config.seed, + dataloader_num_workers=2, + ) + + trainer_kwargs = dict( + model=model, + args=args, + train_dataset=train_ds, + eval_dataset=val_ds, + data_collator=collator, + compute_metrics=compute_metrics, + optimizers=(build_optimizer(config, model), None), + callbacks=[EarlyStoppingCallback(early_stopping_patience=config.early_stopping_patience)], + ) + + # transformers renamed the tokenizer arg to processing_class, support both + if 'processing_class' in inspect.signature(PhraseTrainer.__init__).parameters: + trainer_kwargs['processing_class'] = tokenizer + else: + trainer_kwargs['tokenizer'] = tokenizer + + trainer = PhraseTrainer(**trainer_kwargs) + + trainer.train(resume_from_checkpoint=config.resume or None) + trainer.save_model(str(config.output_dir)) + + # the exporter reads this back to rebuild the model + train_config = config.output_dir / 'train_config.json' + train_config.write_text(json.dumps({ + 'model_name': config.model_name, + 'use_crf': config.use_crf, + 'max_length': config.max_length, + 'labels': LABELS, + }, indent=2)) + + test_metrics = trainer.evaluate(test_ds, metric_key_prefix='test') + click.echo(f'test: {test_metrics}') + + if config.with_isr: + records = list(load_jsonl(config.data_dir / 'test.jsonl')) + isr = evaluate_isr(records, model, tokenizer, config.max_length) + click.echo(f'injection success rate: {isr:.4f}') + test_metrics['test_isr'] = isr + + return test_metrics + + +@click.command() +@click.option('--data-dir', required=True, type=click.Path(exists=True), + help='Directory holding train.jsonl, val.jsonl and test.jsonl') +@click.option('--output-dir', default='model-output', + help='Where to write the trained model and metrics') +@click.option('--model-name', default=MODEL_NAME, help='Base model to finetune') +@click.option('--epochs', default=8, type=int) +@click.option('--base-lr', default=2e-5, type=float, help='Learning rate for the backbone') +@click.option('--head-lr', default=1e-4, type=float, help='Learning rate for the classifier and crf head') +@click.option('--aux-ce-weight', default=0.3, type=float, + help='Weight of the auxiliary weighted cross entropy, 0 for pure CRF') +@click.option('--no-crf', is_flag=True, default=False, + help='Train a plain softmax tagger instead of the CRF head') +@click.option('--with-isr', is_flag=True, default=False, + help='Also report injection success rate, needs scancode installed') +@click.option('--limit', default=0, type=int, + help='Cap examples per split for a quick smoke test, 0 uses everything') +@click.option('--resume', is_flag=True, default=False, + help='Resume from the last checkpoint in output-dir after a crash') +@click.option('--seed', default=42, type=int) +def main(data_dir, output_dir, model_name, epochs, base_lr, head_lr, aux_ce_weight, + no_crf, with_isr, limit, resume, seed): + """Train the required phrase tagger from a BIOES dataset""" + config = Config( + data_dir=Path(data_dir), + output_dir=Path(output_dir), + model_name=model_name, + epochs=epochs, + base_lr=base_lr, + head_lr=head_lr, + aux_ce_weight=aux_ce_weight, + use_crf=not no_crf, + with_isr=with_isr, + limit=limit, + resume=resume, + seed=seed, + ) + run_training(config) + + +if __name__ == '__main__': + main()