MuratcanKoylan/Marketing-Memory-Routing-8B

1"""Quick test of balanced generation for underrepresented categories."""2 3import json4import os5from dotenv import load_dotenv6load_dotenv()7 8import cohere9 10client = cohere.ClientV2(api_key=os.getenv("COHERE_API_KEY"))11 12# Test with underrepresented categories13test_categories = ["company.tools_config", "company.knowledge_artifacts", "none"]14 15for category in test_categories:16    print(f"\n{'='*60}")17    print(f"Testing: {category}")18    print("="*60)19    20    if category == "none":21        prompt = """Generate a marketing conversation that has NO long-term memory value.22 23The conversation should be transactional, vague, or temporary.24Examples: checking status, scheduling, confirming receipt.25 26Generate 4 turns. Start mid-conversation (no greetings).27 28OUTPUT FORMAT (JSON only):29{30  "scenario_id": "none_001",31  "conversation": [32    {"role": "user", "content": "..."},33    {"role": "assistant", "content": "..."}34  ],35  "labels": {36    "categories": ["none"],37    "rationale": "..."38  }39}"""40    else:41        prompt = f"""Generate a marketing conversation that clearly demonstrates: {category}42 43The conversation MUST contain clear signals for this category.444-6 turns, start mid-conversation (no greetings).45 46CRITICAL: The categories array MUST include "{category}".47 48OUTPUT FORMAT (JSON only):49{{50  "scenario_id": "{category.replace('.', '_')}_001",51  "conversation": [52    {{"role": "user", "content": "..."}},53    {{"role": "assistant", "content": "..."}}54  ],55  "labels": {{56    "categories": ["{category}"],57    "rationale": "..."58  }}59}}"""60 61    try:62        response = client.chat(63            messages=[{"role": "user", "content": prompt}],64            temperature=0.7,65            model="command-r-plus-08-2024",66            response_format={"type": "json_object"}67        )68        69        content = response.message.content[0].text70        data = json.loads(content)71        72        output_cats = data.get("labels", {}).get("categories", [])73        print(f"Target: {category}")74        print(f"Output: {output_cats}")75        print(f"Match: {'YES' if category in output_cats else 'NO'}")76        77        if data.get("conversation"):78            print(f"First turn: {data['conversation'][0]['content'][:80]}...")79    except Exception as e:80        print(f"Error: {e}")81 82