From 5f79d5715168d018caf56246c104f42d5b06f50c Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Sat, 1 Aug 2026 00:08:32 +0530
Subject: [PATCH 01/11] pick the rules to tag and mark the predicted phrases

loads the rules once instead of calling get_updatable_rules_by_expression per
expression, that reparses every rule file each time
---
 .../dataset_pipeline/add_ml_phrases.py        | 132 ++++++++++++++++++
 1 file changed, 132 insertions(+)
 create mode 100644 etc/scripts/dataset_pipeline/add_ml_phrases.py

diff --git a/etc/scripts/dataset_pipeline/add_ml_phrases.py b/etc/scripts/dataset_pipeline/add_ml_phrases.py
new file mode 100644
index 0000000000..e70a8a5144
--- /dev/null
+++ b/etc/scripts/dataset_pipeline/add_ml_phrases.py
@@ -0,0 +1,132 @@
+# runs the trained phrase tagger over license rules and marks the required
+# phrases it predicts with {{ }}
+# the rules are written in place, so regen the index afterwards
+import sys
+import unicodedata
+from pathlib import Path
+
+import click
+
+sys.path.insert(0, str(Path(__file__).parent))
+
+from licensedcode.models import get_rules_by_expression
+from licensedcode.required_phrases import add_required_phrase_to_rule
+from licensedcode.required_phrases import find_phrase_spans_in_text
+from licensedcode.required_phrases import RequiredPhraseRuleCandidate
+from licensedcode.tokenize import get_existing_required_phrase_spans
+from licensedcode.tokenize import required_phrase_splitter
+
+from train_model import extract_spans
+
+# the minima gen-new-required-phrases-rules uses to call a phrase good enough
+MIN_TOKENS = 2
+MIN_SINGLE_TOKEN_LEN = 5
+
+# scancode leaves rule texts longer than this alone
+MAX_RULE_TEXT = 4000
+
+
+def words_from_text(text):
+    """Words for the model, tokenized the way build_dataset.py does it"""
+    text = text.replace('\r\n', '\n').replace('\r', '\n')
+    text = unicodedata.normalize('NFKC', text)
+    return required_phrase_splitter(text)
+
+
+def is_updatable(rule):
+    """True if this rule can take new required phrases
+
+    Same checks as get_updatable_rules_by_expression, and also skip rules that
+    already have phrases: their {{ }} are tokens too, so they would shift every
+    word after them and the tags would not line up with the text
+    """
+    if rule.is_from_license:
+        return False
+
+    if len(rule.text) > MAX_RULE_TEXT:
+        return False
+
+    # covers required phrase rules, false positives, tiny texts and more
+    if not rule.is_approx_matchable:
+        return False
+
+    if rule.skip_for_required_phrase_generation:
+        return False
+
+    return not get_existing_required_phrase_spans(rule.text)
+
+
+def select_rules(license_expression=None):
+    """Rules that can take new required phrases, by license expression
+
+    get_updatable_rules_by_expression reloads every rule file each time it is
+    called and skips everything when passed None, so load the rules once here
+    and do the filtering in memory
+    """
+    rules_by_expression = get_rules_by_expression()
+
+    if license_expression:
+        rules = rules_by_expression.get(license_expression)
+        if not rules:
+            raise click.ClickException(f'no rules for license expression: {license_expression}')
+        rules_by_expression = {license_expression: rules}
+
+    selected = {}
+    for expression, rules in rules_by_expression.items():
+        updatable = [rule for rule in rules if is_updatable(rule)]
+        if updatable:
+            selected[expression] = updatable
+
+    return selected
+
+
+def phrases_from_tags(tags, words, truncated=False):
+    """Predicted phrase texts for one rule, longest first
+
+    On a truncated rule we drop a span that runs to the last tag: extract_spans
+    closes whatever is still open at the end, so the phrase would be cut short
+    """
+    phrases = set()
+    for start, end in extract_spans(tags):
+        if end >= len(words):
+            continue
+        if truncated and end == len(tags) - 1:
+            continue
+        phrases.add(' '.join(words[start:end + 1]))
+
+    # longest first, same order required phrases are applied in elsewhere
+    return sorted(phrases, key=lambda phrase: (-len(phrase), phrase))
+
+
+def inject(rule, phrases, counts, dry_run=False, verbose=False):
+    """Mark the good phrases in one rule, True if the rule was written"""
+    # read the source before the loop, add_required_phrase_to_rule overwrites it
+    source = f'{rule.source} ml_model' if rule.source else 'ml_model'
+    written = False
+
+    for phrase in phrases:
+        candidate = RequiredPhraseRuleCandidate.create(rule.license_expression, phrase)
+        if not candidate.is_good(rule, MIN_TOKENS, MIN_SINGLE_TOKEN_LEN):
+            counts['rejected'] += 1
+            continue
+
+        # the words come from NFKC normalized text but we mark up the raw text,
+        # so check the phrase can still be found there
+        if not find_phrase_spans_in_text(rule.text, phrase):
+            counts['not_found'] += 1
+            continue
+
+        updated = add_required_phrase_to_rule(
+            rule=rule,
+            required_phrase=phrase,
+            source=source,
+            debug=verbose,
+            dry_run=dry_run,
+        )
+        if updated:
+            counts['injected'] += 1
+            written = True
+        else:
+            counts['skipped'] += 1
+
+    return written

From 43f1651e25256e866c4ecd8505bcd123e2b3d79b Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Sat, 1 Aug 2026 00:22:18 +0530
Subject: [PATCH 02/11] load the tagger and wire up the cli

reuses validate_and_reindex for the checks after a write, and still prints the
reindex reminder when it is not asked to do it
---
 .../dataset_pipeline/add_ml_phrases.py        | 197 ++++++++++++++++++
 1 file changed, 197 insertions(+)

diff --git a/etc/scripts/dataset_pipeline/add_ml_phrases.py b/etc/scripts/dataset_pipeline/add_ml_phrases.py
index e70a8a5144..80e4f3a893 100644
--- a/etc/scripts/dataset_pipeline/add_ml_phrases.py
+++ b/etc/scripts/dataset_pipeline/add_ml_phrases.py
@@ -1,22 +1,32 @@
 # runs the trained phrase tagger over license rules and marks the required
 # phrases it predicts with {{ }}
 # the rules are written in place, so regen the index afterwards
+import json
+import os
 import sys
 import unicodedata
 from pathlib import Path
 
 import click
 
+# transformers pulls in keras otherwise and blows up on keras 3
+os.environ.setdefault('USE_TF', '0')
+
 sys.path.insert(0, str(Path(__file__).parent))
 
 from licensedcode.models import get_rules_by_expression
 from licensedcode.required_phrases import add_required_phrase_to_rule
 from licensedcode.required_phrases import find_phrase_spans_in_text
 from licensedcode.required_phrases import RequiredPhraseRuleCandidate
+from licensedcode.required_phrases import validate_and_reindex
 from licensedcode.tokenize import get_existing_required_phrase_spans
 from licensedcode.tokenize import required_phrase_splitter
 
 from train_model import extract_spans
+from train_model import ID2LABEL
+from train_model import LABELS
+from train_model import MAX_LENGTH
+from train_model import MODEL_NAME
 
 # the minima gen-new-required-phrases-rules uses to call a phrase good enough
 MIN_TOKENS = 2
@@ -26,6 +36,95 @@
 MAX_RULE_TEXT = 4000
 
 
+class InferenceConfig:
+    """The few settings PhraseTagger reads when it is built
+
+    aux_ce_weight stays 0 so no class weight buffer is created, we are not
+    computing a loss here
+    """
+    aux_ce_weight = 0
+
+    def __init__(self, model_name, use_crf, max_length):
+        self.model_name = model_name
+        self.use_crf = use_crf
+        self.max_length = max_length
+
+
+def load_model(model, hf_token=None):
+    """Tokenizer and tagger with the trained weights, ready to predict
+
+    ``model`` is a local directory or a huggingface repo id
+    """
+    from safetensors.torch import load_file
+    from transformers import AutoTokenizer
+
+    from phrase_model import PhraseTagger
+
+    model_dir = Path(model)
+    if not model_dir.is_dir():
+        from huggingface_hub import snapshot_download
+        model_dir = Path(snapshot_download(repo_id=model, token=hf_token))
+
+    config_file = model_dir / 'train_config.json'
+    if config_file.exists():
+        saved = json.loads(config_file.read_text())
+    else:
+        # older runs did not save it, fall back to how we always trained
+        click.echo(f'no train_config.json in {model_dir}, using the training defaults')
+        saved = {}
+
+    labels = saved.get('labels', LABELS)
+    if len(labels) != len(LABELS):
+        raise click.ClickException(
+            f'this checkpoint has {len(labels)} labels, expected the BIOES {len(LABELS)}'
+        )
+
+    config = InferenceConfig(
+        model_name=saved.get('model_name', MODEL_NAME),
+        use_crf=saved.get('use_crf', True),
+        max_length=saved.get('max_length', MAX_LENGTH),
+    )
+    if not config.use_crf:
+        raise click.ClickException('this tool expects the CRF model')
+
+    tokenizer = AutoTokenizer.from_pretrained(str(model_dir), use_fast=True)
+    if not tokenizer.is_fast:
+        raise click.ClickException('need a fast tokenizer for word_ids, got a slow one')
+
+    tagger = PhraseTagger(config)
+    tagger.load_state_dict(load_file(str(model_dir / 'model.safetensors')), strict=True)
+    # only needed while training and it warns under no_grad
+    tagger.backbone.gradient_checkpointing_disable()
+    tagger.eval()
+
+    return tagger, tokenizer, config.max_length
+
+
+def predict_phrases(tagger, tokenizer, max_length, words):
+    """Phrases the tagger predicts for one rule"""
+    import torch
+
+    encoding = tokenizer(
+        words,
+        is_split_into_words=True,
+        truncation=True,
+        max_length=max_length,
+        return_tensors='pt',
+    )
+    word_ids = encoding.word_ids()
+
+    with torch.no_grad():
+        predicted = tagger.predict_words(
+            encoding['input_ids'],
+            encoding['attention_mask'],
+            word_ids,
+        )
+
+    tags = [ID2LABEL.get(int(label), 'O') for label in predicted]
+    truncated = len(tags) < len(words)
+    return phrases_from_tags(tags, words, truncated=truncated), truncated
+
+
 def words_from_text(text):
     """Words for the model, tokenized the way build_dataset.py does it"""
     text = text.replace('\r\n', '\n').replace('\r', '\n')
@@ -130,3 +229,101 @@ def inject(rule, phrases, counts, dry_run=False, verbose=False):
             counts['skipped'] += 1
 
     return written
+
+
+def process_rules(
+    tagger,
+    tokenizer,
+    max_length,
+    license_expression=None,
+    dry_run=False,
+    limit=0,
+    verbose=False,
+):
+    """Predict and mark phrases in every eligible rule, return the counts"""
+    counts = dict(rules=0, truncated=0, rejected=0, not_found=0, injected=0, skipped=0, written=0)
+
+    selected = select_rules(license_expression=license_expression)
+    total = sum(len(rules) for rules in selected.values())
+    click.echo(f'tagging {total} rules in {len(selected)} license expressions')
+
+    for expression, rules in selected.items():
+        if verbose:
+            click.echo(f'{expression}: {len(rules)} rules')
+
+        for rule in rules:
+            if limit and counts['rules'] >= limit:
+                click.echo(f'stopping at {limit} rules')
+                return counts
+
+            counts['rules'] += 1
+            words = words_from_text(rule.text)
+            if not words:
+                continue
+
+            phrases, truncated = predict_phrases(tagger, tokenizer, max_length, words)
+            if truncated:
+                counts['truncated'] += 1
+
+            if not phrases:
+                continue
+
+            if verbose:
+                click.echo(f'  {rule.identifier}: {phrases}')
+
+            if inject(rule, phrases, counts, dry_run=dry_run, verbose=verbose):
+                counts['written'] += 1
+
+    return counts
+
+
+@click.command()
+@click.option('--model', required=True,
+              help='Trained model directory, or a huggingface repo id to download')
+@click.option('--license-expression', default=None,
+              help='Only tag rules for this license expression, example: apache-2.0')
+@click.option('--dry-run', is_flag=True, default=False,
+              help='Predict and check phrases but do not save any rule')
+@click.option('--limit', default=0, type=int,
+              help='Stop after this many rules, 0 does all of them')
+@click.option('--validate', is_flag=True, default=False,
+              help='Validate all rules and licenses at the end')
+@click.option('--reindex', is_flag=True, default=False,
+              help='Rebuild and cache the license index at the end')
+@click.option('-v', '--verbose', is_flag=True, default=False,
+              help='Print the phrases predicted for each rule')
+@click.help_option('-h', '--help')
+def main(model, license_expression, dry_run, limit, validate, reindex, verbose):
+    """Add required phrases to license rules using the trained phrase tagger"""
+    tagger, tokenizer, max_length = load_model(model, hf_token=os.environ.get('HF_TOKEN'))
+
+    counts = process_rules(
+        tagger=tagger,
+        tokenizer=tokenizer,
+        max_length=max_length,
+        license_expression=license_expression,
+        dry_run=dry_run,
+        limit=limit,
+        verbose=verbose,
+    )
+
+    click.echo('')
+    click.echo(f"rules processed  : {counts['rules']}")
+    click.echo(f"  truncated      : {counts['truncated']}")
+    click.echo(f"phrases injected : {counts['injected']}")
+    click.echo(f"  rejected       : {counts['rejected']}")
+    click.echo(f"  not found      : {counts['not_found']}")
+    click.echo(f"  nothing to add : {counts['skipped']}")
+    click.echo(f"rules written    : {counts['written']}")
+
+    if dry_run:
+        click.echo('dry run, no rules were saved')
+    elif counts['written']:
+        if validate or reindex:
+            validate_and_reindex(validate=validate, reindex=reindex, verbose=verbose)
+        if not reindex:
+            click.echo('run scancode-reindex-licenses to pick up the new required phrases')
+
+
+if __name__ == '__main__':
+    main()

From 7049d052c3a5b8db608b2a9d330dcfaadb5524cb Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Sat, 1 Aug 2026 00:35:37 +0530
Subject: [PATCH 03/11] tests for the phrase injection script

the tagger and tokenizer are faked so this needs no weights and no network
---
 .../dataset_pipeline/test_add_ml_phrases.py   | 225 ++++++++++++++++++
 1 file changed, 225 insertions(+)
 create mode 100644 etc/scripts/dataset_pipeline/test_add_ml_phrases.py

diff --git a/etc/scripts/dataset_pipeline/test_add_ml_phrases.py b/etc/scripts/dataset_pipeline/test_add_ml_phrases.py
new file mode 100644
index 0000000000..e96628916b
--- /dev/null
+++ b/etc/scripts/dataset_pipeline/test_add_ml_phrases.py
@@ -0,0 +1,225 @@
+# tests for add_ml_phrases.py
+# no model and no network here, the tagger is stubbed out
+import sys
+from pathlib import Path
+
+import pytest
+
+sys.path.insert(0, str(Path(__file__).parent))
+import add_ml_phrases
+from add_ml_phrases import inject
+from add_ml_phrases import is_updatable
+from add_ml_phrases import phrases_from_tags
+from add_ml_phrases import predict_phrases
+from add_ml_phrases import process_rules
+from add_ml_phrases import select_rules
+from add_ml_phrases import words_from_text
+from train_model import LABEL2ID
+
+from licensedcode.models import Rule
+
+
+class FakeRule:
+    """Just the flags is_updatable looks at"""
+
+    def __init__(self, text='some license text here', is_from_license=False,
+                 is_approx_matchable=True, skip=False):
+        self.text = text
+        self.is_from_license = is_from_license
+        self.is_approx_matchable = is_approx_matchable
+        self.skip_for_required_phrase_generation = skip
+
+
+class FakeEncoding(dict):
+
+    def __init__(self, word_ids):
+        super().__init__()
+        self._word_ids = word_ids
+        self['input_ids'] = [[0] * len(word_ids)]
+        self['attention_mask'] = [[1] * len(word_ids)]
+
+    def word_ids(self):
+        return self._word_ids
+
+
+class FakeTokenizer:
+    """Maps every word to one subword, truncating at max_length"""
+
+    def __call__(self, words, max_length=512, **kwargs):
+        return FakeEncoding(list(range(len(words)))[:max_length])
+
+
+class StubTagger:
+    """Tags the first three words as one phrase"""
+
+    def predict_words(self, input_ids, attention_mask, word_ids):
+        count = len([w for w in word_ids if w is not None])
+        tags = [LABEL2ID['B-REQ'], LABEL2ID['I-REQ'], LABEL2ID['E-REQ']]
+        return (tags + [LABEL2ID['O']] * count)[:count]
+
+
+class TestWordsFromText:
+
+    def test_matches_the_dataset_tokenizer(self):
+        assert words_from_text('Apache-2.0 License') == ['Apache', '2', '0', 'License']
+
+    def test_normalizes_line_endings(self):
+        assert words_from_text('one\r\ntwo\rthree') == ['one', 'two', 'three']
+
+    def test_applies_nfkc(self):
+        # the fi ligature becomes two characters, otherwise it stays one token
+        assert words_from_text('a\ufb01x') == ['afix']
+
+    def test_empty(self):
+        assert words_from_text('') == []
+
+
+class TestPhrasesFromTags:
+
+    def test_one_phrase(self):
+        words = ['Apache', 'License', 'Version', 'x']
+        tags = ['B-REQ', 'I-REQ', 'E-REQ', 'O']
+        assert phrases_from_tags(tags, words) == ['Apache License Version']
+
+    def test_longest_first_and_deduped(self):
+        words = ['mit', 'license', 'mit', 'license']
+        tags = ['B-REQ', 'E-REQ', 'S-REQ', 'O']
+        assert phrases_from_tags(tags, words) == ['mit license', 'mit']
+
+    def test_drops_a_span_cut_by_truncation(self):
+        words = ['gnu', 'general', 'public']
+        tags = ['O', 'B-REQ', 'I-REQ']
+        assert phrases_from_tags(tags, words, truncated=True) == []
+        assert phrases_from_tags(tags, words, truncated=False) == ['general public']
+
+    def test_ignores_a_span_past_the_words(self):
+        assert phrases_from_tags(['B-REQ', 'E-REQ'], ['only']) == []
+
+    def test_nothing_tagged(self):
+        assert phrases_from_tags(['O', 'O'], ['a', 'b']) == []
+
+
+class TestIsUpdatable:
+
+    def test_plain_rule(self):
+        assert is_updatable(FakeRule())
+
+    def test_skips_rule_from_a_license(self):
+        assert not is_updatable(FakeRule(is_from_license=True))
+
+    def test_skips_long_text(self):
+        assert not is_updatable(FakeRule(text='x' * 4001))
+
+    def test_skips_not_approx_matchable(self):
+        assert not is_updatable(FakeRule(is_approx_matchable=False))
+
+    def test_skips_when_asked_to(self):
+        assert not is_updatable(FakeRule(skip=True))
+
+    def test_skips_rules_that_already_have_phrases(self):
+        assert not is_updatable(FakeRule(text='under the {{mit license}} terms'))
+
+
+def make_rule(text, source=None):
+    rule = Rule(
+        license_expression='mit',
+        identifier='mit_test.RULE',
+        text=text,
+        is_license_reference=True,
+        relevance=100,
+    )
+    rule.source = source
+    return rule
+
+
+def new_counts():
+    return dict(rules=0, truncated=0, rejected=0, not_found=0, injected=0, skipped=0, written=0)
+
+
+TEXT = 'Permission is granted under the MIT License to do things with this'
+
+
+class TestInject:
+
+    def test_marks_a_phrase_and_sets_the_source(self):
+        rule = make_rule(TEXT)
+        counts = new_counts()
+        assert inject(rule, ['MIT License'], counts, dry_run=True)
+        assert counts['injected'] == 1
+        assert '{{MIT License}}' in rule.text
+        assert rule.source == 'ml_model'
+
+    def test_keeps_an_existing_source(self):
+        rule = make_rule(TEXT, source='mit_1.RULE')
+        inject(rule, ['MIT License'], new_counts(), dry_run=True)
+        assert rule.source == 'mit_1.RULE ml_model'
+
+    def test_rejects_a_phrase_is_good_does_not_like(self):
+        rule = make_rule(TEXT)
+        counts = new_counts()
+        assert not inject(rule, ['is'], counts, dry_run=True)
+        assert counts['rejected'] == 1
+        assert '{{' not in rule.text
+
+    def test_counts_a_phrase_that_is_not_in_the_text(self):
+        rule = make_rule(TEXT)
+        counts = new_counts()
+        assert not inject(rule, ['Apache License'], counts, dry_run=True)
+        assert counts['not_found'] == 1
+
+    def test_does_not_mark_the_same_phrase_twice(self):
+        rule = make_rule(TEXT)
+        counts = new_counts()
+        inject(rule, ['MIT License'], counts, dry_run=True)
+        inject(rule, ['MIT License'], counts, dry_run=True)
+        assert counts['injected'] == 1
+        assert counts['skipped'] == 1
+        assert rule.text.count('{{') == 1
+
+
+class TestPredictPhrases:
+
+    def test_predicts_and_reports_no_truncation(self):
+        words = words_from_text(TEXT)
+        phrases, truncated = predict_phrases(StubTagger(), FakeTokenizer(), 512, words)
+        assert phrases == [' '.join(words[:3])]
+        assert not truncated
+
+    def test_reports_truncation(self):
+        words = ['word'] * 20
+        phrases, truncated = predict_phrases(StubTagger(), FakeTokenizer(), 5, words)
+        assert truncated
+
+
+class TestSelectRules:
+
+    def test_filters_and_groups(self, monkeypatch):
+        rules = {
+            'mit': [FakeRule(), FakeRule(is_from_license=True)],
+            'bsd-new': [FakeRule(skip=True)],
+        }
+        monkeypatch.setattr(add_ml_phrases, 'get_rules_by_expression', lambda: rules)
+        selected = select_rules()
+        assert list(selected) == ['mit']
+        assert len(selected['mit']) == 1
+
+    def test_unknown_expression(self, monkeypatch):
+        monkeypatch.setattr(add_ml_phrases, 'get_rules_by_expression', lambda: {'mit': [FakeRule()]})
+        with pytest.raises(Exception):
+            select_rules('nope-1.0')
+
+
+class TestProcessRules:
+
+    def test_dry_run_marks_nothing_on_disk(self, monkeypatch):
+        rule = make_rule(TEXT)
+        monkeypatch.setattr(add_ml_phrases, 'select_rules', lambda license_expression=None: {'mit': [rule]})
+        counts = process_rules(StubTagger(), FakeTokenizer(), 512, dry_run=True)
+        assert counts['rules'] == 1
+        assert counts['injected'] == 1
+
+    def test_limit_stops_early(self, monkeypatch):
+        rules = [make_rule(TEXT) for _ in range(5)]
+        monkeypatch.setattr(add_ml_phrases, 'select_rules', lambda license_expression=None: {'mit': rules})
+        counts = process_rules(StubTagger(), FakeTokenizer(), 512, dry_run=True, limit=2)
+        assert counts['rules'] == 2

From 0a2dca5d60c0bac8925aa153b98fd0f7d6e64c86 Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Sat, 1 Aug 2026 00:37:42 +0530
Subject: [PATCH 04/11] keep predict_phrases next to the decoding helpers

---
 .../dataset_pipeline/add_ml_phrases.py        | 51 ++++++++++---------
 1 file changed, 26 insertions(+), 25 deletions(-)

diff --git a/etc/scripts/dataset_pipeline/add_ml_phrases.py b/etc/scripts/dataset_pipeline/add_ml_phrases.py
index 80e4f3a893..ff21c05c99 100644
--- a/etc/scripts/dataset_pipeline/add_ml_phrases.py
+++ b/etc/scripts/dataset_pipeline/add_ml_phrases.py
@@ -100,31 +100,6 @@ def load_model(model, hf_token=None):
     return tagger, tokenizer, config.max_length
 
 
-def predict_phrases(tagger, tokenizer, max_length, words):
-    """Phrases the tagger predicts for one rule"""
-    import torch
-
-    encoding = tokenizer(
-        words,
-        is_split_into_words=True,
-        truncation=True,
-        max_length=max_length,
-        return_tensors='pt',
-    )
-    word_ids = encoding.word_ids()
-
-    with torch.no_grad():
-        predicted = tagger.predict_words(
-            encoding['input_ids'],
-            encoding['attention_mask'],
-            word_ids,
-        )
-
-    tags = [ID2LABEL.get(int(label), 'O') for label in predicted]
-    truncated = len(tags) < len(words)
-    return phrases_from_tags(tags, words, truncated=truncated), truncated
-
-
 def words_from_text(text):
     """Words for the model, tokenized the way build_dataset.py does it"""
     text = text.replace('\r\n', '\n').replace('\r', '\n')
@@ -197,6 +172,32 @@ def phrases_from_tags(tags, words, truncated=False):
     return sorted(phrases, key=lambda phrase: (-len(phrase), phrase))
 
 
+def predict_phrases(tagger, tokenizer, max_length, words):
+    """Phrases the tagger predicts for one rule"""
+    import torch
+
+    encoding = tokenizer(
+        words,
+        is_split_into_words=True,
+        truncation=True,
+        max_length=max_length,
+        return_tensors='pt',
+    )
+    word_ids = encoding.word_ids()
+
+    with torch.no_grad():
+        predicted = tagger.predict_words(
+            encoding['input_ids'],
+            encoding['attention_mask'],
+            word_ids,
+        )
+
+    tags = [ID2LABEL.get(int(label), 'O') for label in predicted]
+    # fewer tags than words means the rule did not fit in max_length
+    truncated = len(tags) < len(words)
+    return phrases_from_tags(tags, words, truncated=truncated), truncated
+
+
 def inject(rule, phrases, counts, dry_run=False, verbose=False):
     """Mark the good phrases in one rule, True if the rule was written"""
     # read the source before the loop, add_required_phrase_to_rule overwrites it

From 6527fbb67d11d197cc8f5e99897662636c73155c Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Sat, 1 Aug 2026 01:32:23 +0530
Subject: [PATCH 05/11] check the label set instead of just its size

a checkpoint with the same number of labels in another order would map every
prediction to the wrong tag
---
 .../dataset_pipeline/add_ml_phrases.py        | 28 +++++++++++++------
 .../dataset_pipeline/test_add_ml_phrases.py   | 25 ++++++++++++-----
 2 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/etc/scripts/dataset_pipeline/add_ml_phrases.py b/etc/scripts/dataset_pipeline/add_ml_phrases.py
index ff21c05c99..cfe0755305 100644
--- a/etc/scripts/dataset_pipeline/add_ml_phrases.py
+++ b/etc/scripts/dataset_pipeline/add_ml_phrases.py
@@ -51,7 +51,7 @@ def __init__(self, model_name, use_crf, max_length):
 
 
 def load_model(model, hf_token=None):
-    """Tokenizer and tagger with the trained weights, ready to predict
+    """The trained tagger, its tokenizer and the max length it was trained with
 
     ``model`` is a local directory or a huggingface repo id
     """
@@ -73,11 +73,11 @@ def load_model(model, hf_token=None):
         click.echo(f'no train_config.json in {model_dir}, using the training defaults')
         saved = {}
 
+    # the label order is what ID2LABEL maps over, a different one would quietly
+    # turn every prediction into the wrong tag
     labels = saved.get('labels', LABELS)
-    if len(labels) != len(LABELS):
-        raise click.ClickException(
-            f'this checkpoint has {len(labels)} labels, expected the BIOES {len(LABELS)}'
-        )
+    if labels != LABELS:
+        raise click.ClickException(f'this checkpoint was trained on other labels: {labels}')
 
     config = InferenceConfig(
         model_name=saved.get('model_name', MODEL_NAME),
@@ -198,6 +198,19 @@ def predict_phrases(tagger, tokenizer, max_length, words):
     return phrases_from_tags(tags, words, truncated=truncated), truncated
 
 
+def new_counts():
+    """What a run tallies up as it goes"""
+    return dict(
+        rules=0,
+        truncated=0,
+        rejected=0,
+        not_found=0,
+        injected=0,
+        skipped=0,
+        written=0,
+    )
+
+
 def inject(rule, phrases, counts, dry_run=False, verbose=False):
     """Mark the good phrases in one rule, True if the rule was written"""
     # read the source before the loop, add_required_phrase_to_rule overwrites it
@@ -242,7 +255,7 @@ def process_rules(
     verbose=False,
 ):
     """Predict and mark phrases in every eligible rule, return the counts"""
-    counts = dict(rules=0, truncated=0, rejected=0, not_found=0, injected=0, skipped=0, written=0)
+    counts = new_counts()
 
     selected = select_rules(license_expression=license_expression)
     total = sum(len(rules) for rules in selected.values())
@@ -308,8 +321,7 @@ def main(model, license_expression, dry_run, limit, validate, reindex, verbose):
         verbose=verbose,
     )
 
-    click.echo('')
-    click.echo(f"rules processed  : {counts['rules']}")
+    click.echo(f"\nrules processed  : {counts['rules']}")
     click.echo(f"  truncated      : {counts['truncated']}")
     click.echo(f"phrases injected : {counts['injected']}")
     click.echo(f"  rejected       : {counts['rejected']}")
diff --git a/etc/scripts/dataset_pipeline/test_add_ml_phrases.py b/etc/scripts/dataset_pipeline/test_add_ml_phrases.py
index e96628916b..7e01ad2cb5 100644
--- a/etc/scripts/dataset_pipeline/test_add_ml_phrases.py
+++ b/etc/scripts/dataset_pipeline/test_add_ml_phrases.py
@@ -3,12 +3,14 @@
 import sys
 from pathlib import Path
 
+import click
 import pytest
 
 sys.path.insert(0, str(Path(__file__).parent))
 import add_ml_phrases
 from add_ml_phrases import inject
 from add_ml_phrases import is_updatable
+from add_ml_phrases import new_counts
 from add_ml_phrases import phrases_from_tags
 from add_ml_phrases import predict_phrases
 from add_ml_phrases import process_rules
@@ -132,8 +134,12 @@ def make_rule(text, source=None):
     return rule
 
 
-def new_counts():
-    return dict(rules=0, truncated=0, rejected=0, not_found=0, injected=0, skipped=0, written=0)
+def patch_selection(monkeypatch, rules):
+    """Skip the real rule loading, it reads every rule file on disk"""
+    monkeypatch.setattr(
+        add_ml_phrases, 'select_rules',
+        lambda license_expression=None: {'mit': rules},
+    )
 
 
 TEXT = 'Permission is granted under the MIT License to do things with this'
@@ -204,8 +210,9 @@ def test_filters_and_groups(self, monkeypatch):
         assert len(selected['mit']) == 1
 
     def test_unknown_expression(self, monkeypatch):
-        monkeypatch.setattr(add_ml_phrases, 'get_rules_by_expression', lambda: {'mit': [FakeRule()]})
-        with pytest.raises(Exception):
+        rules = {'mit': [FakeRule()]}
+        monkeypatch.setattr(add_ml_phrases, 'get_rules_by_expression', lambda: rules)
+        with pytest.raises(click.ClickException):
             select_rules('nope-1.0')
 
 
@@ -213,13 +220,17 @@ class TestProcessRules:
 
     def test_dry_run_marks_nothing_on_disk(self, monkeypatch):
         rule = make_rule(TEXT)
-        monkeypatch.setattr(add_ml_phrases, 'select_rules', lambda license_expression=None: {'mit': [rule]})
+        patch_selection(monkeypatch, [rule])
         counts = process_rules(StubTagger(), FakeTokenizer(), 512, dry_run=True)
         assert counts['rules'] == 1
         assert counts['injected'] == 1
 
     def test_limit_stops_early(self, monkeypatch):
-        rules = [make_rule(TEXT) for _ in range(5)]
-        monkeypatch.setattr(add_ml_phrases, 'select_rules', lambda license_expression=None: {'mit': rules})
+        patch_selection(monkeypatch, [make_rule(TEXT) for _ in range(5)])
         counts = process_rules(StubTagger(), FakeTokenizer(), 512, dry_run=True, limit=2)
         assert counts['rules'] == 2
+
+    def test_counts_a_truncated_rule(self, monkeypatch):
+        patch_selection(monkeypatch, [make_rule(TEXT)])
+        counts = process_rules(StubTagger(), FakeTokenizer(), 4, dry_run=True)
+        assert counts['truncated'] == 1

From c296f5e32a08be1bc6baa28ac2c26686917f4fa4 Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Sat, 1 Aug 2026 01:45:51 +0530
Subject: [PATCH 06/11] drop the validate and reindex flags

printing the reminder is enough, reindexing is the caller's call
---
 etc/scripts/dataset_pipeline/add_ml_phrases.py | 17 ++---------------
 1 file changed, 2 insertions(+), 15 deletions(-)

diff --git a/etc/scripts/dataset_pipeline/add_ml_phrases.py b/etc/scripts/dataset_pipeline/add_ml_phrases.py
index cfe0755305..b52e24bec3 100644
--- a/etc/scripts/dataset_pipeline/add_ml_phrases.py
+++ b/etc/scripts/dataset_pipeline/add_ml_phrases.py
@@ -18,7 +18,6 @@
 from licensedcode.required_phrases import add_required_phrase_to_rule
 from licensedcode.required_phrases import find_phrase_spans_in_text
 from licensedcode.required_phrases import RequiredPhraseRuleCandidate
-from licensedcode.required_phrases import validate_and_reindex
 from licensedcode.tokenize import get_existing_required_phrase_spans
 from licensedcode.tokenize import required_phrase_splitter
 
@@ -162,8 +161,6 @@ def phrases_from_tags(tags, words, truncated=False):
     """
     phrases = set()
     for start, end in extract_spans(tags):
-        if end >= len(words):
-            continue
         if truncated and end == len(tags) - 1:
             continue
         phrases.add(' '.join(words[start:end + 1]))
@@ -272,9 +269,6 @@ def process_rules(
 
             counts['rules'] += 1
             words = words_from_text(rule.text)
-            if not words:
-                continue
-
             phrases, truncated = predict_phrases(tagger, tokenizer, max_length, words)
             if truncated:
                 counts['truncated'] += 1
@@ -300,14 +294,10 @@ def process_rules(
               help='Predict and check phrases but do not save any rule')
 @click.option('--limit', default=0, type=int,
               help='Stop after this many rules, 0 does all of them')
-@click.option('--validate', is_flag=True, default=False,
-              help='Validate all rules and licenses at the end')
-@click.option('--reindex', is_flag=True, default=False,
-              help='Rebuild and cache the license index at the end')
 @click.option('-v', '--verbose', is_flag=True, default=False,
               help='Print the phrases predicted for each rule')
 @click.help_option('-h', '--help')
-def main(model, license_expression, dry_run, limit, validate, reindex, verbose):
+def main(model, license_expression, dry_run, limit, verbose):
     """Add required phrases to license rules using the trained phrase tagger"""
     tagger, tokenizer, max_length = load_model(model, hf_token=os.environ.get('HF_TOKEN'))
 
@@ -332,10 +322,7 @@ def main(model, license_expression, dry_run, limit, validate, reindex, verbose):
     if dry_run:
         click.echo('dry run, no rules were saved')
     elif counts['written']:
-        if validate or reindex:
-            validate_and_reindex(validate=validate, reindex=reindex, verbose=verbose)
-        if not reindex:
-            click.echo('run scancode-reindex-licenses to pick up the new required phrases')
+        click.echo('run scancode-reindex-licenses to pick up the new required phrases')
 
 
 if __name__ == '__main__':

From f9b00278b5e0c9628d343bb3ab9b73793e145992 Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Sat, 1 Aug 2026 01:45:52 +0530
Subject: [PATCH 07/11] bail out before the backbone when there is nothing to
 tag

---
 etc/scripts/dataset_pipeline/test_add_ml_phrases.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/etc/scripts/dataset_pipeline/test_add_ml_phrases.py b/etc/scripts/dataset_pipeline/test_add_ml_phrases.py
index 7e01ad2cb5..ae014e3e71 100644
--- a/etc/scripts/dataset_pipeline/test_add_ml_phrases.py
+++ b/etc/scripts/dataset_pipeline/test_add_ml_phrases.py
@@ -94,9 +94,6 @@ def test_drops_a_span_cut_by_truncation(self):
         assert phrases_from_tags(tags, words, truncated=True) == []
         assert phrases_from_tags(tags, words, truncated=False) == ['general public']
 
-    def test_ignores_a_span_past_the_words(self):
-        assert phrases_from_tags(['B-REQ', 'E-REQ'], ['only']) == []
-
     def test_nothing_tagged(self):
         assert phrases_from_tags(['O', 'O'], ['a', 'b']) == []
 
@@ -196,6 +193,10 @@ def test_reports_truncation(self):
         phrases, truncated = predict_phrases(StubTagger(), FakeTokenizer(), 5, words)
         assert truncated
 
+    def test_a_rule_with_no_words(self):
+        # nothing to tag, and the tagger never gets as far as the backbone
+        assert predict_phrases(StubTagger(), FakeTokenizer(), 512, []) == ([], False)
+
 
 class TestSelectRules:
 

From da1bdd0feb452f4b24399e1ecac2450d774025cb Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Sat, 8 Aug 2026 23:48:48 +0530
Subject: [PATCH 08/11] cover marking two phrases in one rule

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

diff --git a/etc/scripts/dataset_pipeline/test_add_ml_phrases.py b/etc/scripts/dataset_pipeline/test_add_ml_phrases.py
index ae014e3e71..8af7a92640 100644
--- a/etc/scripts/dataset_pipeline/test_add_ml_phrases.py
+++ b/etc/scripts/dataset_pipeline/test_add_ml_phrases.py
@@ -170,6 +170,13 @@ def test_counts_a_phrase_that_is_not_in_the_text(self):
         assert not inject(rule, ['Apache License'], counts, dry_run=True)
         assert counts['not_found'] == 1
 
+    def test_marks_two_phrases_in_one_rule(self):
+        rule = make_rule(TEXT)
+        counts = new_counts()
+        assert inject(rule, ['MIT License', 'do things'], counts, dry_run=True)
+        assert counts['injected'] == 2
+        assert rule.text.count('{{') == rule.text.count('}}') == 2
+
     def test_does_not_mark_the_same_phrase_twice(self):
         rule = make_rule(TEXT)
         counts = new_counts()

From 45ec541a5b3110cf430dfc3054aa50fa442d0029 Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Sun, 9 Aug 2026 00:12:13 +0530
Subject: [PATCH 09/11] reuse get_base_rules_by_expression for the expression
 filter

---
 .../dataset_pipeline/add_ml_phrases.py        | 17 ++++++--------
 .../dataset_pipeline/test_add_ml_phrases.py   | 23 +++++++++++++++----
 2 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/etc/scripts/dataset_pipeline/add_ml_phrases.py b/etc/scripts/dataset_pipeline/add_ml_phrases.py
index b52e24bec3..12f97e23f8 100644
--- a/etc/scripts/dataset_pipeline/add_ml_phrases.py
+++ b/etc/scripts/dataset_pipeline/add_ml_phrases.py
@@ -14,9 +14,9 @@
 
 sys.path.insert(0, str(Path(__file__).parent))
 
-from licensedcode.models import get_rules_by_expression
 from licensedcode.required_phrases import add_required_phrase_to_rule
 from licensedcode.required_phrases import find_phrase_spans_in_text
+from licensedcode.required_phrases import get_base_rules_by_expression
 from licensedcode.required_phrases import RequiredPhraseRuleCandidate
 from licensedcode.tokenize import get_existing_required_phrase_spans
 from licensedcode.tokenize import required_phrase_splitter
@@ -133,16 +133,13 @@ def select_rules(license_expression=None):
     """Rules that can take new required phrases, by license expression
 
     get_updatable_rules_by_expression reloads every rule file each time it is
-    called and skips everything when passed None, so load the rules once here
-    and do the filtering in memory
+    called and skips everything when passed None, so start from the base mapping
+    and do the filtering here
     """
-    rules_by_expression = get_rules_by_expression()
-
-    if license_expression:
-        rules = rules_by_expression.get(license_expression)
-        if not rules:
-            raise click.ClickException(f'no rules for license expression: {license_expression}')
-        rules_by_expression = {license_expression: rules}
+    try:
+        rules_by_expression = get_base_rules_by_expression(license_expression)
+    except KeyError:
+        raise click.ClickException(f'no rules for license expression: {license_expression}')
 
     selected = {}
     for expression, rules in rules_by_expression.items():
diff --git a/etc/scripts/dataset_pipeline/test_add_ml_phrases.py b/etc/scripts/dataset_pipeline/test_add_ml_phrases.py
index 8af7a92640..8fc331e3bc 100644
--- a/etc/scripts/dataset_pipeline/test_add_ml_phrases.py
+++ b/etc/scripts/dataset_pipeline/test_add_ml_phrases.py
@@ -205,21 +205,34 @@ def test_a_rule_with_no_words(self):
         assert predict_phrases(StubTagger(), FakeTokenizer(), 512, []) == ([], False)
 
 
+def patch_base_rules(monkeypatch, rules):
+    """Stand in for get_base_rules_by_expression, KeyError and all"""
+
+    def base_rules(license_expression=None):
+        if license_expression:
+            return {license_expression: rules[license_expression]}
+        return rules
+
+    monkeypatch.setattr(add_ml_phrases, 'get_base_rules_by_expression', base_rules)
+
+
 class TestSelectRules:
 
     def test_filters_and_groups(self, monkeypatch):
-        rules = {
+        patch_base_rules(monkeypatch, {
             'mit': [FakeRule(), FakeRule(is_from_license=True)],
             'bsd-new': [FakeRule(skip=True)],
-        }
-        monkeypatch.setattr(add_ml_phrases, 'get_rules_by_expression', lambda: rules)
+        })
         selected = select_rules()
         assert list(selected) == ['mit']
         assert len(selected['mit']) == 1
 
+    def test_one_expression(self, monkeypatch):
+        patch_base_rules(monkeypatch, {'mit': [FakeRule()], 'bsd-new': [FakeRule()]})
+        assert list(select_rules('mit')) == ['mit']
+
     def test_unknown_expression(self, monkeypatch):
-        rules = {'mit': [FakeRule()]}
-        monkeypatch.setattr(add_ml_phrases, 'get_rules_by_expression', lambda: rules)
+        patch_base_rules(monkeypatch, {'mit': [FakeRule()]})
         with pytest.raises(click.ClickException):
             select_rules('nope-1.0')
 

From 82cca13966958ee160e859d6f951e35467b6dcf9 Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Sun, 9 Aug 2026 00:42:26 +0530
Subject: [PATCH 10/11] skip the training only buffer when loading weights

aux_ce_weight registers class_weights during training, a strict load into the
inference model choked on it
---
 etc/scripts/dataset_pipeline/add_ml_phrases.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/etc/scripts/dataset_pipeline/add_ml_phrases.py b/etc/scripts/dataset_pipeline/add_ml_phrases.py
index 12f97e23f8..89f89f3936 100644
--- a/etc/scripts/dataset_pipeline/add_ml_phrases.py
+++ b/etc/scripts/dataset_pipeline/add_ml_phrases.py
@@ -91,7 +91,11 @@ def load_model(model, hf_token=None):
         raise click.ClickException('need a fast tokenizer for word_ids, got a slow one')
 
     tagger = PhraseTagger(config)
-    tagger.load_state_dict(load_file(str(model_dir / 'model.safetensors')), strict=True)
+    state = load_file(str(model_dir / 'model.safetensors'))
+    # training registers class_weights for the auxiliary loss, we compute no loss
+    # so the buffer is not there to load into
+    state.pop('class_weights', None)
+    tagger.load_state_dict(state, strict=True)
     # only needed while training and it warns under no_grad
     tagger.backbone.gradient_checkpointing_disable()
     tagger.eval()

From dd74fa6818c3b5e1463ee6e02d9d285fb96473fe Mon Sep 17 00:00:00 2001
From: Kaushik Kumar <kaushikrjpm10@gmail.com>
Date: Mon, 10 Aug 2026 19:26:52 +0530
Subject: [PATCH 11/11] fail early when the checkpoint has no weights file

building the tagger pulls the base model first, so check before that
---
 etc/scripts/dataset_pipeline/add_ml_phrases.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/etc/scripts/dataset_pipeline/add_ml_phrases.py b/etc/scripts/dataset_pipeline/add_ml_phrases.py
index 89f89f3936..00a5aa4901 100644
--- a/etc/scripts/dataset_pipeline/add_ml_phrases.py
+++ b/etc/scripts/dataset_pipeline/add_ml_phrases.py
@@ -90,8 +90,13 @@ def load_model(model, hf_token=None):
     if not tokenizer.is_fast:
         raise click.ClickException('need a fast tokenizer for word_ids, got a slow one')
 
+    # check this before building the backbone, that pulls the base model first
+    weights = model_dir / 'model.safetensors'
+    if not weights.exists():
+        raise click.ClickException(f'no model.safetensors in {model_dir}')
+
     tagger = PhraseTagger(config)
-    state = load_file(str(model_dir / 'model.safetensors'))
+    state = load_file(str(weights))
     # training registers class_weights for the auxiliary loss, we compute no loss
     # so the buffer is not there to load into
     state.pop('class_weights', None)