Refactor required phrase dataset extraction

Signed-off-by: Kaushik Kumar <kaushikrjpm10@gmail.com>

Author: Kaushik-Kumar-CEG

etc/scripts/dataset_pipeline/build_dataset.py

@@ -1,163 +1,204 @@
-# extracts required phrases from .RULE files
-# outputs a JSONL dataset for NER model training
+# -*- 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
-import unicodedata
-from collections import Counter
 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 is_* flag set on the rule"""
-    for flag in ('is_license_text', 'is_license_notice', 'is_license_reference',
-                 'is_license_tag', 'is_license_intro', 'is_license_clue',
-                 'is_false_positive'):
-        if getattr(rule, flag, False):
+    """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):
-    """Tag each word token with a BIOES label based on {{ }} markers"""
+    """Return rule text tokens and their required phrase BIOES labels."""
     tokens = []
     labels = []
     in_phrase = False
-    count = 0
+    phrase_length = 0
 
-    for tok in required_phrase_splitter(text):
-        if tok == '{{':
+    for token in required_phrase_splitter(text):
+        if token == '{{':
             in_phrase = True
-            count = 0
+            phrase_length = 0
             continue
-        if tok == '}}':
-            if in_phrase and count > 0:
-                labels[-1] = 'S-REQ' if count == 1 else 'E-REQ'
+
+        if token == '}}':
+            if in_phrase and phrase_length:
+                labels[-1] = 'S-REQ' if phrase_length == 1 else 'E-REQ'
             in_phrase = False
-            count = 0
+            phrase_length = 0
             continue
-        tokens.append(tok)
+
+        tokens.append(token)
         if in_phrase:
-            labels.append('B-REQ' if count == 0 else 'I-REQ')
-            count += 1
+            labels.append('B-REQ' if phrase_length == 0 else 'I-REQ')
+            phrase_length += 1
         else:
             labels.append('O')
 
-    assert len(tokens) == len(labels), f'token/label mismatch: {len(tokens)} vs {len(labels)}'
     return tokens, labels
 
 
-def assign_splits(results, threshold=50):
-    """80/10/10 split by license expression to prevent data leakage.
-    Expressions with >= threshold rules get split per-rule via hash,
-    rare ones stay together in one split"""
-    expr_counts = Counter(e['license_expression'] for e in results)
-    heavy = {e for e, c in expr_counts.items() if c >= threshold}
+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]
 
-    light_exprs = sorted((e for e in expr_counts if e not in heavy),
-                         key=lambda x: (-expr_counts[x], x))
-    total = sum(expr_counts[e] for e in light_exprs)
-    targets = {'train': 0.8 * total, 'val': 0.1 * total, 'test': 0.1 * total}
-    filled = {'train': 0, 'val': 0, 'test': 0}
-    assignment = {}
-    for expr in light_exprs:
-        best = min(targets, key=lambda s: filled[s] / max(targets[s], 1))
-        assignment[expr] = best
-        filled[best] += expr_counts[expr]
+        splits[split].append(record)
 
-    return heavy, assignment
+    return splits
 
 
 @click.command()
-@click.option('--rules-dir', type=click.Path(exists=True), default=None,
-              help='Path to rules directory (defaults to repo rules dir)')
-@click.option('--output-dir', default='dataset-output',
-              help='Output directory for train/val/test JSONL files')
+@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 required phrases from rule files for NER training"""
-    if not rules_dir:
-        repo_rules = Path(__file__).resolve().parents[3] / 'src' / 'licensedcode' / 'data' / 'rules'
-        rules_dir = str(repo_rules) if repo_rules.is_dir() else default_rules_data_dir
-
-    rules_path = Path(rules_dir)
-    out_dir = Path(output_dir)
-    out_dir.mkdir(parents=True, exist_ok=True)
-
-    total_rules = 0
-    annotated = 0
-    results = []
+    """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 rf in sorted(rules_path.glob('*.RULE')):
-        try:
-            rule = Rule.from_file(rule_file=str(rf))
-        except Exception as e:
-            click.echo(f'  skipping {rf.name}: {e}', err=True)
-            continue
-        total_rules += 1
-
-        if getattr(rule, 'is_required_phrase', False):
-            continue
-
-        text = rule.text or ''
-        if not text:
-            continue
-
-        # normalize line endings and unicode
-        text = text.replace('\r\n', '\n').replace('\r', '\n')
-        text = unicodedata.normalize('NFKC', text)
-
-        phrases = list(get_required_phrase_verbatim(text))
-        if not phrases:
-            continue
-
-        tokens, bioes_labels = tag_tokens(text)
-
-        # strip markers for the clean text field
-        clean_text = text.replace('{{', '').replace('}}', '')
-
-        annotated += 1
-        results.append({
-            'identifier': rule.identifier,
-            'license_expression': rule.license_expression or '',
-            'rule_type': get_rule_type(rule),
-            'text': clean_text,
-            'tokens': tokens,
-            'bioes_labels': bioes_labels,
-        })
-
-    # split by license expression and write
-    heavy, assignment = assign_splits(results)
-    splits = {'train': [], 'val': [], 'test': []}
-    for entry in results:
-        expr = entry['license_expression']
-        if expr in heavy:
-            bucket = int(hashlib.md5(entry['identifier'].encode('utf-8')).hexdigest(), 16) % 100
-            if bucket < 80:
-                splits['train'].append(entry)
-            elif bucket < 90:
-                splits['val'].append(entry)
-            else:
-                splits['test'].append(entry)
-        else:
-            splits[assignment[expr]].append(entry)
-
-    for name, records in splits.items():
-        path = out_dir / f'{name}.jsonl'
-        with open(path, 'w', encoding='utf-8') as f:
-            for entry in records:
-                f.write(json.dumps(entry, ensure_ascii=False) + '\n')
+    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: {total_rules}')
-    click.echo(f'  annotated: {annotated}')
-    click.echo(f'  train: {len(splits["train"])}  val: {len(splits["val"])}  test: {len(splits["test"])}')
-    click.echo(f'  output: {out_dir}')
+    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__':