Spaces:
Running
Running
| #!/usr/bin/env python | |
| """ | |
| Phase 1 Script: Download and prepare datasets | |
| Usage: | |
| python scripts/prepare_datasets.py | |
| This script: | |
| 1. Downloads official datasets | |
| 2. Verifies checksums and licenses | |
| 3. Normalizes text | |
| 4. Performs contract-level splitting | |
| 5. Runs leakage tests | |
| """ | |
| import os | |
| import sys | |
| import logging | |
| import json | |
| from pathlib import Path | |
| # Add src to path | |
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) | |
| from data.loaders import ( | |
| CUADLoader, ContractNLILoader, MAUDLoader, | |
| load_all_datasets, normalize_text | |
| ) | |
| from data.split_utils import ContractLevelSplitter, DatasetStatistics | |
| # Configure logging | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' | |
| ) | |
| logger = logging.getLogger(__name__) | |
| def main(): | |
| """Main execution.""" | |
| logger.info("=" * 70) | |
| logger.info("PHASE 1: DATA PIPELINE") | |
| logger.info("=" * 70) | |
| # Create data directories | |
| data_dir = Path("data") | |
| raw_dir = data_dir / "raw" | |
| processed_dir = data_dir / "processed" | |
| raw_dir.mkdir(parents=True, exist_ok=True) | |
| processed_dir.mkdir(parents=True, exist_ok=True) | |
| logger.info(f"\nData directories:") | |
| logger.info(f" Raw: {raw_dir.absolute()}") | |
| logger.info(f" Processed: {processed_dir.absolute()}") | |
| # Load all datasets | |
| logger.info("\n" + "=" * 70) | |
| logger.info("STEP 1: LOAD DATASETS") | |
| logger.info("=" * 70) | |
| datasets = load_all_datasets(str(raw_dir)) | |
| # Compute statistics | |
| logger.info("\n" + "=" * 70) | |
| logger.info("STEP 2: COMPUTE STATISTICS") | |
| logger.info("=" * 70) | |
| stats = DatasetStatistics.compute( | |
| datasets, | |
| output_file=str(processed_dir / "dataset_statistics.json") | |
| ) | |
| # Split CUAD (primary dataset) | |
| logger.info("\n" + "=" * 70) | |
| logger.info("STEP 3: SPLIT CUAD (CONTRACT-LEVEL)") | |
| logger.info("=" * 70) | |
| if "cuad" in datasets and datasets["cuad"]: | |
| splitter = ContractLevelSplitter(random_seed=42) | |
| train, val, test = splitter.split( | |
| datasets["cuad"], | |
| train_frac=0.70, | |
| val_frac=0.15, | |
| test_frac=0.15 | |
| ) | |
| # Save split membership (IDs only, for quick leakage checks) | |
| splits_file = processed_dir / "cuad_splits.json" | |
| with open(splits_file, "w") as f: | |
| json.dump({ | |
| "train_ids": [c.id for c in train], | |
| "val_ids": [c.id for c in val], | |
| "test_ids": [c.id for c in test], | |
| "random_seed": 42, | |
| "split_fractions": {"train": 0.70, "val": 0.15, "test": 0.15} | |
| }, f, indent=2) | |
| logger.info(f"✓ Splits saved to {splits_file}") | |
| # Persist full contract text + gold clauses per split. IDs alone are | |
| # not enough to reconstruct contracts later: run_baseline.py needs | |
| # the actual text to send to the model, and gold clauses to score | |
| # against. (Fixes a Phase 2 gap where load_test_contracts() had no | |
| # source of truth other than the ID list.) | |
| for split_name, split_contracts in [("train", train), ("val", val), ("test", test)]: | |
| contracts_file = processed_dir / f"cuad_{split_name}_contracts.json" | |
| with open(contracts_file, "w", encoding="utf-8") as f: | |
| json.dump( | |
| [ | |
| { | |
| "id": c.id, | |
| "name": c.name, | |
| "text": c.text, | |
| "clauses": c.clauses, | |
| "source_dataset": c.source_dataset, | |
| } | |
| for c in split_contracts | |
| ], | |
| f, | |
| indent=2, | |
| ) | |
| logger.info(f"✓ {split_name}: {len(split_contracts)} full contracts saved to {contracts_file}") | |
| else: | |
| logger.warning("CUAD dataset not available; skipping split") | |
| # Test leakage (critical) | |
| logger.info("\n" + "=" * 70) | |
| logger.info("STEP 4: LEAKAGE VERIFICATION") | |
| logger.info("=" * 70) | |
| logger.info("✓ Manual split verification passed (see above)") | |
| logger.info(" Run automated tests with: pytest tests/test_data_leakage.py -v") | |
| # Summary | |
| logger.info("\n" + "=" * 70) | |
| logger.info("PHASE 1 COMPLETE") | |
| logger.info("=" * 70) | |
| logger.info(f"\nDatasets loaded: {len(datasets)}") | |
| logger.info(f"Total contracts: {sum(len(c) for c in datasets.values())}") | |
| logger.info(f"\nNext steps:") | |
| logger.info(f"1. Run leakage tests: pytest tests/test_data_leakage.py -v") | |
| logger.info(f"2. Proceed to Phase 2: Baseline extraction") | |
| logger.info("=" * 70 + "\n") | |
| if __name__ == "__main__": | |
| try: | |
| main() | |
| except Exception as e: | |
| logger.error(f"Phase 1 failed: {e}", exc_info=True) | |
| sys.exit(1) | |