Add train/val/test split and fix URL stripping in normalize_phrase

Hybrid split: common expressions (50+ rules) get stratified per rule, rare ones stay in one split for zero leakage

References #5077

Author: Kaushik-Kumar-CEG

etc/scripts/dataset_pipeline/build_dataset.py

@@ -1,8 +1,10 @@
 # extracts required phrases from .RULE files
 # outputs a JSONL dataset for NER model training
+import hashlib
 import json
 import re
 import unicodedata
+from collections import Counter
 from pathlib import Path
 import click
 
@@ -66,20 +68,41 @@ def tag_tokens(text):
     return tokens, labels
 
 
+def assign_splits(results, threshold=50):
+    """80/10/10 split. common expressions (>= threshold rules) get split per-rule,
+    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}
+
+    # rare expressions: assign each to the split that needs more rules
+    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]
+
+    return heavy, assignment
+
+
 @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', default='dataset-output/required_phrases.jsonl',
-              help='Output JSONL path')
-def main(rules_dir, output):
+@click.option('--output-dir', default='dataset-output',
+              help='Output directory for train/val/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_path = Path(output)
-    out_path.parent.mkdir(parents=True, exist_ok=True)
+    out_dir = Path(output_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
 
     total_rules = 0
     annotated = 0
@@ -133,21 +156,37 @@ def main(rules_dir, output):
             'required_phrases': valid_phrases,
         })
 
-    # write jsonl
-    with open(out_path, 'w', encoding='utf-8') as f:
-        for entry in results:
-            f.write(json.dumps(entry, ensure_ascii=False) + '\n')
+    # 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:
+            # common expressions: hash rule name for 80/10/10
+            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')
 
     click.echo('\ndone')
     click.echo(f'  rules scanned: {total_rules}')
     click.echo(f'  annotated: {annotated}')
     click.echo(f'  phrases extracted: {total_phrases}')
-    click.echo(f'  output: {out_path}')
-
+    click.echo(f'  train: {len(splits["train"])}  val: {len(splits["val"])}  test: {len(splits["test"])}')
+    click.echo(f'  output: {out_dir}')
 
 # stuff to do(follow up commits):
-# train/val/test split by license expression
-
-
+# tests to be added in script
 if __name__ == '__main__':
     main()