the-jashthakkar/CodeMode

1"""2Script to generate positive pairs and triplets from code chunks.3 4This script loads code chunks and generates:51. Positive Pairs: (question, code) with 4-5 variations per sample62. Triplets: (anchor_question, positive_code, negative_code)7 8Usage:9    python -m scripts.run_pairs_triplets_pipeline --chunks <path> --output <dir>10    python -m scripts.run_pairs_triplets_pipeline --help11 12Examples:13    # Generate from local chunks with default settings14    python -m scripts.run_pairs_triplets_pipeline \\15        --chunks data/processed/chunks/Local_saved_files/chunks.jsonl \\16        --output data/processed/training17 18    # Generate from repository chunks  19    python -m scripts.run_pairs_triplets_pipeline \\20        --chunks data/processed/repos/langgraph_20260116_123638/langgraph_chunks.jsonl \\21        --output data/processed/training/langgraph22 23    # Custom settings24    python -m scripts.run_pairs_triplets_pipeline \\25        --chunks data/processed/chunks/Local_saved_files/chunks.jsonl \\26        --output data/processed/training \\27        --pairs 100 --triplets 100 --variance 528"""29 30import sys31from pathlib import Path32 33# Add project root to path34PROJECT_ROOT = Path(__file__).parent.parent35sys.path.insert(0, str(PROJECT_ROOT))36 37from src.task_3_data_engineering.export.pairs_triplets_generator import (38    generate_pairs_and_triplets,39    main as cli_main40)41 42 43def run_default_pipeline():44    """Run with default settings for the available chunks."""45    46    # Try multiple possible chunk locations47    possible_paths = [48        PROJECT_ROOT / "data" / "processed" / "chunks" / "Local_saved_files" / "chunks.jsonl",49        PROJECT_ROOT / "data" / "processed" / "chunks" / "sample_code" / "chunks.jsonl",50    ]51    52    # Find all chunks.jsonl files in chunks folder subdirectories53    chunks_dir = PROJECT_ROOT / "data" / "processed" / "chunks"54    if chunks_dir.exists():55        for subdir in chunks_dir.iterdir():56            if subdir.is_dir():57                chunks_file = subdir / "chunks.jsonl"58                if chunks_file.exists() and chunks_file not in possible_paths:59                    possible_paths.append(chunks_file)60    61    # Find repository chunks62    repos_dir = PROJECT_ROOT / "data" / "processed" / "repos"63    if repos_dir.exists():64        for repo_dir in repos_dir.iterdir():65            if repo_dir.is_dir():66                for jsonl_file in repo_dir.glob("*_chunks.jsonl"):67                    possible_paths.append(jsonl_file)68    69    chunks_path = None70    for path in possible_paths:71        if path.exists():72            chunks_path = path73            break74    75    if chunks_path is None:76        print("❌ No chunks files found. Please specify a chunks file with --chunks")77        print("\nPossible locations checked:")78        for p in possible_paths[:5]:79            print(f"   - {p}")80        return81    82    output_dir = PROJECT_ROOT / "data" / "processed" / "training"83    84    print("=" * 60)85    print("🚀 Positive Pairs & Triplets Generator")86    print("=" * 60)87    print(f"\n📂 Chunks Path: {chunks_path}")88    print(f"📁 Output Dir: {output_dir}")89    print(f"📊 Settings: pairs=100, triplets=100, variance=5")90    print("\n" + "-" * 60)91    92    pairs, triplets = generate_pairs_and_triplets(93        chunks_path=chunks_path,94        output_dir=output_dir,95        num_pairs=100,96        num_triplets=100,97        variance=5,98        export_format="both"99    )100    101    print("\n" + "=" * 60)102    print("✅ Pipeline Complete!")103    print("=" * 60)104    print(f"\n📁 Output files saved to: {output_dir}")105    print("   - positive_pairs.jsonl")106    print("   - positive_pairs.json")107    print("   - triplets.jsonl")108    print("   - triplets.json")109 110 111if __name__ == "__main__":112    import argparse113    114    # Check if any arguments provided115    if len(sys.argv) > 1:116        # Use CLI with provided arguments117        cli_main()118    else:119        # Run with defaults120        run_default_pipeline()121