From 6bc547c69b256c5e39df22c3589a91810519fe9a Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Fri, 3 Jul 2026 16:50:50 +0530
Subject: [PATCH 01/12] Add ML requirements for required phrase training

---
 etc/requirements-ml.txt | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)
 create mode 100644 etc/requirements-ml.txt

diff --git a/etc/requirements-ml.txt b/etc/requirements-ml.txt
new file mode 100644
index 0000000000..4d1e3e9461
--- /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.2
+sentencepiece==0.2.0
+protobuf>=3.20
+accelerate==0.33.0
+
+# CRF layer on top of the token classifier
+pytorch-crf==0.7.2
+
+# span level scoring
+seqeval==1.2.2
+
+# ONNX export for CPU inference
+onnx==1.16.2
+onnxruntime==1.18.1

From 0b6792eb9af0a86f7d58daf34b2c71089a6c8af1 Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Fri, 3 Jul 2026 16:50:50 +0530
Subject: [PATCH 02/12] Add training script with data loading and metrics

---
 etc/scripts/dataset_pipeline/train_model.py | 363 ++++++++++++++++++++
 1 file changed, 363 insertions(+)
 create mode 100644 etc/scripts/dataset_pipeline/train_model.py

diff --git a/etc/scripts/dataset_pipeline/train_model.py b/etc/scripts/dataset_pipeline/train_model.py
new file mode 100644
index 0000000000..d21f0c8bb1
--- /dev/null
+++ b/etc/scripts/dataset_pipeline/train_model.py
@@ -0,0 +1,363 @@
+# 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 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
+    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
+
+    # turn off to train a plain softmax tagger for ablation
+    use_crf: bool = True
+    # extra weighted cross entropy added to the CRF loss, 0 turns it off
+    aux_ce_weight: float = 0.0
+    # 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):
+        self.examples = []
+        self.truncated = 0
+        for record in load_jsonl(path):
+            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):
+    """Entity level F1 plus a strict rule level exact match"""
+    from seqeval.metrics import f1_score
+    from seqeval.metrics import precision_score
+    from seqeval.metrics import recall_score
+    from seqeval.scheme import IOBES
+
+    predictions, labels = eval_pred
+    pred_tags = []
+    true_tags = []
+    exact = 0
+    for pred_row, label_row in zip(predictions, labels):
+        predicted, actual = decode_row(pred_row, label_row)
+        pred_tags.append(predicted)
+        true_tags.append(actual)
+        if extract_spans(predicted) == extract_spans(actual):
+            exact += 1
+
+    scored = dict(mode='strict', scheme=IOBES)
+    return {
+        'f1': f1_score(true_tags, pred_tags, **scored),
+        'precision': precision_score(true_tags, pred_tags, **scored),
+        'recall': recall_score(true_tags, pred_tags, **scored),
+        '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)
+    train_ds = PhraseDataset(config.data_dir / 'train.jsonl', tokenizer, config.max_length)
+    val_ds = PhraseDataset(config.data_dir / 'val.jsonl', tokenizer, config.max_length)
+    test_ds = PhraseDataset(config.data_dir / 'test.jsonl', tokenizer, config.max_length)
+
+    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 = PhraseTrainer(
+        model=model,
+        args=args,
+        train_dataset=train_ds,
+        eval_dataset=val_ds,
+        data_collator=collator,
+        tokenizer=tokenizer,
+        compute_metrics=compute_metrics,
+        optimizers=(build_optimizer(config, model), None),
+        callbacks=[EarlyStoppingCallback(early_stopping_patience=config.early_stopping_patience)],
+    )
+
+    trainer.train()
+    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('--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('--seed', default=42, type=int)
+def main(data_dir, output_dir, model_name, epochs, no_crf, with_isr, 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,
+        use_crf=not no_crf,
+        with_isr=with_isr,
+        seed=seed,
+    )
+    run_training(config)
+
+
+if __name__ == '__main__':
+    main()

From 7de2b5860ef43835538d959f0f90db8bee4c57f8 Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Fri, 3 Jul 2026 16:50:51 +0530
Subject: [PATCH 03/12] Add DeBERTa tagger with CRF head

---
 etc/scripts/dataset_pipeline/phrase_model.py | 177 +++++++++++++++++++
 1 file changed, 177 insertions(+)
 create mode 100644 etc/scripts/dataset_pipeline/phrase_model.py

diff --git a/etc/scripts/dataset_pipeline/phrase_model.py b/etc/scripts/dataset_pipeline/phrase_model.py
new file mode 100644
index 0000000000..e4152d58a1
--- /dev/null
+++ b/etc/scripts/dataset_pipeline/phrase_model.py
@@ -0,0 +1,177 @@
+# 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)
+
+        self.backbone = AutoModel.from_pretrained(config.model_name)
+        self.backbone.gradient_checkpointing_enable()
+
+        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.base_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})
+
+    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'])

From 03c14c8dc43f35201de3b97adcba8f1491d95127 Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Fri, 3 Jul 2026 16:50:51 +0530
Subject: [PATCH 04/12] Add ONNX export with numpy viterbi decoding

---
 etc/scripts/dataset_pipeline/export_onnx.py | 189 ++++++++++++++++++++
 1 file changed, 189 insertions(+)
 create mode 100644 etc/scripts/dataset_pipeline/export_onnx.py

diff --git a/etc/scripts/dataset_pipeline/export_onnx.py b/etc/scripts/dataset_pipeline/export_onnx.py
new file mode 100644
index 0000000000..f9780baa28
--- /dev/null
+++ b/etc/scripts/dataset_pipeline/export_onnx.py
@@ -0,0 +1,189 @@
+# 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')
+    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()

From 95f1264e0605d4d9c2755c1daa3992190703cce0 Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Fri, 3 Jul 2026 16:50:52 +0530
Subject: [PATCH 05/12] Add tests for the training helpers

---
 .../dataset_pipeline/test_train_model.py      | 136 ++++++++++++++++++
 1 file changed, 136 insertions(+)
 create mode 100644 etc/scripts/dataset_pipeline/test_train_model.py

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..87abe6e1c1
--- /dev/null
+++ b/etc/scripts/dataset_pipeline/test_train_model.py
@@ -0,0 +1,136 @@
+# 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
+
+import pytest
+
+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():
+    pytest.importorskip('seqeval')
+    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():
+    pytest.importorskip('seqeval')
+    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

From fd2a7e010f4c91fde569f2887481a2b04dec0833 Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Sat, 4 Jul 2026 22:36:15 +0530
Subject: [PATCH 06/12] Require a fast tokenizer since word_ids needs it

---
 etc/scripts/dataset_pipeline/train_model.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/etc/scripts/dataset_pipeline/train_model.py b/etc/scripts/dataset_pipeline/train_model.py
index d21f0c8bb1..6c78704002 100644
--- a/etc/scripts/dataset_pipeline/train_model.py
+++ b/etc/scripts/dataset_pipeline/train_model.py
@@ -257,7 +257,9 @@ def run_training(config):
     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)
+    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)
     val_ds = PhraseDataset(config.data_dir / 'val.jsonl', tokenizer, config.max_length)
     test_ds = PhraseDataset(config.data_dir / 'test.jsonl', tokenizer, config.max_length)

From a89dd67bcba54510bca75c19962783debf9cabc1 Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Sun, 5 Jul 2026 00:32:47 +0530
Subject: [PATCH 07/12] Add a limit flag for quick smoke tests

---
 etc/scripts/dataset_pipeline/train_model.py | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/etc/scripts/dataset_pipeline/train_model.py b/etc/scripts/dataset_pipeline/train_model.py
index 6c78704002..3761833025 100644
--- a/etc/scripts/dataset_pipeline/train_model.py
+++ b/etc/scripts/dataset_pipeline/train_model.py
@@ -41,6 +41,9 @@ class Config:
     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
+
     # turn off to train a plain softmax tagger for ablation
     use_crf: bool = True
     # extra weighted cross entropy added to the CRF loss, 0 turns it off
@@ -105,10 +108,12 @@ def align_labels(tokens, word_labels, tokenizer, max_length):
 class PhraseDataset:
     """Reads a BIOES JSONL split and encodes each rule for the model"""
 
-    def __init__(self, path, tokenizer, max_length):
+    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:
@@ -260,9 +265,9 @@ def run_training(config):
     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)
-    val_ds = PhraseDataset(config.data_dir / 'val.jsonl', tokenizer, config.max_length)
-    test_ds = PhraseDataset(config.data_dir / 'test.jsonl', tokenizer, config.max_length)
+    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:
@@ -346,8 +351,10 @@ def run_training(config):
               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('--seed', default=42, type=int)
-def main(data_dir, output_dir, model_name, epochs, no_crf, with_isr, seed):
+def main(data_dir, output_dir, model_name, epochs, no_crf, with_isr, limit, seed):
     """Train the required phrase tagger from a BIOES dataset"""
     config = Config(
         data_dir=Path(data_dir),
@@ -356,6 +363,7 @@ def main(data_dir, output_dir, model_name, epochs, no_crf, with_isr, seed):
         epochs=epochs,
         use_crf=not no_crf,
         with_isr=with_isr,
+        limit=limit,
         seed=seed,
     )
     run_training(config)

From db9fe3450d70a91b9725f6d799efc348e401c568 Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Mon, 6 Jul 2026 00:00:09 +0530
Subject: [PATCH 08/12] Score spans without seqeval and add lr, aux loss and
 resume flags

---
 etc/scripts/dataset_pipeline/train_model.py | 67 +++++++++++++++------
 1 file changed, 47 insertions(+), 20 deletions(-)

diff --git a/etc/scripts/dataset_pipeline/train_model.py b/etc/scripts/dataset_pipeline/train_model.py
index 3761833025..aa973b7951 100644
--- a/etc/scripts/dataset_pipeline/train_model.py
+++ b/etc/scripts/dataset_pipeline/train_model.py
@@ -1,6 +1,7 @@
 # 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
@@ -34,6 +35,8 @@ class Config:
     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
@@ -43,11 +46,14 @@ class Config:
 
     # 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
-    # extra weighted cross entropy added to the CRF loss, 0 turns it off
-    aux_ce_weight: float = 0.0
+    # 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
 
@@ -180,28 +186,31 @@ def decode_row(pred_row, label_row):
 
 
 def compute_metrics(eval_pred):
-    """Entity level F1 plus a strict rule level exact match"""
-    from seqeval.metrics import f1_score
-    from seqeval.metrics import precision_score
-    from seqeval.metrics import recall_score
-    from seqeval.scheme import IOBES
+    """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
-    pred_tags = []
-    true_tags = []
+    tp = fp = fn = 0
     exact = 0
     for pred_row, label_row in zip(predictions, labels):
         predicted, actual = decode_row(pred_row, label_row)
-        pred_tags.append(predicted)
-        true_tags.append(actual)
-        if extract_spans(predicted) == extract_spans(actual):
+        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
 
-    scored = dict(mode='strict', scheme=IOBES)
+    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_score(true_tags, pred_tags, **scored),
-        'precision': precision_score(true_tags, pred_tags, **scored),
-        'recall': recall_score(true_tags, pred_tags, **scored),
+        'f1': f1,
+        'precision': precision,
+        'recall': recall,
         'exact_match': exact / len(predictions) if len(predictions) else 0.0,
     }
 
@@ -304,19 +313,26 @@ def run_training(config):
         dataloader_num_workers=2,
     )
 
-    trainer = PhraseTrainer(
+    trainer_kwargs = dict(
         model=model,
         args=args,
         train_dataset=train_ds,
         eval_dataset=val_ds,
         data_collator=collator,
-        tokenizer=tokenizer,
         compute_metrics=compute_metrics,
         optimizers=(build_optimizer(config, model), None),
         callbacks=[EarlyStoppingCallback(early_stopping_patience=config.early_stopping_patience)],
     )
 
-    trainer.train()
+    # 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
@@ -347,23 +363,34 @@ def run_training(config):
               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, no_crf, with_isr, limit, seed):
+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)

From cc852f7f74f3cf719b5daef3031ce46ef9e449fc Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Mon, 6 Jul 2026 00:00:10 +0530
Subject: [PATCH 09/12] Fit deberta on a single gpu with fp32 and the 8 bit
 optimizer

---
 etc/scripts/dataset_pipeline/phrase_model.py | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/etc/scripts/dataset_pipeline/phrase_model.py b/etc/scripts/dataset_pipeline/phrase_model.py
index e4152d58a1..f1e0c12014 100644
--- a/etc/scripts/dataset_pipeline/phrase_model.py
+++ b/etc/scripts/dataset_pipeline/phrase_model.py
@@ -25,8 +25,13 @@ def __init__(self, config):
         self.aux_ce_weight = config.aux_ce_weight
         self.num_labels = len(LABELS)
 
-        self.backbone = AutoModel.from_pretrained(config.model_name)
+        # 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)
@@ -140,7 +145,7 @@ def build_optimizer(config, model):
 
     def rate_for(name):
         if name.startswith('classifier') or name.startswith('crf'):
-            return config.base_lr
+            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))
@@ -154,7 +159,13 @@ def rate_for(name):
         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})
 
-    return AdamW(groups, lr=config.base_lr, eps=config.adam_epsilon, betas=(0.9, 0.999))
+    # 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):

From dda4768c7571b018eb1cce4acde4e9407f217f6b Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Mon, 6 Jul 2026 00:00:10 +0530
Subject: [PATCH 10/12] Skip the training only class weights when loading for
 export

---
 etc/scripts/dataset_pipeline/export_onnx.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/etc/scripts/dataset_pipeline/export_onnx.py b/etc/scripts/dataset_pipeline/export_onnx.py
index f9780baa28..12cd54a887 100644
--- a/etc/scripts/dataset_pipeline/export_onnx.py
+++ b/etc/scripts/dataset_pipeline/export_onnx.py
@@ -83,6 +83,8 @@ def load_tagger(model_dir, train_config):
         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()
 

From 3e45fdb81c1856089579855f9798780e1824a118 Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Mon, 6 Jul 2026 00:00:10 +0530
Subject: [PATCH 11/12] Drop the seqeval dependency

---
 etc/requirements-ml.txt                          | 3 ---
 etc/scripts/dataset_pipeline/test_train_model.py | 4 ----
 2 files changed, 7 deletions(-)

diff --git a/etc/requirements-ml.txt b/etc/requirements-ml.txt
index 4d1e3e9461..fe8222aa5b 100644
--- a/etc/requirements-ml.txt
+++ b/etc/requirements-ml.txt
@@ -13,9 +13,6 @@ accelerate==0.33.0
 # CRF layer on top of the token classifier
 pytorch-crf==0.7.2
 
-# span level scoring
-seqeval==1.2.2
-
 # ONNX export for CPU inference
 onnx==1.16.2
 onnxruntime==1.18.1
diff --git a/etc/scripts/dataset_pipeline/test_train_model.py b/etc/scripts/dataset_pipeline/test_train_model.py
index 87abe6e1c1..1b0919de9c 100644
--- a/etc/scripts/dataset_pipeline/test_train_model.py
+++ b/etc/scripts/dataset_pipeline/test_train_model.py
@@ -3,8 +3,6 @@
 import sys
 from pathlib import Path
 
-import pytest
-
 sys.path.insert(0, str(Path(__file__).parent))
 from train_model import align_labels
 from train_model import decode_row
@@ -113,7 +111,6 @@ def test_viterbi_avoids_forbidden_transition():
 
 
 def test_compute_metrics_perfect_and_exact():
-    pytest.importorskip('seqeval')
     from train_model import compute_metrics
 
     # two rules, word level predictions and labels line up perfectly
@@ -125,7 +122,6 @@ def test_compute_metrics_perfect_and_exact():
 
 
 def test_compute_metrics_counts_a_miss():
-    pytest.importorskip('seqeval')
     from train_model import compute_metrics
 
     # predict O everywhere, the gold has one phrase, so f1 is 0

From 35763b1daace73b5ca00eb849fca50116cac4bb5 Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Mon, 6 Jul 2026 00:11:13 +0530
Subject: [PATCH 12/12] Relax ml pins to what was validated and add
 bitsandbytes

---
 etc/requirements-ml.txt | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/etc/requirements-ml.txt b/etc/requirements-ml.txt
index fe8222aa5b..af9d7c8945 100644
--- a/etc/requirements-ml.txt
+++ b/etc/requirements-ml.txt
@@ -5,14 +5,17 @@
 # install with: pip install -r etc/requirements-ml.txt
 
 torch>=2.0
-transformers==4.44.2
-sentencepiece==0.2.0
+transformers>=4.44
+sentencepiece>=0.2
 protobuf>=3.20
-accelerate==0.33.0
+accelerate>=0.33
 
 # CRF layer on top of the token classifier
-pytorch-crf==0.7.2
+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.2
-onnxruntime==1.18.1
+onnx>=1.16
+onnxruntime>=1.18