""" Phase 1: Data splitting with contract-level leakage prevention This module implements contract-level train/val/test splitting, ensuring that no contract text appears in more than one split. CRITICAL: This prevents data leakage and ensures honest evaluation. """ import logging import random from typing import List, Tuple, Set, Dict from pathlib import Path import json from data.loaders import Contract logger = logging.getLogger(__name__) class ContractLevelSplitter: """Split contracts into train/val/test with zero-leakage guarantee.""" def __init__(self, random_seed: int = 42): """ Initialize splitter. Args: random_seed: Seed for reproducibility """ self.random_seed = random_seed random.seed(random_seed) def split( self, contracts: List[Contract], train_frac: float = 0.70, val_frac: float = 0.15, test_frac: float = 0.15, ) -> Tuple[List[Contract], List[Contract], List[Contract]]: """ Split contracts at the contract level (never split a contract across splits). Args: contracts: List of contracts to split train_frac: Fraction for training val_frac: Fraction for validation test_frac: Fraction for testing Returns: Tuple of (train_contracts, val_contracts, test_contracts) Raises: AssertionError: If fractions don't sum to 1.0 or leakage is detected """ assert abs(train_frac + val_frac + test_frac - 1.0) < 1e-6, \ f"Fractions must sum to 1.0, got {train_frac + val_frac + test_frac}" # Shuffle contracts shuffled = list(contracts) random.shuffle(shuffled) n = len(shuffled) train_end = int(n * train_frac) val_end = train_end + int(n * val_frac) train = shuffled[:train_end] val = shuffled[train_end:val_end] test = shuffled[val_end:] # Verify zero leakage self._verify_no_leakage(train, val, test) logger.info(f"Split {n} contracts:") logger.info(f" Train: {len(train)} ({len(train)/n*100:.1f}%)") logger.info(f" Val: {len(val)} ({len(val)/n*100:.1f}%)") logger.info(f" Test: {len(test)} ({len(test)/n*100:.1f}%)") return train, val, test @staticmethod def _verify_no_leakage( train: List[Contract], val: List[Contract], test: List[Contract] ) -> None: """ Verify that no contract appears in multiple splits. Raises: AssertionError: If leakage is detected """ train_ids = {c.id for c in train} val_ids = {c.id for c in val} test_ids = {c.id for c in test} # Check for overlaps train_val_overlap = train_ids & val_ids val_test_overlap = val_ids & test_ids train_test_overlap = train_ids & test_ids assert len(train_val_overlap) == 0, \ f"LEAKAGE: {len(train_val_overlap)} contracts in both train and val: {train_val_overlap}" assert len(val_test_overlap) == 0, \ f"LEAKAGE: {len(val_test_overlap)} contracts in both val and test: {val_test_overlap}" assert len(train_test_overlap) == 0, \ f"LEAKAGE: {len(train_test_overlap)} contracts in both train and test: {train_test_overlap}" # Check for duplicate text (additional verification) train_texts = {c.text for c in train} val_texts = {c.text for c in val} test_texts = {c.text for c in test} # Duplicates within splits are OK; duplicates across splits are NOT OK assert len(train_texts & val_texts) == 0, \ "LEAKAGE: Duplicate contract text in train and val" assert len(val_texts & test_texts) == 0, \ "LEAKAGE: Duplicate contract text in val and test" assert len(train_texts & test_texts) == 0, \ "LEAKAGE: Duplicate contract text in train and test" logger.info("✓ Leakage verification passed: zero contract overlap across splits") class DatasetStatistics: """Compute and log dataset statistics.""" @staticmethod def compute( datasets: Dict[str, List[Contract]], output_file: str = None ) -> Dict: """ Compute statistics across all datasets. Args: datasets: Dict of {name: [contracts]} output_file: Optional path to save statistics as JSON Returns: Dict of statistics """ stats = {} for dataset_name, contracts in datasets.items(): if not contracts: continue # Basic counts num_contracts = len(contracts) num_clauses = sum(len(c.clauses) for c in contracts) # Text statistics text_lengths = [len(c.text) for c in contracts] clause_counts = [len(c.clauses) for c in contracts] # Category distribution (for CUAD/MAUD) category_counts = {} for contract in contracts: for clause in contract.clauses: cat = clause.get("category", "unknown") category_counts[cat] = category_counts.get(cat, 0) + 1 stats[dataset_name] = { "num_contracts": num_contracts, "num_clauses": num_clauses, "clauses_per_contract": { "mean": sum(clause_counts) / len(clause_counts), "median": sorted(clause_counts)[len(clause_counts)//2], "min": min(clause_counts), "max": max(clause_counts), }, "text_length_chars": { "mean": sum(text_lengths) / len(text_lengths), "median": sorted(text_lengths)[len(text_lengths)//2], "min": min(text_lengths), "max": max(text_lengths), }, "categories": category_counts, } # Log summary logger.info("\n" + "=" * 70) logger.info("DATASET STATISTICS") logger.info("=" * 70) for dataset_name, ds_stats in stats.items(): logger.info(f"\n{dataset_name}:") logger.info(f" Contracts: {ds_stats['num_contracts']}") logger.info(f" Clauses: {ds_stats['num_clauses']}") logger.info(f" Clauses/contract (mean): {ds_stats['clauses_per_contract']['mean']:.1f}") logger.info(f" Text length chars (mean): {ds_stats['text_length_chars']['mean']:.0f}") if ds_stats['categories']: logger.info(f" Categories: {len(ds_stats['categories'])}") sorted_cats = sorted( ds_stats['categories'].items(), key=lambda x: x[1], reverse=True ) for cat, count in sorted_cats[:5]: logger.info(f" - {cat}: {count}") logger.info("=" * 70 + "\n") # Save to file if requested if output_file: with open(output_file, "w") as f: json.dump(stats, f, indent=2) logger.info(f"Statistics saved to {output_file}") return stats if __name__ == "__main__": logging.basicConfig(level=logging.INFO) from data.loaders import load_all_datasets # Load all datasets datasets = load_all_datasets() # Example: split CUAD train/test if "cuad_train" in datasets and datasets["cuad_train"]: splitter = ContractLevelSplitter(random_seed=42) train, val, test = splitter.split( datasets["cuad_train"], train_frac=0.70, val_frac=0.15, test_frac=0.15 ) logger.info(f"✓ Successfully split CUAD: {len(train)} train, {len(val)} val, {len(test)} test") # Compute statistics stats = DatasetStatistics.compute(datasets, output_file="data/statistics.json")