MuratcanKoylan/Marketing-Memory-Routing-8B

1"""2Balanced Dataset Generation with Concurrent API Calls3 4Generates 10 items simultaneously per batch for faster generation.5"""6 7import json8import random9import time10import sys11import asyncio12import os13from typing import List, Dict, Any, Optional14from datetime import datetime15from concurrent.futures import ThreadPoolExecutor16import cohere17from dotenv import load_dotenv18 19load_dotenv()20 21# Target counts per category (balanced)22CATEGORY_TARGETS = {23    "company.brand_core": 77,24    "company.strategic_signatures": 77,25    "company.knowledge_artifacts": 77,26    "company.business_priorities": 77,27    "company.tools_config": 77,28    "company.performance_context": 77,29    "user.communication_style": 77,30    "user.strategic_approach": 77,31    "user.role_context": 77,32    "user.workflow_patterns": 77,33    "user.session_history": 77,34    "user.interaction_preferences": 77,35    "none": 77,36}37 38CATEGORY_EXAMPLES = {39    "company.brand_core": {40        "signals": ["brand voice is warm", "primary color is #2563EB", "never use jargon", "tagline is..."],41    },42    "company.strategic_signatures": {43        "signals": ["always prioritize retention", "80/20 rule", "never launch without testing"],44    },45    "company.knowledge_artifacts": {46        "signals": ["style guide says", "playbook recommends", "SOP for launches", "template includes"],47    },48    "company.business_priorities": {49        "signals": ["Q4 focus is", "this quarter's target", "holiday campaign", "prioritizing APAC"],50    },51    "company.tools_config": {52        "signals": ["Slack webhook URL", "HubSpot sync", "API key is", "Zapier integration"],53    },54    "company.performance_context": {55        "signals": ["24% open rate", "CTR improved by", "retrospective showed", "conversion dropped"],56    },57    "user.communication_style": {58        "signals": ["prefer bullet points", "keep it under 200 words", "casual tone", "data-driven"],59    },60    "user.strategic_approach": {61        "signals": ["prioritize speed over perfection", "test fast fail fast", "customer feedback"],62    },63    "user.role_context": {64        "signals": ["As VP of Marketing", "report to CMO", "budget authority up to", "manage team of"],65    },66    "user.workflow_patterns": {67        "signals": ["review drafts Monday", "don't send Friday", "async via Slack", "weekly sync Tuesday"],68    },69    "user.session_history": {70        "signals": ["as we discussed yesterday", "continuing from last", "proposal we started"],71    },72    "user.interaction_preferences": {73        "signals": ["push back on my ideas", "give me options", "be direct", "ask clarifying questions"],74    },75    "none": {76        "signals": ["what time is meeting", "checking status", "confirming receipt", "quick question"],77    },78}79 80 81class BalancedAsyncGenerator:82    def __init__(self):83        self.api_key = os.getenv("COHERE_API_KEY")84        if not self.api_key:85            raise ValueError("COHERE_API_KEY not found")86        self.client = cohere.ClientV2(api_key=self.api_key)87        self.model = "command-r-plus-08-2024"88        self.executor = ThreadPoolExecutor(max_workers=10)89    90    def _extract_text(self, response) -> Optional[str]:91        if not response or not getattr(response, "message", None):92            return None93        blocks = getattr(response.message, "content", []) or []94        for block in blocks:95            text = getattr(block, "text", None)96            if isinstance(text, str) and text.strip():97                return text98        return None99    100    def _generate_sync(self, category: str) -> Optional[Dict]:101        """Synchronous generation for a single category."""102        signals = CATEGORY_EXAMPLES.get(category, {}).get("signals", [])103        signals_text = "\n".join(f"- {s}" for s in signals[:4])104        105        if category == "none":106            prompt = f"""Generate a marketing conversation with NO long-term memory value.107Transactional, vague, or temporary only. Examples: status check, scheduling, confirming.1084-6 turns, no greetings, start mid-conversation.109 110OUTPUT (JSON only):111{{"scenario_id": "none_{random.randint(100,999)}", "conversation": [{{"role": "user", "content": "..."}}, {{"role": "assistant", "content": "..."}}], "labels": {{"categories": ["none"], "persistence_horizon": "short", "memory_scope": "none", "rationale": "..."}}, "metadata": {{"primary_category": "none", "turn_count": 4}}}}"""112        else:113            prompt = f"""Generate a marketing conversation demonstrating: {category}114 115SIGNALS FOR THIS CATEGORY:116{signals_text}117 118REQUIREMENTS:1191. MUST contain clear signals for {category}1202. 4-6 turns, no greetings, start mid-conversation1213. Include specific details (names, numbers, dates)122 123CRITICAL: categories array MUST include "{category}"124 125OUTPUT (JSON only):126{{"scenario_id": "{category.replace('.', '_')}_{random.randint(100,999)}", "conversation": [{{"role": "user", "content": "..."}}, {{"role": "assistant", "content": "..."}}], "labels": {{"categories": ["{category}"], "persistence_horizon": "long", "memory_scope": "company", "rationale": "..."}}, "metadata": {{"primary_category": "{category}", "turn_count": 4}}}}"""127 128        try:129            response = self.client.chat(130                messages=[{"role": "user", "content": prompt}],131                temperature=0.7,132                model=self.model,133                response_format={"type": "json_object"}134            )135            136            content = self._extract_text(response)137            if not content:138                return None139            140            if content.startswith("```json"):141                content = content[7:]142            if content.endswith("```"):143                content = content[:-3]144            145            data = json.loads(content.strip())146            147            # Validate target category is present148            categories = data.get("labels", {}).get("categories", [])149            if category.lower() not in [c.lower() for c in categories]:150                return None151            152            # Clean: Remove "none" if other categories exist153            if len(categories) > 1 and "none" in [c.lower() for c in categories]:154                data["labels"]["categories"] = [c for c in categories if c.lower() != "none"]155            156            return data157            158        except Exception as e:159            return None160    161    async def generate_batch(self, categories: List[str]) -> List[Dict]:162        """Generate a batch of items concurrently."""163        loop = asyncio.get_event_loop()164        tasks = [165            loop.run_in_executor(self.executor, self._generate_sync, cat)166            for cat in categories167        ]168        results = await asyncio.gather(*tasks, return_exceptions=True)169        return [r for r in results if isinstance(r, dict)]170 171 172async def run_balanced_generation_async():173    """Run balanced generation with concurrent batches."""174    175    generator = BalancedAsyncGenerator()176    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")177    output_file = f"synthetic_data/balanced_dataset_{timestamp}.jsonl"178    179    # Track progress per category180    category_counts = {cat: 0 for cat in CATEGORY_TARGETS}181    all_data = []182    183    print("=" * 70, flush=True)184    print("BALANCED CONCURRENT DATASET GENERATION", flush=True)185    print("=" * 70, flush=True)186    print(f"Target per category: 77", flush=True)187    print(f"Total categories: {len(CATEGORY_TARGETS)}", flush=True)188    print(f"Expected total: {77 * len(CATEGORY_TARGETS)}", flush=True)189    print(f"Batch size: 10 concurrent requests", flush=True)190    print(flush=True)191    192    batch_num = 0193    194    while True:195        # Find categories that still need examples196        needed = []197        for cat, target in CATEGORY_TARGETS.items():198            remaining = target - category_counts[cat]199            needed.extend([cat] * min(remaining, 2))  # Up to 2 per category per batch200        201        if not needed:202            break203        204        # Take up to 10 for this batch205        batch_categories = needed[:10]206        batch_num += 1207        208        print(f"\n[Batch {batch_num}] Generating {len(batch_categories)} items...", flush=True)209        210        results = await generator.generate_batch(batch_categories)211        212        # Process results213        for result in results:214            if result:215                primary = result.get("metadata", {}).get("primary_category") or \216                         result.get("labels", {}).get("categories", ["unknown"])[0]217                218                if primary in category_counts:219                    category_counts[primary] += 1220                    all_data.append(result)221                    222                    # Save incrementally223                    with open(output_file, "a") as f:224                        f.write(json.dumps(result) + "\n")225        226        # Progress report227        total_done = sum(category_counts.values())228        total_target = sum(CATEGORY_TARGETS.values())229        print(f"  Success: {len(results)}/{len(batch_categories)} | Total: {total_done}/{total_target}", flush=True)230        231        # Show category progress every 10 batches232        if batch_num % 10 == 0:233            print("\n  Category Progress:", flush=True)234            for cat, count in sorted(category_counts.items()):235                target = CATEGORY_TARGETS[cat]236                bar = "█" * (count * 20 // target) + "░" * (20 - count * 20 // target)237                print(f"    {cat:<35} [{bar}] {count}/{target}", flush=True)238        239        # Rate limit: wait 3 seconds between batches240        await asyncio.sleep(3)241    242    # Final summary243    print("\n" + "=" * 70, flush=True)244    print("GENERATION COMPLETE", flush=True)245    print("=" * 70, flush=True)246    print(f"\nFinal Distribution:", flush=True)247    for cat, count in sorted(category_counts.items(), key=lambda x: -x[1]):248        pct = count / len(all_data) * 100 if all_data else 0249        print(f"  {cat:<40} {count:>4} ({pct:.1f}%)", flush=True)250    251    print(f"\nTotal examples: {len(all_data)}", flush=True)252    print(f"Output file: {output_file}", flush=True)253    254    return output_file255 256 257if __name__ == "__main__":258    asyncio.run(run_balanced_generation_async())259 260