diff --git a/etc/scripts/dataset_pipeline/build_dataset.py b/etc/scripts/dataset_pipeline/build_dataset.py new file mode 100644 index 0000000000..72320c4e75 --- /dev/null +++ b/etc/scripts/dataset_pipeline/build_dataset.py @@ -0,0 +1,205 @@ +# -*- coding: utf-8 -*- +# +# Copyright (c) nexB Inc. and others. All rights reserved. +# ScanCode is a trademark of nexB Inc. +# SPDX-License-Identifier: Apache-2.0 +# See http://www.apache.org/licenses/LICENSE-2.0 for the license text. +# See https://github.com/nexB/scancode-toolkit for support or download. +# See https://aboutcode.org for more information about nexB OSS projects. +# + +"""Build a BIOES dataset from required phrases marked in license rules.""" + +from collections import Counter +import hashlib +import json +from pathlib import Path +import unicodedata + +import click + +from licensedcode.models import Rule +from licensedcode.models import rules_data_dir as default_rules_data_dir +from licensedcode.required_phrases import get_required_phrase_verbatim +from licensedcode.tokenize import get_existing_required_phrase_spans +from licensedcode.tokenize import required_phrase_splitter + + +def get_rule_type(rule): + """Return the first license rule type set on ``rule``.""" + for flag in rule.license_flag_names: + if getattr(rule, flag): + return flag + if rule.is_false_positive: + return 'is_false_positive' + return 'unknown' + + +def tag_tokens(text): + """Return rule text tokens and their required phrase BIOES labels.""" + tokens = [] + labels = [] + in_phrase = False + phrase_length = 0 + + for token in required_phrase_splitter(text): + if token == '{{': + in_phrase = True + phrase_length = 0 + continue + + if token == '}}': + if in_phrase and phrase_length: + labels[-1] = 'S-REQ' if phrase_length == 1 else 'E-REQ' + in_phrase = False + phrase_length = 0 + continue + + tokens.append(token) + if in_phrase: + labels.append('B-REQ' if phrase_length == 0 else 'I-REQ') + phrase_length += 1 + else: + labels.append('O') + + return tokens, labels + + +def build_record(rule): + """Return a dataset record for an annotated rule, or None.""" + if rule.is_required_phrase or not rule.text: + return + + text = rule.text.replace('\r\n', '\n').replace('\r', '\n') + text = unicodedata.normalize('NFKC', text) + + # Fail on invalid nested, empty, or dangling required phrase markers. + get_existing_required_phrase_spans(text) + if not any(get_required_phrase_verbatim(text)): + return + + tokens, bioes_labels = tag_tokens(text) + return { + 'identifier': rule.identifier, + 'license_expression': rule.license_expression or '', + 'rule_type': get_rule_type(rule), + 'text': text.replace('{{', '').replace('}}', ''), + 'tokens': tokens, + 'bioes_labels': bioes_labels, + } + + +def split_records(records, common_expression_threshold=50): + """ + Return train, validation, and test records using a hybrid split. + + Keep rare license expressions in one split. Distribute records from common + expressions by identifier so each split represents their varied rule text. + """ + expression_counts = Counter( + record['license_expression'] + for record in records + ) + common_expressions = { + expression + for expression, count in expression_counts.items() + if count >= common_expression_threshold + } + + rare_expressions = sorted( + ( + expression + for expression in expression_counts + if expression not in common_expressions + ), + key=lambda expression: (-expression_counts[expression], expression), + ) + rare_record_count = sum( + expression_counts[expression] + for expression in rare_expressions + ) + targets = { + 'train': 0.8 * rare_record_count, + 'val': 0.1 * rare_record_count, + 'test': 0.1 * rare_record_count, + } + assigned_counts = {name: 0 for name in targets} + rare_assignments = {} + + for expression in rare_expressions: + split = min( + targets, + key=lambda name: assigned_counts[name] / targets[name], + ) + rare_assignments[expression] = split + assigned_counts[split] += expression_counts[expression] + + splits = {name: [] for name in targets} + for record in records: + expression = record['license_expression'] + if expression in common_expressions: + identifier = record['identifier'].encode('utf-8') + bucket = int(hashlib.md5(identifier).hexdigest(), 16) % 100 + if bucket < 80: + split = 'train' + elif bucket < 90: + split = 'val' + else: + split = 'test' + else: + split = rare_assignments[expression] + + splits[split].append(record) + + return splits + + +@click.command() +@click.option( + '--rules-dir', + type=click.Path(exists=True, file_okay=False), + default=None, + help='Path to rules directory (defaults to the ScanCode rules directory).', +) +@click.option( + '--output-dir', + type=click.Path(file_okay=False), + default='dataset-output', + help='Output directory for train, validation, and test JSONL files.', +) +def main(rules_dir, output_dir): + """Extract marked required phrases into a BIOES training dataset.""" + rules_path = Path(rules_dir or default_rules_data_dir) + rule_files = sorted(rules_path.glob('*.RULE')) + records = [] + + click.echo(f'scanning rules from: {rules_path}') + for rule_file in rule_files: + rule = Rule.from_file(rule_file=str(rule_file)) + record = build_record(rule) + if record: + records.append(record) + + splits = split_records(records) + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + for split_name, records_in_split in splits.items(): + split_file = output_path / f'{split_name}.jsonl' + with split_file.open('w', encoding='utf-8') as output: + for record in records_in_split: + output.write(json.dumps(record, ensure_ascii=False) + '\n') + + click.echo('\ndone') + click.echo(f' rules scanned: {len(rule_files)}') + click.echo(f' annotated: {len(records)}') + click.echo( + f' train: {len(splits["train"])} ' + f'val: {len(splits["val"])} ' + f'test: {len(splits["test"])}' + ) + click.echo(f' output: {output_path}') + + +if __name__ == '__main__': + main() diff --git a/tests/licensedcode/test_build_dataset.py b/tests/licensedcode/test_build_dataset.py new file mode 100644 index 0000000000..2b7a9573bf --- /dev/null +++ b/tests/licensedcode/test_build_dataset.py @@ -0,0 +1,199 @@ +# -*- coding: utf-8 -*- +# +# Copyright (c) nexB Inc. and others. All rights reserved. +# ScanCode is a trademark of nexB Inc. +# SPDX-License-Identifier: Apache-2.0 +# See http://www.apache.org/licenses/LICENSE-2.0 for the license text. +# See https://github.com/nexB/scancode-toolkit for support or download. +# See https://aboutcode.org for more information about nexB OSS projects. +# + +import json + +from click.testing import CliRunner +import pytest + +from licensedcode.models import InvalidRule +from licensedcode.models import Rule +from licensedcode.tokenize import InvalidRuleRequiredPhrase + +from etc.scripts.dataset_pipeline.build_dataset import build_record +from etc.scripts.dataset_pipeline.build_dataset import main +from etc.scripts.dataset_pipeline.build_dataset import split_records +from etc.scripts.dataset_pipeline.build_dataset import tag_tokens + + +def make_rule( + identifier='mit_test.RULE', + license_expression='mit', + text='Licensed under the {{MIT License}}.', + **kwargs, +): + return Rule( + identifier=identifier, + license_expression=license_expression, + text=text, + is_license_notice=True, + **kwargs, + ) + + +def test_tag_tokens_assigns_bioes_labels(): + tokens, labels = tag_tokens( + 'Use {{MIT}} or the {{Apache License Version}} terms.' + ) + + assert tokens == [ + 'Use', 'MIT', 'or', 'the', 'Apache', 'License', 'Version', 'terms' + ] + assert labels == [ + 'O', 'S-REQ', 'O', 'O', 'B-REQ', 'I-REQ', 'E-REQ', 'O' + ] + + +def test_build_record_returns_normalized_rule_data(): + rule = make_rule(text='Licensed under the {{ＭＩＴ License}}.\rTerms') + + record = build_record(rule) + + assert record == { + 'identifier': 'mit_test.RULE', + 'license_expression': 'mit', + 'rule_type': 'is_license_notice', + 'text': 'Licensed under the MIT License.\nTerms', + 'tokens': ['Licensed', 'under', 'the', 'MIT', 'License', 'Terms'], + 'bioes_labels': ['O', 'O', 'O', 'B-REQ', 'E-REQ', 'O'], + } + + +def test_build_record_skips_unannotated_rules_and_rejects_invalid_markers(): + assert build_record(make_rule(text='Licensed under the MIT License.')) is None + assert build_record(make_rule(text='')) is None + assert build_record(make_rule(is_required_phrase=True)) is None + + invalid_texts = ( + 'Empty {{}} marker', + 'Opening {{dangling marker', + 'Closing dangling}} marker', + 'Valid {{MIT}} and {{dangling', + ) + for text in invalid_texts: + with pytest.raises(InvalidRuleRequiredPhrase): + build_record(make_rule(text=text)) + + +def test_split_records_uses_the_hybrid_split_deterministically(): + records = [ + { + 'identifier': f'common_{index}.RULE', + 'license_expression': 'common', + } + for index in range(50) + ] + for expression in ('rare-a', 'rare-b', 'rare-c'): + records.extend( + { + 'identifier': f'{expression}_{index}.RULE', + 'license_expression': expression, + } + for index in range(5) + ) + + splits = split_records(records) + + assert splits == split_records(records) + assert sum(len(split) for split in splits.values()) == len(records) + split_by_identifier = { + record['identifier']: name + for name, split in splits.items() + for record in split + } + assert split_by_identifier['common_20.RULE'] == 'train' + assert split_by_identifier['common_43.RULE'] == 'val' + assert split_by_identifier['common_3.RULE'] == 'test' + assert split_by_identifier['rare-a_0.RULE'] == 'train' + assert split_by_identifier['rare-b_0.RULE'] == 'val' + assert split_by_identifier['rare-c_0.RULE'] == 'test' + + for expression in ('rare-a', 'rare-b', 'rare-c'): + containing_splits = [ + name + for name, split in splits.items() + if any(record['license_expression'] == expression for record in split) + ] + assert len(containing_splits) == 1 + + +def test_main_writes_the_complete_dataset(tmp_path): + rules_dir = tmp_path / 'rules' + output_dir = tmp_path / 'dataset' + rules_dir.mkdir() + + rules = [ + make_rule( + identifier='mit_test.RULE', + license_expression='mit', + text='Licensed under the {{MIT License}}.', + ), + make_rule( + identifier='apache_test.RULE', + license_expression='apache-2.0', + text='Licensed under the {{Apache License}}.', + ), + make_rule( + identifier='bsd_test.RULE', + license_expression='bsd-new', + text='Licensed under the {{BSD License}}.', + ), + make_rule( + identifier='unmarked_test.RULE', + text='Licensed under the MIT License.', + ), + make_rule( + identifier='required_phrase_test.RULE', + text='MIT License', + is_required_phrase=True, + ), + ] + for rule in rules: + rule.dump(str(rules_dir)) + + result = CliRunner().invoke( + main, + ['--rules-dir', str(rules_dir), '--output-dir', str(output_dir)], + ) + + assert result.exit_code == 0, result.output + split_files = sorted(path.name for path in output_dir.glob('*.jsonl')) + assert split_files == ['test.jsonl', 'train.jsonl', 'val.jsonl'] + + records = [] + for split_file in output_dir.glob('*.jsonl'): + records.extend( + json.loads(line) + for line in split_file.read_text(encoding='utf-8').splitlines() + ) + + assert {record['identifier'] for record in records} == { + 'apache_test.RULE', + 'bsd_test.RULE', + 'mit_test.RULE', + } + assert all(len(record['tokens']) == len(record['bioes_labels']) for record in records) + assert all('{{' not in record['text'] and '}}' not in record['text'] for record in records) + + +def test_main_fails_when_a_rule_cannot_be_loaded(tmp_path): + rules_dir = tmp_path / 'rules' + output_dir = tmp_path / 'dataset' + rules_dir.mkdir() + (rules_dir / 'broken.RULE').write_text('', encoding='utf-8') + + result = CliRunner().invoke( + main, + ['--rules-dir', str(rules_dir), '--output-dir', str(output_dir)], + ) + + assert isinstance(result.exception, InvalidRule) + assert 'broken.RULE' in str(result.exception) + assert not output_dir.exists()