Initial commit (clean history)HEADmain

1 files changed, 607 insertions, 0 deletions

diff --git a/scripts/pilot_runner_v1.py b/scripts/pilot_runner_v1.py
new file mode 100644
index 0000000..fbb2876
--- /dev/null
+++ b/scripts/pilot_runner_v1.py
@@ -0,0 +1,607 @@
+#!/usr/bin/env python3
+"""
+Pilot Runner v1 - Style-Aware Judge + Gating Logic
+
+Upgrade from v0:
+- StylePrefs: User style preferences (length, bullets, language)
+- style_judge: Checks style conformance, not just non-empty
+- compute_feedback_for_turn: gating=1 only for preference-related turns
+- Extended queries: ~10 turns with preference/task mix
+
+Goal: Verify that:
+1. sat_t varies based on style violations (not always 1)
+2. gating=1 only on preference turns, 0 on regular tasks
+3. RL updates happen when gating=1 and reward != baseline
+4. Over turns, model may adapt to preferences (sat_t improves)
+"""
+
+import sys
+import os
+import json
+from datetime import datetime
+from dataclasses import dataclass, asdict, field
+from typing import List, Dict, Any, Optional, Tuple
+
+# Add src to path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../src"))
+
+from personalization.serving import PersonalizedLLM, Feedback, AssistantResponse
+
+
+# =============================================================================
+# Style Preferences
+# =============================================================================
+
+@dataclass
+class StylePrefs:
+    """User's style preferences for the judge to check."""
+    require_short: bool = False
+    max_chars: int = 300
+    require_bullets: bool = False
+    lang: str = "en"  # "en" or "zh"
+
+
+# =============================================================================
+# Style-Aware Judge
+# =============================================================================
+
+@dataclass
+class JudgeResult:
+    """Output from the judge for one turn."""
+    sat_t: float        # Satisfaction score [0, 1]
+    sev_t: float        # Severity of violations [0, 1]
+    prog_t: float       # Task progress [0, 1]
+    violations: List[str]  # List of violated constraints
+
+
+def style_judge(
+    query: str,
+    answer: str,
+    task_type: str,
+    prefs: StylePrefs,
+) -> JudgeResult:
+    """
+    Style-aware judge that checks:
+    - Empty/too short answer
+    - Length constraint (max_chars)
+    - Bullet point requirement
+    - Language preference
+    - Code block for code tasks
+    
+    Returns:
+        JudgeResult with sat_t, sev_t, prog_t, and violations list.
+    """
+    violations: List[str] = []
+    text = (answer or "").strip()
+
+    # 0) Empty answer - immediate fail
+    if not text or len(text) < 5:
+        violations.append("empty_answer")
+        return JudgeResult(
+            sat_t=0.0,
+            sev_t=1.0,
+            prog_t=0.0,
+            violations=violations,
+        )
+
+    # 1) Length preference
+    if prefs.require_short:
+        if len(text) > prefs.max_chars:
+            violations.append("too_long")
+
+    # 2) Bullet preference (only for general/list tasks, not pure preference statements)
+    if prefs.require_bullets and task_type in ("general", "list"):
+        has_bullets = ("- " in text) or ("• " in text) or ("* " in text) or ("\n- " in text)
+        if not has_bullets:
+            violations.append("no_bullets")
+
+    # 3) Language preference (rough heuristic)
+    if prefs.lang == "zh":
+        # For Chinese preference, check if answer has enough non-ASCII chars
+        ascii_count = sum(c.isascii() for c in text)
+        ascii_ratio = ascii_count / max(1, len(text))
+        if ascii_ratio > 0.7:  # Too much ASCII = probably not Chinese
+            violations.append("wrong_lang")
+    elif prefs.lang == "en":
+        # For English preference, check if answer is mostly ASCII
+        ascii_count = sum(c.isascii() for c in text)
+        ascii_ratio = ascii_count / max(1, len(text))
+        if ascii_ratio < 0.5:  # Too little ASCII = probably not English
+            violations.append("wrong_lang")
+
+    # 4) Code task: must have code markers
+    prog_t = 1.0
+    if task_type == "code":
+        has_code = ("```" in text) or ("def " in text) or ("function " in text)
+        if not has_code:
+            violations.append("no_code_block")
+            prog_t = 0.0
+
+    # 5) Compute sat_t and sev_t from violations
+    if not violations:
+        sat_t = 1.0
+        sev_t = 0.0
+    else:
+        # Each violation costs 0.3, minimum 0
+        sat_t = max(0.0, 1.0 - 0.3 * float(len(violations)))
+        # Hard violations trigger sev_t=1
+        hard_violations = {"empty_answer", "too_long", "wrong_lang"}
+        sev_t = 1.0 if any(v in hard_violations for v in violations) else 0.0
+
+    return JudgeResult(
+        sat_t=sat_t,
+        sev_t=sev_t,
+        prog_t=prog_t,
+        violations=violations,
+    )
+
+
+# =============================================================================
+# Feedback Computation (reward + gating)
+# =============================================================================
+
+def compute_feedback_for_turn(
+    turn_id: int,
+    query: str,
+    query_type: str,
+    task_type: str,
+    judge_result: JudgeResult,
+) -> Tuple[float, float]:
+    """
+    Convert JudgeResult into (reward, gating):
+    - reward = sat_t (style satisfaction)
+    - gating = 1 only if this turn is preference-related (declared or complained)
+    
+    Args:
+        turn_id: The turn index
+        query: The user's query text
+        query_type: "preference" or "task" from query metadata
+        task_type: "general", "list", "code", etc.
+        judge_result: The judge's evaluation
+    
+    Returns:
+        (reward, gating) tuple
+    """
+    reward = judge_result.sat_t
+
+    # Gating logic: only allow RL update on preference-related turns
+    # 1. Explicit preference declaration (query_type == "preference")
+    # 2. Complaint about not following preference
+    lower_q = (query or "").lower()
+    
+    is_pref_turn = (
+        query_type == "preference"
+        or "i prefer" in lower_q
+        or "my preference" in lower_q
+        or "please use" in lower_q
+        or "please keep" in lower_q
+        or "you didn't follow" in lower_q
+        or "you forgot" in lower_q
+        or "remember that i" in lower_q
+        or "i told you" in lower_q
+        or "i asked for" in lower_q
+    )
+
+    if is_pref_turn:
+        gating = 1.0
+    else:
+        gating = 0.0
+
+    return reward, gating
+
+
+# =============================================================================
+# Extended Queries for Pilot v1 (~10 turns)
+# =============================================================================
+
+def get_pilot_v1_queries() -> List[Dict[str, Any]]:
+    """
+    Extended query set for pilot v1.
+    Mix of preference declarations and tasks.
+    Tests: length constraint, bullet points, task completion.
+    """
+    return [
+        # Turn 0: Declare length preference
+        {
+            "query": "I prefer short, concise answers. Please keep responses under 200 characters.",
+            "type": "preference",
+            "task_type": "general",
+        },
+        # Turn 1: Task that should be short
+        {
+            "query": "What are three tips for better sleep?",
+            "type": "task",
+            "task_type": "list",
+        },
+        # Turn 2: Declare bullet preference
+        {
+            "query": "I also prefer bullet points when listing things. Please use bullet points.",
+            "type": "preference",
+            "task_type": "general",
+        },
+        # Turn 3: Task that should use bullets
+        {
+            "query": "What are the main benefits of regular exercise?",
+            "type": "task",
+            "task_type": "list",
+        },
+        # Turn 4: Another task (test if preferences stick)
+        {
+            "query": "Name five popular programming languages.",
+            "type": "task",
+            "task_type": "list",
+        },
+        # Turn 5: Complaint if needed (always include to test gating)
+        {
+            "query": "Remember that I asked for short answers with bullet points. Can you list three healthy breakfast ideas?",
+            "type": "preference",
+            "task_type": "list",
+        },
+        # Turn 6: Regular task
+        {
+            "query": "What is the capital of France?",
+            "type": "task",
+            "task_type": "general",
+        },
+        # Turn 7: Task requiring list
+        {
+            "query": "What are four seasons of the year?",
+            "type": "task",
+            "task_type": "list",
+        },
+        # Turn 8: Another preference reminder
+        {
+            "query": "I prefer concise bullet points. Please list three types of renewable energy.",
+            "type": "preference",
+            "task_type": "list",
+        },
+        # Turn 9: Final task - test memory
+        {
+            "query": "Summarize what you know about my communication preferences.",
+            "type": "task",
+            "task_type": "general",
+        },
+    ]
+
+
+# =============================================================================
+# Logging
+# =============================================================================
+
+@dataclass
+class TurnLog:
+    """Log entry for one turn."""
+    turn_id: int
+    query: str
+    query_type: str
+    task_type: str
+    answer: str
+    answer_length: int
+    sat_t: float
+    sev_t: float
+    prog_t: float
+    violations: List[str]
+    reward: float
+    gating: float
+    z_long_norm_before: float
+    z_long_norm_after: float
+    z_short_norm_before: float
+    z_short_norm_after: float
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    num_memories_retrieved: int
+    num_prefs_extracted: int
+
+
+def log_to_jsonl(logs: List[TurnLog], filepath: str):
+    """Save logs to JSONL file."""
+    os.makedirs(os.path.dirname(filepath), exist_ok=True)
+    with open(filepath, "w") as f:
+        for log in logs:
+            f.write(json.dumps(asdict(log)) + "\n")
+
+
+# =============================================================================
+# Pilot Runner v1
+# =============================================================================
+
+def run_pilot_v1(
+    llm: PersonalizedLLM,
+    user_id: str = "pilot_user_v1",
+    prefs: Optional[StylePrefs] = None,
+    queries: Optional[List[Dict[str, Any]]] = None,
+) -> List[TurnLog]:
+    """
+    Run pilot v1 with style-aware judge and gating.
+    
+    Args:
+        llm: PersonalizedLLM instance
+        user_id: User identifier
+        prefs: Style preferences for this user
+        queries: Query list (defaults to get_pilot_v1_queries)
+    
+    Returns:
+        List of TurnLog entries
+    """
+    if prefs is None:
+        # Default preferences: short + bullets + English
+        prefs = StylePrefs(
+            require_short=True,
+            max_chars=200,
+            require_bullets=True,
+            lang="en",
+        )
+    
+    if queries is None:
+        queries = get_pilot_v1_queries()
+    
+    logs: List[TurnLog] = []
+    
+    print(f"\n{'='*60}")
+    print(f"PILOT v1 SESSION: user_id={user_id}, turns={len(queries)}")
+    print(f"Preferences: short={prefs.require_short}, max_chars={prefs.max_chars}, bullets={prefs.require_bullets}, lang={prefs.lang}")
+    print(f"{'='*60}")
+    
+    # Reset user for clean start
+    print(f"\n[Pilot] Resetting user: {user_id}")
+    llm.reset_user(user_id)
+    
+    # Start session
+    print(f"[Pilot] Starting session")
+    llm.reset_session(user_id)
+    
+    # Get initial state
+    state_before = llm.get_user_state_summary(user_id)
+    print(f"[Pilot] Initial state: z_long={state_before['z_long_norm']:.6f}, z_short={state_before['z_short_norm']:.6f}")
+    
+    for turn_id, q_info in enumerate(queries):
+        query = q_info["query"]
+        query_type = q_info.get("type", "task")
+        task_type = q_info.get("task_type", "general")
+        
+        print(f"\n{'─'*60}")
+        print(f"Turn {turn_id} [{query_type}]")
+        print(f"{'─'*60}")
+        print(f"[Query] {query}")
+        
+        # Get state before
+        state_before = llm.get_user_state_summary(user_id)
+        z_long_before = state_before["z_long_norm"]
+        z_short_before = state_before["z_short_norm"]
+        
+        # Apply feedback for previous turn (from turn 1 onwards)
+        if turn_id > 0 and len(logs) > 0:
+            prev_log = logs[-1]
+            prev_query = queries[turn_id - 1]
+            
+            # Re-judge the previous answer with current context
+            # (In practice we already have the result, but this shows the flow)
+            feedback = Feedback(
+                user_id=user_id,
+                turn_id=turn_id - 1,
+                reward=prev_log.reward,
+                gating=prev_log.gating,
+                meta={
+                    "sat_t": prev_log.sat_t,
+                    "sev_t": prev_log.sev_t,
+                    "prog_t": prev_log.prog_t,
+                    "violations": prev_log.violations,
+                    "task_type": prev_log.task_type,
+                    "source": "pilot_v1",
+                }
+            )
+            print(f"[Feedback] turn={turn_id-1}, reward={feedback.reward:.2f}, gating={feedback.gating:.1f}")
+            llm.apply_feedback(feedback)
+        
+        # Chat
+        resp: AssistantResponse = llm.chat(user_id, query)
+        
+        # Truncate answer for display
+        answer_display = resp.answer[:150] + "..." if len(resp.answer) > 150 else resp.answer
+        print(f"[Answer] ({len(resp.answer)} chars) {answer_display}")
+        print(f"[Usage] prompt={resp.usage.prompt_tokens}, completion={resp.usage.completion_tokens}")
+        
+        # Judge with style preferences
+        judge_result = style_judge(query, resp.answer, task_type, prefs)
+        print(f"[Judge] sat={judge_result.sat_t:.2f}, sev={judge_result.sev_t:.1f}, prog={judge_result.prog_t:.1f}")
+        if judge_result.violations:
+            print(f"[Judge] violations={judge_result.violations}")
+        
+        # Compute feedback for THIS turn (will be applied next turn)
+        reward, gating = compute_feedback_for_turn(
+            turn_id=turn_id,
+            query=query,
+            query_type=query_type,
+            task_type=task_type,
+            judge_result=judge_result,
+        )
+        print(f"[Feedback] reward={reward:.2f}, gating={gating:.1f} (computed for this turn)")
+        
+        # Get state after
+        state_after = llm.get_user_state_summary(user_id)
+        z_long_after = state_after["z_long_norm"]
+        z_short_after = state_after["z_short_norm"]
+        
+        # Debug info
+        num_memories = len(resp.debug.selected_memory_ids) if resp.debug else 0
+        num_prefs = len(resp.debug.extracted_preferences) if resp.debug else 0
+        
+        z_long_delta = z_long_after - z_long_before
+        z_short_delta = z_short_after - z_short_before
+        print(f"[State] z_long: {z_long_before:.6f} → {z_long_after:.6f} (Δ={z_long_delta:+.6f})")
+        print(f"[State] z_short: {z_short_before:.6f} → {z_short_after:.6f} (Δ={z_short_delta:+.6f})")
+        print(f"[Debug] memories={num_memories}, prefs_extracted={num_prefs}")
+        
+        # Log
+        log = TurnLog(
+            turn_id=turn_id,
+            query=query,
+            query_type=query_type,
+            task_type=task_type,
+            answer=resp.answer,
+            answer_length=len(resp.answer),
+            sat_t=judge_result.sat_t,
+            sev_t=judge_result.sev_t,
+            prog_t=judge_result.prog_t,
+            violations=judge_result.violations,
+            reward=reward,
+            gating=gating,
+            z_long_norm_before=z_long_before,
+            z_long_norm_after=z_long_after,
+            z_short_norm_before=z_short_before,
+            z_short_norm_after=z_short_after,
+            prompt_tokens=resp.usage.prompt_tokens,
+            completion_tokens=resp.usage.completion_tokens,
+            total_tokens=resp.usage.total_tokens,
+            num_memories_retrieved=num_memories,
+            num_prefs_extracted=num_prefs,
+        )
+        logs.append(log)
+    
+    # Apply final feedback for last turn
+    if len(logs) > 0:
+        last_log = logs[-1]
+        feedback = Feedback(
+            user_id=user_id,
+            turn_id=len(queries) - 1,
+            reward=last_log.reward,
+            gating=last_log.gating,
+            meta={"source": "pilot_v1", "final": True}
+        )
+        print(f"\n[Final Feedback] turn={len(queries)-1}, reward={feedback.reward:.2f}, gating={feedback.gating:.1f}")
+        llm.apply_feedback(feedback)
+    
+    return logs
+
+
+def print_summary_v1(logs: List[TurnLog], prefs: StylePrefs):
+    """Print summary statistics for pilot v1."""
+    print(f"\n{'='*60}")
+    print("PILOT v1 SUMMARY")
+    print(f"{'='*60}")
+    
+    total_turns = len(logs)
+    if total_turns == 0:
+        print("No turns to summarize.")
+        return
+    
+    # Basic stats
+    avg_sat = sum(l.sat_t for l in logs) / total_turns
+    avg_prog = sum(l.prog_t for l in logs) / total_turns
+    total_tokens = sum(l.total_tokens for l in logs)
+    total_prompt = sum(l.prompt_tokens for l in logs)
+    total_completion = sum(l.completion_tokens for l in logs)
+    
+    # Gating stats
+    gated_turns = [l for l in logs if l.gating > 0]
+    non_gated_turns = [l for l in logs if l.gating == 0]
+    
+    print(f"\n--- Turn Statistics ---")
+    print(f"Total turns: {total_turns}")
+    print(f"Gated turns (RL active): {len(gated_turns)}")
+    print(f"Non-gated turns (RL skipped): {len(non_gated_turns)}")
+    
+    print(f"\n--- Satisfaction ---")
+    print(f"Average sat_t (all): {avg_sat:.3f}")
+    if gated_turns:
+        avg_sat_gated = sum(l.sat_t for l in gated_turns) / len(gated_turns)
+        print(f"Average sat_t (gated only): {avg_sat_gated:.3f}")
+    print(f"Average prog_t: {avg_prog:.3f}")
+    
+    print(f"\n--- Token Usage ---")
+    print(f"Total tokens: {total_tokens}")
+    print(f"  Prompt: {total_prompt}")
+    print(f"  Completion: {total_completion}")
+    print(f"Avg tokens/turn: {total_tokens / total_turns:.1f}")
+    
+    # Violations breakdown
+    print(f"\n--- Violations ---")
+    from collections import Counter
+    all_violations = [v for l in logs for v in l.violations]
+    if all_violations:
+        print(f"Total violations: {len(all_violations)}")
+        for v, count in Counter(all_violations).most_common():
+            print(f"  {v}: {count}")
+    else:
+        print("No violations")
+    
+    # Answer length analysis
+    print(f"\n--- Answer Lengths (max_chars={prefs.max_chars}) ---")
+    lengths = [l.answer_length for l in logs]
+    over_limit = sum(1 for l in lengths if l > prefs.max_chars)
+    print(f"Min: {min(lengths)}, Max: {max(lengths)}, Avg: {sum(lengths)/len(lengths):.1f}")
+    print(f"Over limit: {over_limit}/{total_turns}")
+    
+    # RL Health Check
+    print(f"\n--- RL Health Check ---")
+    z_long_changes = [abs(l.z_long_norm_after - l.z_long_norm_before) for l in logs]
+    z_short_changes = [abs(l.z_short_norm_after - l.z_short_norm_before) for l in logs]
+    any_z_long_change = any(c > 1e-6 for c in z_long_changes)
+    any_z_short_change = any(c > 1e-6 for c in z_short_changes)
+    
+    print(f"z_long changed: {any_z_long_change} (max Δ: {max(z_long_changes):.6f})")
+    print(f"z_short changed: {any_z_short_change} (max Δ: {max(z_short_changes):.6f})")
+    
+    if any_z_long_change or any_z_short_change:
+        print("✓ User vectors ARE being updated by RL")
+    else:
+        print("✗ WARNING: User vectors NOT changing")
+        print("  Check: gating=1 on some turns? reward != baseline?")
+    
+    # Per-turn detail table
+    print(f"\n--- Turn-by-Turn Summary ---")
+    print(f"{'Turn':>4} {'Type':>10} {'Len':>5} {'sat':>5} {'gate':>5} {'violations'}")
+    print("-" * 60)
+    for l in logs:
+        viol_str = ",".join(l.violations) if l.violations else "-"
+        print(f"{l.turn_id:>4} {l.query_type:>10} {l.answer_length:>5} {l.sat_t:>5.2f} {l.gating:>5.1f} {viol_str}")
+
+
+def main():
+    print("=" * 60)
+    print("PILOT RUNNER v1 - Style-Aware Judge + Gating")
+    print("=" * 60)
+    print(f"Started at: {datetime.now().isoformat()}")
+    
+    # Define user preferences
+    prefs = StylePrefs(
+        require_short=True,
+        max_chars=200,
+        require_bullets=True,
+        lang="en",
+    )
+    print(f"\n[Config] User preferences: {prefs}")
+    
+    # Initialize LLM
+    print("\n[Init] Loading PersonalizedLLM...")
+    llm = PersonalizedLLM(
+        user_store_path="data/users/user_store_pilot_v1.npz",
+        only_own_memories=True,
+        enable_preference_extraction=True,
+        enable_rl_updates=True,
+    )
+    
+    # Run pilot
+    user_id = "pilot_user_v1"
+    logs = run_pilot_v1(llm, user_id=user_id, prefs=prefs)
+    
+    # Summary
+    print_summary_v1(logs, prefs)
+    
+    # Save logs
+    log_path = f"data/logs/pilot_v1_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jsonl"
+    log_to_jsonl(logs, log_path)
+    print(f"\n[Logs] Saved to: {log_path}")
+    
+    # Final state
+    final_state = llm.get_user_state_summary(user_id)
+    print(f"\n[Final State] {final_state}")
+    
+    print(f"\nCompleted at: {datetime.now().isoformat()}")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()
+