MuratcanKoylan/Marketing-Memory-Routing-8B

1"""2Balanced Dataset Generation Script3 4This script generates a balanced training dataset with:51. STRICT category enforcement - the model MUST output the target category62. Equal distribution across all categories73. Improved prompts for underrepresented categories8"""9 10import json11import random12import time13import sys14import asyncio15import os16from typing import List, Dict, Any, Optional17from datetime import datetime18import cohere19from dotenv import load_dotenv20 21load_dotenv()22 23# BALANCED DISTRIBUTION - Equal weight for all categories24BALANCED_DISTRIBUTION = {25    "company.brand_core": 80,26    "company.strategic_signatures": 80,27    "company.knowledge_artifacts": 80,28    "company.business_priorities": 80,29    "company.tools_config": 80,30    "company.performance_context": 80,31    "user.communication_style": 80,32    "user.strategic_approach": 80,33    "user.role_context": 80,34    "user.workflow_patterns": 80,35    "user.session_history": 80,36    "user.interaction_preferences": 80,37    "none": 80,38}39 40# Category-specific examples and signals for better generation41CATEGORY_EXAMPLES = {42    "company.brand_core": {43        "description": "Brand voice, values, positioning, visual identity, tone guidelines",44        "example_signals": [45            "Our brand voice is warm and conversational",46            "We always use sentence case for headlines",47            "Our primary color is #2563EB",48            "We never use corporate jargon",49            "Our tagline is 'Simplify Everything'"50        ],51        "example_conversation": "USER: Remember, our brand personality is 'friendly expert' - knowledgeable but approachable."52    },53    "company.strategic_signatures": {54        "description": "Decision frameworks, strategic heuristics, recurring patterns in how the company operates",55        "example_signals": [56            "We always prioritize retention over acquisition",57            "Our 80/20 rule: 80% proven tactics, 20% experiments",58            "We never launch without A/B testing",59            "Customer lifetime value drives all decisions"60        ],61        "example_conversation": "USER: Our strategic principle is 'land and expand' - start small with enterprises then grow."62    },63    "company.knowledge_artifacts": {64        "description": "Style guides, playbooks, SOPs, documented processes, templates",65        "example_signals": [66            "Here's our content style guide",67            "The campaign playbook says...",68            "According to our SOP for launches",69            "Our template for proposals includes..."70        ],71        "example_conversation": "USER: I'm attaching our updated brand guidelines PDF. Make sure all content follows section 3.2."72    },73    "company.business_priorities": {74        "description": "Quarterly goals, seasonal campaigns, current OKRs, active initiatives",75        "example_signals": [76            "Q4 focus is enterprise expansion",77            "This quarter's target is 500 MQLs",78            "Holiday campaign launches December 1st",79            "We're prioritizing APAC market this quarter"80        ],81        "example_conversation": "USER: For Q1, we're shifting focus entirely to the SMB segment. All campaigns should target companies under 100 employees."82    },83    "company.tools_config": {84        "description": "Integrations, API keys, workflow settings, tool configurations",85        "example_signals": [86            "The Slack webhook URL is...",87            "Configure HubSpot to sync with...",88            "The API key for analytics is...",89            "Set up the Zapier integration to..."90        ],91        "example_conversation": "USER: Here's the API key for our analytics dashboard: sk-xxx-123. Make sure it syncs every 6 hours."92    },93    "company.performance_context": {94        "description": "Campaign metrics, retrospectives, learnings, performance data",95        "example_signals": [96            "Last campaign had 24% open rate",97            "CTR improved by 15% after the redesign",98            "The retrospective showed we need more testing",99            "Conversion rate dropped after the price change"100        ],101        "example_conversation": "USER: The email campaign results are in: 28% open rate, 4.2% CTR. That's our best performance this year."102    },103    "user.communication_style": {104        "description": "Preferred tone, verbosity, format expectations, writing style",105        "example_signals": [106            "I prefer bullet points over paragraphs",107            "Keep responses under 200 words",108            "Use casual, friendly tone with me",109            "I like data-driven explanations"110        ],111        "example_conversation": "USER: Just so you know, I prefer concise bullet points. No need for lengthy explanations with me."112    },113    "user.strategic_approach": {114        "description": "Personal priorities, success definitions, decision-making style",115        "example_signals": [116            "I always prioritize speed over perfection",117            "My philosophy is test fast, fail fast",118            "I measure success by customer feedback",119            "I believe in data-driven decisions only"120        ],121        "example_conversation": "USER: My approach is always 'done is better than perfect'. I'd rather ship and iterate."122    },123    "user.role_context": {124        "description": "Title, scope, decision authority, reporting structure",125        "example_signals": [126            "As VP of Marketing, I approve all campaigns",127            "I report directly to the CMO",128            "My budget authority is up to $50k",129            "I manage a team of 12 marketers"130        ],131        "example_conversation": "USER: Just for context, I'm the Director of Growth and I have final say on all acquisition campaigns."132    },133    "user.workflow_patterns": {134        "description": "Review cadence, collaboration norms, meeting schedules",135        "example_signals": [136            "I review drafts every Monday morning",137            "Don't send me anything on Fridays",138            "I prefer async communication via Slack",139            "Weekly sync is Tuesdays at 2pm"140        ],141        "example_conversation": "USER: My review schedule is Monday mornings only. Anything sent Friday won't be seen until next week."142    },143    "user.session_history": {144        "description": "Immediate context, recent asks, current working session",145        "example_signals": [146            "As we discussed yesterday...",147            "Continuing from our last conversation",148            "The proposal we started earlier",149            "Following up on the draft you sent"150        ],151        "example_conversation": "USER: Let's pick up where we left off yesterday on the Johnson account proposal."152    },153    "user.interaction_preferences": {154        "description": "Coaching style, feedback expectations, collaboration preferences",155        "example_signals": [156            "I want you to push back on my ideas",157            "Give me options, not just one answer",158            "Be direct with feedback, don't sugarcoat",159            "I prefer you ask clarifying questions"160        ],161        "example_conversation": "USER: I want you to challenge my assumptions. If you think I'm wrong, tell me directly."162    },163    "none": {164        "description": "Transactional, vague, or temporary content with no memory value",165        "example_signals": [166            "What time is the meeting?",167            "Can you check the status?",168            "Just confirming receipt",169            "Quick question about the attachment"170        ],171        "example_conversation": "USER: Hey, what's the status on that thing we discussed? Just checking in."172    }173}174 175class BalancedDataGenerator:176    def __init__(self, api_key: Optional[str] = None):177        self.api_key = api_key or os.getenv("COHERE_API_KEY")178        if not self.api_key:179            raise ValueError("COHERE_API_KEY not found")180        self.client = cohere.ClientV2(api_key=self.api_key)181        self.model = "command-r-plus-08-2024"182    183    def _extract_text(self, response) -> Optional[str]:184        if not response or not getattr(response, "message", None):185            return None186        blocks = getattr(response.message, "content", []) or []187        for block in blocks:188            text = getattr(block, "text", None)189            if isinstance(text, str) and text.strip():190                return text191        return None192    193    def generate_for_category(self, category: str, max_retries: int = 3) -> Optional[Dict]:194        """Generate a conversation that MUST contain the specified category."""195        196        cat_info = CATEGORY_EXAMPLES.get(category, {})197        description = cat_info.get("description", category)198        example_signals = cat_info.get("example_signals", [])199        example_conv = cat_info.get("example_conversation", "")200        201        # Build a very specific prompt202        if category == "none":203            prompt = f"""Generate a realistic marketing conversation that has NO long-term memory value.204 205The conversation should be:206- Transactional (checking status, scheduling, confirming)207- Vague or generic (no specific details worth remembering)208- Temporary (only relevant for this moment)209 210Examples of "none" conversations:211- "What time is the meeting tomorrow?"212- "Just confirming you received the file"213- "Quick status check on the project"214- "Can you resend that link?"215 216Generate a 4-6 turn conversation between USER and ASSISTANT.217Start mid-conversation (no greetings).218 219OUTPUT FORMAT (JSON only):220{{221  "scenario_id": "none_{random.randint(100,999)}",222  "conversation": [223    {{"role": "user", "content": "..."}},224    {{"role": "assistant", "content": "..."}}225  ],226  "labels": {{227    "categories": ["none"],228    "persistence_horizon": "short",229    "memory_scope": "none",230    "rationale": "This conversation is transactional/temporary with no memory value"231  }},232  "metadata": {{233    "primary_category": "none",234    "turn_count": 4235  }}236}}"""237        else:238            prompt = f"""Generate a marketing conversation that clearly demonstrates the category: {category}239 240CATEGORY DEFINITION:241{description}242 243SIGNALS THAT INDICATE THIS CATEGORY:244{chr(10).join(f"- {s}" for s in example_signals[:4])}245 246EXAMPLE UTTERANCE:247{example_conv}248 249REQUIREMENTS:2501. The conversation MUST contain clear signals for {category}2512. The USER should explicitly state information that maps to this category2523. Make it natural and realistic - embed the signals organically2534. 4-6 turns, start mid-conversation (no greetings)2545. Include specific, concrete details (names, numbers, dates)255 256CRITICAL: The output categories array MUST include "{category}" as the primary category.257You may include 1 additional category if naturally present, but {category} MUST be there.258 259OUTPUT FORMAT (JSON only):260{{261  "scenario_id": "{category.replace('.', '_')}_{random.randint(100,999)}",262  "conversation": [263    {{"role": "user", "content": "..."}},264    {{"role": "assistant", "content": "..."}}265  ],266  "labels": {{267    "categories": ["{category}"],268    "persistence_horizon": "long|medium|short",269    "memory_scope": "company|user",270    "rationale": "Explanation of why {category} applies"271  }},272  "metadata": {{273    "primary_category": "{category}",274    "turn_count": 4275  }}276}}"""277 278        for attempt in range(max_retries):279            try:280                response = self.client.chat(281                    messages=[{"role": "user", "content": prompt}],282                    temperature=0.7,283                    model=self.model,284                    response_format={"type": "json_object"}285                )286                287                content = self._extract_text(response)288                if not content:289                    continue290                291                # Clean JSON292                if content.startswith("```json"):293                    content = content[7:]294                if content.endswith("```"):295                    content = content[:-3]296                297                data = json.loads(content.strip())298                299                # VALIDATE: Ensure target category is present300                categories = data.get("labels", {}).get("categories", [])301                if category.lower() not in [c.lower() for c in categories]:302                    print(f"  Warning: Target {category} not in output {categories}. Retrying...")303                    continue304                305                # Clean: Remove "none" if other categories exist306                if len(categories) > 1 and "none" in [c.lower() for c in categories]:307                    data["labels"]["categories"] = [c for c in categories if c.lower() != "none"]308                309                return data310                311            except Exception as e:312                print(f"  Attempt {attempt+1} failed: {e}")313                time.sleep(5 * (attempt + 1))314        315        return None316 317 318async def generate_balanced_dataset(output_dir: str = "synthetic_data", target_per_category: int = 80):319    """Generate a balanced dataset with equal examples per category."""320    321    os.makedirs(output_dir, exist_ok=True)322    generator = BalancedDataGenerator()323    324    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")325    output_file = f"{output_dir}/balanced_dataset_{timestamp}.jsonl"326    log_file = f"{output_dir}/balanced_generation_log_{timestamp}.txt"327    328    all_data = []329    category_counts = {cat: 0 for cat in BALANCED_DISTRIBUTION.keys()}330    331    print("=" * 70, flush=True)332    print("BALANCED DATASET GENERATION", flush=True)333    print("=" * 70, flush=True)334    print(f"Target per category: {target_per_category}", flush=True)335    print(f"Total categories: {len(BALANCED_DISTRIBUTION)}", flush=True)336    print(f"Expected total: {target_per_category * len(BALANCED_DISTRIBUTION)}", flush=True)337    print(flush=True)338    339    with open(log_file, "w") as log:340        log.write(f"Balanced Generation Started: {timestamp}\n")341        log.write(f"Target per category: {target_per_category}\n\n")342        343        for category in BALANCED_DISTRIBUTION.keys():344            print(f"\n--- Generating {target_per_category} examples for: {category} ---", flush=True)345            log.write(f"\n=== {category} ===\n")346            log.flush()347            348            for i in range(target_per_category):349                result = generator.generate_for_category(category)350                351                if result:352                    all_data.append(result)353                    category_counts[category] += 1354                    355                    # Save incrementally356                    with open(output_file, "a") as f:357                        f.write(json.dumps(result) + "\n")358                    359                    if (i + 1) % 10 == 0:360                        print(f"  Progress: {i+1}/{target_per_category}", flush=True)361                        log.write(f"  {i+1}/{target_per_category} complete\n")362                        log.flush()363                else:364                    print(f"  Failed: {i+1}", flush=True)365                    log.write(f"  Failed to generate example {i+1}\n")366                    log.flush()367                368                # Rate limiting369                await asyncio.sleep(0.5)370            371            print(f"  Completed: {category_counts[category]}/{target_per_category}", flush=True)372    373    # Final summary374    print("\n" + "=" * 70)375    print("GENERATION COMPLETE")376    print("=" * 70)377    print(f"\nCategory Distribution:")378    for cat, count in sorted(category_counts.items(), key=lambda x: -x[1]):379        pct = count / len(all_data) * 100 if all_data else 0380        print(f"  {cat:<40} {count:>4} ({pct:.1f}%)")381    382    print(f"\nTotal examples: {len(all_data)}")383    print(f"Output file: {output_file}")384    385    return output_file386 387 388if __name__ == "__main__":389    target = int(sys.argv[1]) if len(sys.argv) > 1 else 80390    asyncio.run(generate_balanced_dataset(target_per_category=target))391 392