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..00a5aa4901 --- /dev/null +++ b/etc/scripts/dataset_pipeline/add_ml_phrases.py @@ -0,0 +1,335 @@ +# 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.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 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 +MIN_SINGLE_TOKEN_LEN = 5 + +# scancode leaves rule texts longer than this alone +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): + """The trained tagger, its tokenizer and the max length it was trained with + + ``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 = {} + + # 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 labels != LABELS: + raise click.ClickException(f'this checkpoint was trained on other labels: {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') + + # 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(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) + tagger.load_state_dict(state, 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 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 start from the base mapping + and do the filtering here + """ + 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(): + 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 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 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 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 + 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 + + +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 = new_counts() + + 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) + 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('-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, 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(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 pick up 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..8fc331e3bc --- /dev/null +++ b/etc/scripts/dataset_pipeline/test_add_ml_phrases.py @@ -0,0 +1,257 @@ +# tests for add_ml_phrases.py +# no model and no network here, the tagger is stubbed out +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 +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_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 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' + + +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_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() + 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 + + 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) + + +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): + patch_base_rules(monkeypatch, { + '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_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): + patch_base_rules(monkeypatch, {'mit': [FakeRule()]}) + with pytest.raises(click.ClickException): + select_rules('nope-1.0') + + +class TestProcessRules: + + def test_dry_run_marks_nothing_on_disk(self, monkeypatch): + rule = make_rule(TEXT) + 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): + 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