training/quick_test.py

"""
Quick test of the RL model.
"""

import asyncio
import json
import os
from dotenv import load_dotenv

load_dotenv()

import tinker
from tinker import types
from tinker_cookbook import renderers
from tinker_cookbook.tokenizer_utils import get_tokenizer

BASE_MODEL = "meta-llama/Llama-3.1-8B"

# Run with both SFT and RL (most iterations)
RL_CHECKPOINT = "tinker://398393e1-7182-555d-aa1b-7ddf23892338:train:0/sampler_weights/rl_iter_005"

# SFT from the same run
SFT_CHECKPOINT = "tinker://398393e1-7182-555d-aa1b-7ddf23892338:train:0/sampler_weights/sft_final_sampler"

VALID_CATEGORIES = {
    "company.brand_core", "company.strategic_signatures", "company.knowledge_artifacts",
    "company.business_priorities", "company.tools_config", "company.performance_context",
    "user.communication_style", "user.strategic_approach", "user.role_context",
    "user.workflow_patterns", "user.session_history", "user.interaction_preferences",
    "none"
}


SYSTEM_PROMPT = """You route marketing conversations into structured memory categories.

Available categories:
- company.brand_core: Voice, values, positioning, identity anchors (Long >1y)
- company.strategic_signatures: Decision frameworks, strategic heuristics (Long >1y)
- company.knowledge_artifacts: Docs, style guides, playbooks (Long >1y)
- company.business_priorities: Quarterly/seasonal goals, active campaigns (Short <3m)
- company.tools_config: Integrations, API keys, workflow settings (Medium ~6m)
- company.performance_context: Campaign metrics, retrospectives, learnings (Rolling ~6m)
- user.communication_style: Tone, verbosity, format expectations (Long >1y)
- user.strategic_approach: Personal priorities, success definitions (Long >1y)
- user.role_context: Title, scope, decision authority (Medium ~1y)
- user.workflow_patterns: Review cadence, collaboration norms (Medium ~1y)
- user.session_history: Immediate context, recent asks (Short <2w)
- user.interaction_preferences: Coaching style, feedback expectations (Evolving)
- none: Irrelevant, vague, or transactional content

Respond with comma-separated categories. Use 'none' only if no other category applies."""


async def test_model(checkpoint: str, name: str, test_examples: list):
    """Test a model on examples."""
    print(f"\n{'='*60}")
    print(f"TESTING: {name}")
    print(f"Checkpoint: {checkpoint}")
    print(f"{'='*60}")
    
    service_client = tinker.ServiceClient()
    tokenizer = get_tokenizer(BASE_MODEL)
    renderer = renderers.get_renderer(name="llama3", tokenizer=tokenizer)
    
    sampling_client = service_client.create_sampling_client(model_path=checkpoint)
    stop_sequences = renderer.get_stop_sequences()
    
    results = []
    
    for i, example in enumerate(test_examples):
        messages = example.get("messages", [])
        gold = example.get("categories", [])
        
        # Build prompt with system message (matching training format)
        conversation_text = ""
        for m in messages:
            role = m["role"].upper()
            conversation_text += f"{role}: {m['content']}\n"
        
        prompt_messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": f"Conversation:\n{conversation_text}"}
        ]
        
        prompt = renderer.build_generation_prompt(prompt_messages)
        params = types.SamplingParams(max_tokens=100, temperature=0.1, stop=stop_sequences)
        
        result = sampling_client.sample(prompt=prompt, sampling_params=params, num_samples=1).result()
        response, success = renderer.parse_response(result.sequences[0].tokens)
        predicted = response["content"] if success else ""
        
        # Parse prediction
        predicted_set = set([c.strip().lower() for c in predicted.split(",") 
                           if c.strip().lower() in VALID_CATEGORIES])
        gold_set = set([c.lower() for c in gold])
        
        any_match = len(predicted_set & gold_set) > 0 if gold_set else (len(predicted_set) == 0)
        exact_match = predicted_set == gold_set
        
        results.append({
            "any_match": any_match,
            "exact_match": exact_match,
            "predicted": predicted,
            "gold": gold
        })
        
        # Show first 5 examples
        if i < 5:
            print(f"\nExample {i+1}:")
            print(f"  Gold: {gold}")
            print(f"  Pred: {predicted}")
            print(f"  Match: {'Yes' if any_match else 'No'}")
    
    # Summary
    any_match_rate = sum(r["any_match"] for r in results) / len(results) if results else 0
    exact_match_rate = sum(r["exact_match"] for r in results) / len(results) if results else 0
    
    print(f"\n--- Results ({len(results)} examples) ---")
    print(f"Any Match:   {any_match_rate:.1%}")
    print(f"Exact Match: {exact_match_rate:.1%}")
    
    return {"any_match": any_match_rate, "exact_match": exact_match_rate}


async def main():
    # First, preprocess data
    print("=" * 60)
    print("LOADING TEST DATA")
    print("=" * 60)
    
    data = []
    with open("synthetic_data/training_dataset_1000.jsonl", "r") as f:
        for line in f:
            item = json.loads(line)
            messages = []
            for turn in item.get("conversation", []):
                if isinstance(turn, dict):
                    messages.append({"role": turn["role"], "content": turn["content"]})
            
            # Extract categories - handle nested labels structure
            labels = item.get("labels", {})
            if isinstance(labels, dict):
                categories = labels.get("categories", [])
            elif isinstance(labels, list):
                categories = labels
            else:
                categories = []
            
            if not categories:
                # Parse from scenario_id
                scenario_id = item.get("scenario_id", "")
                if "." in scenario_id:
                    cat = scenario_id.split("_")[0]
                    categories = [cat]
            
            data.append({
                "messages": messages,
                "categories": categories
            })
    
    print(f"Total examples: {len(data)}")
    
    # Use last 50 as test
    test_data = data[-50:]
    print(f"Test examples: {len(test_data)}")
    
    # Test RL model
    rl_results = await test_model(RL_CHECKPOINT, "RL Model (5 iters)", test_data)
    
    # Test SFT model for comparison
    sft_results = await test_model(SFT_CHECKPOINT, "SFT Model", test_data)
    
    print("\n" + "=" * 60)
    print("COMPARISON")
    print("=" * 60)
    print(f"SFT Any Match:   {sft_results['any_match']:.1%}")
    print(f"RL Any Match:    {rl_results['any_match']:.1%}")
    print(f"Improvement:     {(rl_results['any_match'] - sft_results['any_match'])*100:+.1f}pp")


if __name__ == "__main__":
    asyncio.run(main())