Spaces:
Running
Running
| """ | |
| 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 | |
| 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.""" | |
| 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") | |