the-jashthakkar/CodeMode

1"""2Generate training datasets for ALL frameworks automatically.3 4This script auto-discovers all chunk files and processes them,5generating separate datasets for each framework PLUS a combined dataset.6 7Usage:8    python scripts/generate_all_frameworks.py9    10Output Structure:11    data/processed/training_crewai/12        - positive_pairs.json13        - triplets.json14    data/processed/training_langgraph/15        - positive_pairs.json16        - triplets.json17    data/processed/training_combined/18        - positive_pairs.json  (ALL frameworks merged)19        - triplets.json        (ALL frameworks merged)20"""21 22import sys23import json24from pathlib import Path25from typing import List, Tuple26from dataclasses import asdict27 28# Add project root to path29PROJECT_ROOT = Path(__file__).parent.parent30sys.path.insert(0, str(PROJECT_ROOT))31 32from src.task_3_data_engineering.export.pairs_triplets_generator import (33    generate_pairs_and_triplets,34    PositivePair,35    Triplet36)37 38 39def discover_all_chunk_files() -> List[Tuple[Path, str]]:40    """41    Discover all chunk files in the workspace.42    43    Returns:44        List of (chunk_path, framework_name) tuples45    """46    chunk_files = []47    48    # Check local chunks49    local_paths = [50        PROJECT_ROOT / "data" / "processed" / "chunks" / "Local_saved_files" / "chunks.jsonl",51        PROJECT_ROOT / "data" / "processed" / "chunks" / "sample_code" / "chunks.jsonl",52    ]53    54    for path in local_paths:55        if path.exists():56            # Extract framework from parent directory or use "local"57            if "Local_saved_files" in str(path):58                framework = "crewai"59            elif "sample_code" in str(path):60                framework = "sample"61            else:62                framework = path.parent.name63            chunk_files.append((path, framework))64    65    # Check repository chunks66    repos_dir = PROJECT_ROOT / "data" / "processed" / "repos"67    if repos_dir.exists():68        for repo_dir in repos_dir.iterdir():69            if repo_dir.is_dir():70                for jsonl_file in repo_dir.glob("*_chunks.jsonl"):71                    # Extract framework from filename or directory72                    framework = jsonl_file.stem.replace("_chunks", "").split("_")[0]73                    chunk_files.append((jsonl_file, framework))74    75    return chunk_files76 77 78def merge_datasets(all_pairs: List[List[PositivePair]], 79                   all_triplets: List[List[Triplet]], 80                   output_dir: Path) -> None:81    """Merge all framework datasets into combined files (JSON + JSONL)."""82    output_dir.mkdir(parents=True, exist_ok=True)83    84    # Flatten lists85    combined_pairs = []86    for pairs in all_pairs:87        combined_pairs.extend(pairs)88    89    combined_triplets = []90    for triplets in all_triplets:91        combined_triplets.extend(triplets)92    93    # Export combined positive pairs - JSON94    pairs_json_path = output_dir / "positive_pairs.json"95    with open(pairs_json_path, "w", encoding="utf-8") as f:96        json.dump([asdict(p) for p in combined_pairs], f, indent=2, ensure_ascii=False)97    print(f"✅ Combined positive pairs (JSON): {pairs_json_path}")98    99    # Export combined positive pairs - JSONL100    pairs_jsonl_path = output_dir / "positive_pairs.jsonl"101    with open(pairs_jsonl_path, "w", encoding="utf-8") as f:102        for p in combined_pairs:103            f.write(json.dumps(asdict(p), ensure_ascii=False) + "\n")104    print(f"✅ Combined positive pairs (JSONL): {pairs_jsonl_path}")105    106    # Export combined triplets - JSON107    triplets_json_path = output_dir / "triplets.json"108    with open(triplets_json_path, "w", encoding="utf-8") as f:109        json.dump([asdict(t) for t in combined_triplets], f, indent=2, ensure_ascii=False)110    print(f"✅ Combined triplets (JSON): {triplets_json_path}")111    112    # Export combined triplets - JSONL113    triplets_jsonl_path = output_dir / "triplets.jsonl"114    with open(triplets_jsonl_path, "w", encoding="utf-8") as f:115        for t in combined_triplets:116            f.write(json.dumps(asdict(t), ensure_ascii=False) + "\n")117    print(f"✅ Combined triplets (JSONL): {triplets_jsonl_path}")118    119    return len(combined_pairs), len(combined_triplets)120 121 122def main():123    """Generate datasets for all discovered frameworks + combined dataset."""124    print("=" * 80)125    print("🚀 MULTI-FRAMEWORK TRAINING DATA GENERATOR")126    print("=" * 80)127    128    # Discover all chunk files129    print("\n🔍 Discovering chunk files...")130    chunk_files = discover_all_chunk_files()131    132    if not chunk_files:133        print("❌ No chunk files found!")134        print("\nPlease ensure chunks exist in:")135        print("  - data/processed/chunks/Local_saved_files/")136        print("  - data/processed/repos/*/")137        return138    139    print(f"✅ Found {len(chunk_files)} chunk file(s):\n")140    for path, framework in chunk_files:141        print(f"   📦 {framework}: {path.name}")142    143    # Process each framework144    print("\n" + "=" * 80)145    print("🔄 PROCESSING INDIVIDUAL FRAMEWORKS")146    print("=" * 80 + "\n")147    148    results = []149    all_pairs = []150    all_triplets = []151    152    for i, (chunks_path, framework) in enumerate(chunk_files, 1):153        print(f"\n[{i}/{len(chunk_files)}] Processing {framework.upper()}...")154        print("-" * 60)155        156        output_dir = PROJECT_ROOT / "data" / "processed" / f"training_{framework}"157        158        try:159            pairs, triplets = generate_pairs_and_triplets(160                chunks_path=chunks_path,161                output_dir=output_dir,162                num_pairs=100,163                num_triplets=100,164                variance=5,165                export_format="both"  # JSON + JSONL166            )167            168            # Collect for combined dataset169            all_pairs.append(pairs)170            all_triplets.append(triplets)171            172            results.append({173                "framework": framework,174                "status": "✅ SUCCESS",175                "pairs": len(pairs),176                "variations": sum(len(p.variations) for p in pairs),177                "triplets": len(triplets),178                "output": output_dir179            })180            181        except Exception as e:182            results.append({183                "framework": framework,184                "status": f"❌ FAILED: {str(e)}",185                "output": output_dir186            })187    188    # Create combined dataset189    print("\n" + "=" * 80)190    print("🔗 CREATING COMBINED DATASET (ALL FRAMEWORKS)")191    print("=" * 80 + "\n")192    193    combined_dir = PROJECT_ROOT / "data" / "processed" / "training_combined"194    total_pairs, total_triplets = merge_datasets(all_pairs, all_triplets, combined_dir)195    196    # Final summary197    print("\n" + "=" * 80)198    print("📊 FINAL SUMMARY")199    print("=" * 80 + "\n")200    201    print("INDIVIDUAL FRAMEWORK DATASETS:")202    print("-" * 40)203    for result in results:204        print(f"\n📦 {result['framework'].upper()}")205        print(f"   Status: {result['status']}")206        if "pairs" in result:207            print(f"   - positive_pairs.json: {result['pairs']} docs ({result['variations']} variations)")208            print(f"   - triplets.json: {result['triplets']} docs")209        print(f"   📁 {result['output']}")210    211    print("\n\nCOMBINED DATASET (ALL FRAMEWORKS):")212    print("-" * 40)213    print(f"📁 {combined_dir}")214    print(f"   - positive_pairs.json: {total_pairs} docs")215    print(f"   - triplets.json: {total_triplets} docs")216    217    # File count summary218    successful = sum(1 for r in results if "SUCCESS" in r["status"])219    total_files = (successful * 4) + 4  # 4 per framework + 4 combined220    221    print(f"\n\n📄 TOTAL FILES GENERATED: {total_files}")222    print(f"   - {successful} frameworks × 4 files = {successful * 4} files")223    print(f"   - Combined dataset = 4 files")224    print("=" * 80)225 226 227if __name__ == "__main__":228    main()229