From f506f9335bde7fb402a121339071d113c9b8aa23 Mon Sep 17 00:00:00 2001 From: Kaushik Kumar Date: Mon, 7 Sep 2026 19:00:57 +0530 Subject: [PATCH] Add hardened required phrase model inference Signed-off-by: Kaushik Kumar --- .../dataset_pipeline/add_ml_phrases.py | 306 +++++++++++++++ .../dataset_pipeline/test_add_ml_phrases.py | 349 ++++++++++++++++++ 2 files changed, 655 insertions(+) create mode 100644 etc/scripts/dataset_pipeline/add_ml_phrases.py create mode 100644 etc/scripts/dataset_pipeline/test_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..fa54c27fd6 --- /dev/null +++ b/etc/scripts/dataset_pipeline/add_ml_phrases.py @@ -0,0 +1,306 @@ +# Run the trained phrase tagger over license rules and mark its predictions. +import json +import os +import sys +import unicodedata +from numbers import Integral +from pathlib import Path + +import click + +# Avoid importing TensorFlow through Transformers. +os.environ.setdefault("USE_TF", "0") + +sys.path.insert(0, str(Path(__file__).parent)) + +from licensedcode.models import rules_data_dir +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 + +from train_model import extract_spans +from train_model import ID2LABEL +from train_model import load_final_model + + +MIN_TOKENS = 2 +MIN_SINGLE_TOKEN_LEN = 5 +MAX_RULE_TEXT = 4000 + + +def load_model(model, hf_token=None): + """Load and validate a local or Hugging Face Final_Model.""" + 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)) + + tagger, tokenizer = load_final_model(model_dir, offline=True) + config = json.loads((model_dir / "train_config.json").read_text(encoding="utf-8")) + return tagger, tokenizer, config["max_length"] + + +def words_from_text(text): + """Return words tokenized as they are in the training dataset.""" + text = text.replace("\r\n", "\n").replace("\r", "\n") + return required_phrase_splitter(unicodedata.normalize("NFKC", text)) + + +def is_updatable(rule): + """Return True if a rule can receive predicted required phrases.""" + if rule.is_from_license: + return False + if len(rule.text) > MAX_RULE_TEXT: + return False + 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): + """Return eligible rules grouped by license expression.""" + try: + rules_by_expression = get_base_rules_by_expression(license_expression) + except KeyError: + raise click.ClickException( + f"No rules for license expression: {license_expression}" + ) from None + + 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 _word_counts(word_ids): + counts = {} + for word_id in word_ids: + if word_id is not None: + counts[word_id] = counts.get(word_id, 0) + 1 + return counts + + +def encode_words(tokenizer, words, max_length): + """Encode the longest complete-word prefix and report truncation.""" + call = dict(is_split_into_words=True, add_special_tokens=True) + full = tokenizer(words, truncation=False, **call) + encoding = tokenizer(words, truncation=True, max_length=max_length, **call) + + full_counts = _word_counts(full.word_ids()) + retained_counts = _word_counts(encoding.word_ids()) + covered_words = max(retained_counts, default=-1) + 1 + complete_words = covered_words + + if covered_words and retained_counts[covered_words - 1] != full_counts[covered_words - 1]: + complete_words -= 1 + encoding = tokenizer(words[:complete_words], truncation=False, **call) + + if not complete_words: + raise ValueError("Tokenizer retained no complete words") + if len(encoding["input_ids"]) > max_length: + raise ValueError("Complete-word encoding exceeds the model maximum length") + + return encoding, complete_words < len(words) + + +def phrases_from_tags(tags, words): + """Return unique predicted phrase texts, longest first.""" + phrases = { + " ".join(words[start : end + 1]) + for start, end in extract_spans(tags) + } + return sorted(phrases, key=lambda phrase: (-len(phrase), phrase)) + + +def predict_phrases(tagger, tokenizer, max_length, words): + """Return predicted phrases and whether the rule was truncated.""" + if not words: + return [], False + + import torch + + encoding, truncated = encode_words(tokenizer, words, max_length) + word_ids = encoding.word_ids() + input_ids = torch.tensor([encoding["input_ids"]], dtype=torch.long) + attention_mask = torch.tensor([encoding["attention_mask"]], dtype=torch.long) + + with torch.no_grad(): + predicted = tagger.predict_words(input_ids, attention_mask, word_ids) + + word_count = len(set(word_id for word_id in word_ids if word_id is not None)) + if len(predicted) != word_count: + raise ValueError("Model returned a different number of labels than encoded words") + + tags = [] + for label in predicted: + if isinstance(label, bool) or not isinstance(label, Integral) or label not in ID2LABEL: + raise ValueError(f"Model returned an invalid label ID: {label!r}") + tags.append(ID2LABEL[int(label)]) + + return phrases_from_tags(tags, words), truncated + + +def new_counts(): + 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): + """Validate and add predicted phrases, writing the rule at most once.""" + candidates = [] + 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 + if not find_phrase_spans_in_text(rule.text, phrase): + counts["not_found"] += 1 + continue + candidates.append(phrase) + + if not candidates: + return False + + original_text = rule.text + original_source = rule.source + source = f"{original_source} ml_model" if original_source else "ml_model" + + for phrase in candidates: + updated = add_required_phrase_to_rule( + rule=rule, + required_phrase=phrase, + source=source, + debug=verbose, + dry_run=True, + ) + if updated: + counts["injected"] += 1 + else: + counts["skipped"] += 1 + + if rule.text == original_text: + return False + if not dry_run: + rule.dump(rules_data_dir) + return True + + +def process_rules( + selected, + tagger, + tokenizer, + max_length, + dry_run=False, + limit=0, + verbose=False, +): + """Predict and mark phrases in selected rules and return run counts.""" + counts = new_counts() + 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) + 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="Final model directory or Hugging Face repository.", +) +@click.option( + "--license-expression", + help="Only update rules for this license expression.", +) +@click.option( + "--dry-run", + is_flag=True, + help="Predict and validate phrases without saving rules.", +) +@click.option( + "--limit", + default=0, + type=click.IntRange(min=0), + help="Stop after this many rules; zero processes all rules.", +) +@click.option( + "-v", + "--verbose", + is_flag=True, + help="Print predictions for each rule.", +) +@click.help_option("-h", "--help") +def main(model, license_expression, dry_run, limit, verbose): + """Add model-predicted required phrases to license rules.""" + selected = select_rules(license_expression=license_expression) + if not selected: + click.echo("No eligible rules found") + return + + tagger, tokenizer, max_length = load_model( + model, + hf_token=os.environ.get("HF_TOKEN"), + ) + counts = process_rules( + selected=selected, + tagger=tagger, + tokenizer=tokenizer, + max_length=max_length, + dry_run=dry_run, + limit=limit, + verbose=verbose, + ) + + 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']}") + 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"]: + click.echo("Run scancode-reindex-licenses to use the new required phrases") + + +if __name__ == "__main__": + main() 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..7d5e985860 --- /dev/null +++ b/etc/scripts/dataset_pipeline/test_add_ml_phrases.py @@ -0,0 +1,349 @@ +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +import click +import pytest +from click.testing import CliRunner + +sys.path.insert(0, str(Path(__file__).parent)) + +import add_ml_phrases +from add_ml_phrases import encode_words +from add_ml_phrases import inject +from add_ml_phrases import is_updatable +from add_ml_phrases import load_model +from add_ml_phrases import main +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 +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: + 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__( + input_ids=list(range(len(word_ids))), + attention_mask=[1] * len(word_ids), + ) + self._word_ids = word_ids + + def word_ids(self): + return list(self._word_ids) + + +class FakeTokenizer: + """Tokenize words with configurable subword counts.""" + + is_fast = True + + def __init__(self, subwords=None): + self.subwords = subwords or {} + + def __call__(self, words, truncation, max_length=None, **kwargs): + word_ids = [None] + for index, word in enumerate(words): + word_ids.extend([index] * self.subwords.get(word, 1)) + word_ids.append(None) + if truncation and len(word_ids) > max_length: + word_ids = word_ids[: max_length - 1] + [None] + return FakeEncoding(word_ids) + + +class StubTagger: + def __init__(self, labels): + self.labels = labels + self.calls = 0 + + def predict_words(self, input_ids, attention_mask, word_ids): + self.calls += 1 + count = len(set(word_id for word_id in word_ids if word_id is not None)) + labels = self.labels + [LABEL2ID["O"]] * count + return labels[:count] + + +class TestLoadModel: + def test_loads_a_valid_local_final_model(self, tmp_path, monkeypatch): + config = tmp_path / "train_config.json" + config.write_text(json.dumps({"max_length": 256}), encoding="utf-8") + tagger = object() + tokenizer = object() + calls = [] + + def load_final_model(model_dir, offline): + calls.append((model_dir, offline)) + return tagger, tokenizer + + monkeypatch.setattr(add_ml_phrases, "load_final_model", load_final_model) + assert load_model(tmp_path) == (tagger, tokenizer, 256) + assert calls == [(tmp_path, True)] + + def test_downloads_before_using_the_strict_loader(self, tmp_path, monkeypatch): + model_dir = tmp_path / "snapshot" + model_dir.mkdir() + (model_dir / "train_config.json").write_text( + json.dumps({"max_length": 512}), encoding="utf-8" + ) + downloads = [] + hub = SimpleNamespace( + snapshot_download=lambda repo_id, token: downloads.append((repo_id, token)) + or str(model_dir) + ) + monkeypatch.setitem(sys.modules, "huggingface_hub", hub) + monkeypatch.setattr( + add_ml_phrases, + "load_final_model", + lambda path, offline: (object(), object()), + ) + + load_model("owner/model", hf_token="secret") + assert downloads == [("owner/model", "secret")] + + +class TestWordsFromText: + def test_uses_dataset_tokenization(self): + assert words_from_text("Apache-2.0 License") == ["Apache", "2", "0", "License"] + + def test_normalizes_line_endings_and_unicode(self): + assert words_from_text("a\ufb01x\r\ntwo\rthree") == ["afix", "two", "three"] + + +class TestEncodeWords: + def test_keeps_all_complete_words(self): + encoding, truncated = encode_words(FakeTokenizer(), ["one", "two"], 10) + assert encoding.word_ids() == [None, 0, 1, None] + assert not truncated + + def test_removes_a_partially_truncated_word(self): + tokenizer = FakeTokenizer({"many": 3}) + encoding, truncated = encode_words(tokenizer, ["one", "many", "three"], 4) + assert encoding.word_ids() == [None, 0, None] + assert truncated + + def test_rejects_when_no_complete_word_fits(self): + with pytest.raises(ValueError, match="no complete words"): + encode_words(FakeTokenizer({"many": 4}), ["many"], 3) + + +class TestPhrasesFromTags: + def test_returns_longest_unique_phrases_first(self): + words = ["mit", "license", "mit", "other"] + tags = ["B-REQ", "E-REQ", "S-REQ", "O"] + assert phrases_from_tags(tags, words) == ["mit license", "mit"] + + def test_rejects_invalid_bioes(self): + with pytest.raises(ValueError, match="invalid BIOES"): + phrases_from_tags(["B-REQ", "O"], ["one", "two"]) + + +class TestPredictPhrases: + def test_predicts_from_complete_words(self): + labels = [LABEL2ID["B-REQ"], LABEL2ID["E-REQ"], LABEL2ID["O"]] + phrases, truncated = predict_phrases( + StubTagger(labels), FakeTokenizer(), 10, ["MIT", "License", "text"] + ) + assert phrases == ["MIT License"] + assert not truncated + + def test_keeps_a_valid_span_at_the_truncation_boundary(self): + labels = [LABEL2ID["O"], LABEL2ID["S-REQ"]] + phrases, truncated = predict_phrases( + StubTagger(labels), FakeTokenizer(), 4, ["one", "two", "three"] + ) + assert phrases == ["two"] + assert truncated + + @pytest.mark.parametrize("label", [99, 1.5, True]) + def test_rejects_an_invalid_label_id(self, label): + with pytest.raises(ValueError, match="invalid label ID"): + predict_phrases(StubTagger([label]), FakeTokenizer(), 10, ["one"]) + + def test_rejects_the_wrong_number_of_labels(self): + tagger = StubTagger([]) + tagger.predict_words = lambda *args: [] + with pytest.raises(ValueError, match="different number"): + predict_phrases(tagger, FakeTokenizer(), 10, ["one"]) + + def test_empty_text_does_not_call_the_model(self): + tagger = StubTagger([]) + assert predict_phrases(tagger, FakeTokenizer(), 10, []) == ([], False) + assert tagger.calls == 0 + + +class TestIsUpdatable: + @pytest.mark.parametrize( + "rule", + [ + FakeRule(is_from_license=True), + FakeRule(text="x" * 4001), + FakeRule(is_approx_matchable=False), + FakeRule(skip=True), + FakeRule(text="under the {{mit license}} terms"), + ], + ) + def test_excludes_ineligible_rules(self, rule): + assert not is_updatable(rule) + + def test_accepts_a_plain_rule(self): + assert is_updatable(FakeRule()) + + +class TestSelectRules: + def test_filters_and_groups_rules(self, monkeypatch): + monkeypatch.setattr( + add_ml_phrases, + "get_base_rules_by_expression", + lambda expression: { + "mit": [FakeRule(), FakeRule(is_from_license=True)], + "bsd-new": [FakeRule(skip=True)], + }, + ) + selected = select_rules() + assert list(selected) == ["mit"] + assert len(selected["mit"]) == 1 + + def test_forwards_the_expression_filter(self, monkeypatch): + expressions = [] + + def get_rules(expression): + expressions.append(expression) + return {expression: [FakeRule()]} + + monkeypatch.setattr(add_ml_phrases, "get_base_rules_by_expression", get_rules) + assert list(select_rules("mit")) == ["mit"] + assert expressions == ["mit"] + + def test_reports_an_unknown_expression(self, monkeypatch): + def get_rules(expression): + raise KeyError(expression) + + monkeypatch.setattr(add_ml_phrases, "get_base_rules_by_expression", get_rules) + with pytest.raises(click.ClickException, match="No rules"): + select_rules("unknown") + + +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 + + +TEXT = "Permission is granted under the MIT License to do things with this" + + +class TestInject: + def test_marks_two_phrases_and_preserves_source(self): + rule = make_rule(TEXT, source="mit_1.RULE") + 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 + assert rule.source == "mit_1.RULE ml_model" + + def test_rejects_an_unsuitable_phrase(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_not_found_in_the_rule(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_writes_a_rule_once(self, monkeypatch): + rule = make_rule(TEXT) + writes = [] + monkeypatch.setattr(Rule, "dump", lambda self, directory: writes.append(directory)) + assert inject(rule, ["MIT License", "do things"], new_counts()) + assert writes == [add_ml_phrases.rules_data_dir] + +class TestProcessRules: + def test_processes_selected_rules(self): + rule = make_rule(TEXT) + labels = [LABEL2ID["B-REQ"], LABEL2ID["I-REQ"], LABEL2ID["E-REQ"]] + counts = process_rules( + selected={"mit": [rule]}, + tagger=StubTagger(labels), + tokenizer=FakeTokenizer(), + max_length=50, + dry_run=True, + ) + assert counts["rules"] == 1 + assert counts["injected"] == 1 + assert counts["written"] == 1 + + def test_limit_stops_before_another_rule(self): + rules = [make_rule(TEXT) for _ in range(3)] + labels = [LABEL2ID["B-REQ"], LABEL2ID["I-REQ"], LABEL2ID["E-REQ"]] + counts = process_rules( + {"mit": rules}, StubTagger(labels), FakeTokenizer(), 50, dry_run=True, limit=2 + ) + assert counts["rules"] == 2 + + +class TestCommand: + def test_does_not_load_a_model_when_no_rules_are_eligible(self, monkeypatch): + monkeypatch.setattr(add_ml_phrases, "select_rules", lambda **kwargs: {}) + + def fail(*args, **kwargs): + raise AssertionError("model should not load") + + monkeypatch.setattr(add_ml_phrases, "load_model", fail) + result = CliRunner().invoke(main, ["--model", "unused"]) + assert result.exit_code == 0 + assert "No eligible rules found" in result.output + + def test_wires_selection_loading_and_processing(self, monkeypatch, tmp_path): + selected = {"mit": [object()]} + tagger = object() + tokenizer = object() + calls = [] + monkeypatch.setattr(add_ml_phrases, "select_rules", lambda **kwargs: selected) + monkeypatch.setattr( + add_ml_phrases, "load_model", lambda *args, **kwargs: (tagger, tokenizer, 256) + ) + + def process(**kwargs): + calls.append(kwargs) + return new_counts() + + monkeypatch.setattr(add_ml_phrases, "process_rules", process) + result = CliRunner().invoke( + main, + ["--model", str(tmp_path), "--license-expression", "mit", "--dry-run"], + ) + assert result.exit_code == 0 + assert calls[0]["selected"] is selected + assert calls[0]["tagger"] is tagger + assert calls[0]["tokenizer"] is tokenizer + assert calls[0]["max_length"] == 256 + assert calls[0]["dry_run"] is True