MuratcanKoylan/Marketing-Memory-Routing-8B

1import json2import random3import time4import sys5from typing import List, Dict, Any6from synthetic_data.pipeline import SyntheticDataPipeline7from synthetic_data.validate import validate_synthetic_data8 9CATEGORY_DISTRIBUTION = {10    "company.brand_core": 0.10,11    "company.strategic_signatures": 0.08,12    "company.knowledge_artifacts": 0.08,13    "company.business_priorities": 0.10,14    "company.tools_config": 0.07,15    "company.performance_context": 0.09,16    "user.communication_style": 0.10,17    "user.strategic_approach": 0.09,18    "user.role_context": 0.07,19    "user.workflow_patterns": 0.08,20    "user.session_history": 0.06,21    "user.interaction_preferences": 0.08,22    "none": 0.1023}24 25def run_pipeline_batches(total_items: int = 100, batch_size: int = 10):26    pipeline = SyntheticDataPipeline()27    categories = list(CATEGORY_DISTRIBUTION.keys())28    weights = list(CATEGORY_DISTRIBUTION.values())29    30    all_data = []31    num_batches = max(1, total_items // batch_size)32    33    print(f"Starting generation of {total_items} items in {num_batches} batches (Size: {batch_size})...")34 35    for batch_num in range(1, num_batches + 1):36        print(f"\n=== Processing Batch {batch_num}/{num_batches} ===")37        batch_data = []38        39        while len(batch_data) < batch_size:40            category = random.choices(categories, weights=weights, k=1)[0]41            current_count = len(batch_data) + 142            print(f"  Generating item {current_count}/{batch_size} (Category: {category})...")43            44            # Determine if we should add a distractor (30% chance)45            distractor = None46            if random.random() < 0.30 and category != "none":47                 possible_distractors = [c for c in categories if c != category and c != "none"]48                 if possible_distractors:49                     distractor = random.choice(possible_distractors)50 51            persistence = _get_persistence_for_category(category)52            turns = random.randint(4, 10)53            54            scenario = pipeline.generate_scenario_spec(55                category=category,56                distractor=distractor,57                persistence=persistence,58                turns=turns59            )60            61            if not scenario:62                print(f"    Failed to generate scenario for {category}. Retrying...")63                time.sleep(20)64                continue65                66            conversation = pipeline.generate_conversation(scenario, turn_count=turns)67            68            if conversation:69                batch_data.append(conversation)70                print(f"    Generated: {conversation.get('scenario_id', 'Unknown ID')}")71            else:72                 print(f"    Failed to generate conversation for {category}. Retrying...")73                 time.sleep(20)74                 continue75            76            print("    Sleeping for 15s to avoid rate limits...")77            time.sleep(15)78        79        # Save batch80        batch_filename = f"synthetic_data/batch_{batch_num:02d}.json"81        with open(batch_filename, "w") as f:82            json.dump(batch_data, f, indent=2)83        print(f"  Saved batch to {batch_filename}")84        85        # Validate batch86        print("  Validating batch...")87        metrics = validate_synthetic_data(batch_filename)88        print(json.dumps(metrics, indent=2))89        90        all_data.extend(batch_data)91        92    # Save all data93    with open("synthetic_data/all_generated_data_100.json", "w") as f:94        json.dump(all_data, f, indent=2)95    print(f"\nCompleted. Total items generated: {len(all_data)}")96    print("Full dataset saved to synthetic_data/all_generated_data_100.json")97 98def _get_persistence_for_category(category: str) -> str:99    if "brand_core" in category or "strategic_signatures" in category or "knowledge_artifacts" in category or "communication_style" in category or "strategic_approach" in category:100        return "long"101    elif "tools_config" in category or "role_context" in category or "workflow_patterns" in category:102        return "medium"103    elif "business_priorities" in category or "session_history" in category:104        return "short"105    elif "performance_context" in category:106        return "rolling"107    elif "interaction_preferences" in category:108        return "evolving"109    elif "none" in category:110        return "short"111    return "medium" 112 113if __name__ == "__main__":114    total = int(sys.argv[1]) if len(sys.argv) > 1 else 100115    batch = int(sys.argv[2]) if len(sys.argv) > 2 else 10116    run_pipeline_batches(total, batch)117