docs/v3_thinking_control_patch.md

V3 Patch: 3 Changes for Task-Aware Thinking Control

Overview

These 3 changes go into the v3 notebook. Each change is a precise cell modification.

Research basis:

• OptimalThinkingBench (2508.13141): "Don't overthink" → -23% tokens, +7.7pp accuracy on Qwen3
• Mid-Think (2601.07036): task-specific thinking control in GRPO → +2.6pp AIME, -15% train time
• L1 (2503.04697): token budgets in prompts work when trained with RL reward signal
• User's proven extraction prompt: XML-tagged structure + few-shot + schema enforcement

---

CHANGE 1: Replace SYSTEM_PT with task-aware system prompts (Cell 3)

REMOVE this block from Cell 3:

SYSTEM_PT = (
    "Você é um assistente de IA especializado em análise de e-commerce brasileiro. "
    "Você compreende avaliações de clientes em português e padrões de comércio brasileiro."
)

REPLACE with:

# ══════════════════════════════════════════════════════════════════════════════
# v3: TASK-AWARE SYSTEM PROMPTS
# ══════════════════════════════════════════════════════════════════════════════
# Research basis:
#   - OptimalThinkingBench (2508.13141): "Don't overthink" → -23% tokens, +7.7pp accuracy on Qwen3
#   - Mid-Think (2601.07036): task-specific thinking control in GRPO → +2.6pp AIME, -15% train time
#   - L1 (2503.04697): token budgets in prompts work when trained with RL reward signal
#   - User's proven extraction prompt: XML-tagged structure + few-shot + schema enforcement

COMPLAINT_CATEGORIES_STR = ", ".join(sorted(VALID_CATEGORIES))

SYSTEM_EXTRACTION = (
    "Você é um motor de extração de dados de e-commerce brasileiro. "
    "Retorne APENAS um objeto JSON válido, sem nenhum texto antes ou depois. "
    "NÃO USE blocos de código markdown (` `` json). "
    "O primeiro caractere da sua resposta deve ser { e o último deve ser }. "
    "Campos não mencionados na avaliação devem ser null — nunca invente valores. "
    "Sem explicação. Sem comentários. Não pense em excesso."
)

SYSTEM_SQL = (
    "Você é um assistente de IA especializado em análise de e-commerce brasileiro. "
    "Você compreende avaliações de clientes em português e padrões de comércio brasileiro.\n\n"
    "Para consultas e análises de dados: pense brevemente sobre a estrutura necessária, "
    "depois apresente a resposta de forma direta com números e dados concretos. "
    "Seja conciso no raciocínio. Não pense em excesso."
)

SYSTEM_INSIGHTS = (
    "Você é um assistente de IA especializado em análise de e-commerce brasileiro. "
    "Você compreende avaliações de clientes em português e padrões de comércio brasileiro.\n\n"
    "Para análises estratégicas: raciocine de forma estruturada e concisa, "
    "focando nos pontos principais e recomendações acionáveis. "
    "Use no máximo 500 tokens para raciocinar antes de responder."
)

SYSTEM_PUSH = (
    "Você é um assistente de IA especializado em análise de e-commerce brasileiro. "
    "Você compreende avaliações de clientes em português e padrões de comércio brasileiro.\n\n"
    "Para notificações push: seja direto e criativo. "
    "A notificação deve ter no máximo 120 caracteres. "
    "Responda diretamente sem pensar em excesso."
)

# Legacy fallback — used only in cells that don't have task context
SYSTEM_PT = (
    "Você é um assistente de IA especializado em análise de e-commerce brasileiro. "
    "Você compreende avaliações de clientes em português e padrões de comércio brasileiro."
)

def get_system_prompt(task_type: str) -> str:
    """Return task-optimized system prompt."""
    return {
        "extraction": SYSTEM_EXTRACTION,
        "sql_qa": SYSTEM_SQL,
        "insights": SYSTEM_INSIGHTS,
        "push": SYSTEM_PUSH,
    }.get(task_type, SYSTEM_PT)

# ── Think token budgets per task (for reward function) ────────────────────────
# These are soft targets — the reward function nudges, not enforces
THINK_BUDGETS = {
    "extraction": 150,   # Extraction barely needs thinking — pattern matching
    "push": 100,         # Push is creative writing, not reasoning
    "sql_qa": 400,       # SQL benefits from brief query planning
    "insights": 800,     # Insights need structured multi-step analysis
}

print("✓ v3 Task-aware system prompts defined")
print(f"  extraction: '{SYSTEM_EXTRACTION[:60]}...'")
print(f"  sql_qa:     '{SYSTEM_SQL[:60]}...'")
print(f"  insights:   '{SYSTEM_INSIGHTS[:60]}...'")
print(f"  push:       '{SYSTEM_PUSH[:60]}...'")

---

CHANGE 2: Add reward_think_efficiency() to Cell 6 (Reward Functions)

ADD this function right before commerce_reward_fn in Cell 6:

def reward_think_efficiency(completion: str, task_type: str) -> float:
    """
    Reward concise thinking, penalize bloated <think> blocks.
    
    v3 NEW — Research basis:
      - OptimalThinkingBench (2508.13141): overthinking hurts accuracy on simple tasks
      - L1 (2503.04697): token budget rewards teach models to control reasoning length
      - Train Long Think Short (2508.08940): triangular length reward around target budget
    
    Returns: -0.05 to +0.1 (small component — nudge, not dominate)
    """
    think_match = re.search(r"<think>(.*?)</think>", completion, re.DOTALL)
    budget = THINK_BUDGETS.get(task_type, 500)
    
    if not think_match:
        # No think block at all
        if task_type in ("extraction", "push"):
            return 0.1   # Great — these tasks don't need thinking
        else:
            return 0.0   # Neutral for analytical tasks
    
    think_content = think_match.group(1).strip()
    think_chars = len(think_content)  # chars as proxy (cheaper than tokenizing)
    # Rough conversion: ~4 chars per token for Portuguese
    think_tokens_approx = think_chars / 4
    
    if think_tokens_approx <= budget:
        # Within budget — reward proportional to how concise
        return 0.1
    elif think_tokens_approx <= budget * 2:
        # Over budget but not catastrophic — linear decay
        overshoot = (think_tokens_approx - budget) / budget
        return 0.1 * (1.0 - overshoot)  # 0.1 → 0.0
    else:
        # Way over budget — mild penalty
        return -0.05

MODIFY commerce_reward_fn dispatch block:

Current code (REMOVE):

        if task == "extraction":
            rewards.append(reward_extraction(comp_text))
        elif task == "sql_qa":
            rewards.append(reward_sql_qa(comp_text))
        elif task == "insights":
            rewards.append(reward_insights(comp_text))
        elif task == "push":
            rewards.append(reward_push(comp_text))
        else:
            r = 0.15 if has_think_block(comp_text) else 0.0
            r += 0.2 if comp_text.strip() else 0.0
            rewards.append(r)

New code (REPLACE WITH):

        if task == "extraction":
            task_r = reward_extraction(comp_text)
        elif task == "sql_qa":
            task_r = reward_sql_qa(comp_text)
        elif task == "insights":
            task_r = reward_insights(comp_text)
        elif task == "push":
            task_r = reward_push(comp_text)
        else:
            task_r = 0.15 if has_think_block(comp_text) else 0.0
            task_r += 0.2 if comp_text.strip() else 0.0

        # v3: Think efficiency bonus/penalty (small weight — nudge, not dominate)
        think_r = reward_think_efficiency(comp_text, task)
        rewards.append(task_r + think_r)

---

CHANGE 3: Wire system prompts into data preparation and eval

Cell 7 (Calibration) — add helper + use in loop:

Add this helper function after loading by_type:

def inject_task_system_prompt(msgs, task_type):
    """Replace generic system prompt with task-specific one."""
    new_msgs = []
    system_prompt = get_system_prompt(task_type)
    has_system = False
    for m in msgs:
        if m["role"] == "system":
            new_msgs.append({"role": "system", "content": system_prompt})
            has_system = True
        else:
            new_msgs.append(m)
    if not has_system:
        new_msgs.insert(0, {"role": "system", "content": system_prompt})
    return new_msgs

Then in the calibration loop, inject the task-aware prompt before template application:

for i, msgs in enumerate(cal_samples):
    # Determine task type from user content
    user_text = " ".join(m["content"] for m in msgs if m["role"] == "user")
    task = _classify_task_type(user_text)
    
    # v3: Inject task-aware system prompt
    msgs = inject_task_system_prompt(msgs, task)
    
    text = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
    # ... rest of loop unchanged

Cell 8 (Dataset Preparation) — inject into train/eval records:

In prepare_grpo_datasets_v3, after building train_records and eval_records (before creating HF Datasets), add:

    # v3: Inject task-aware system prompts into each training record
    for i, record in enumerate(train_records):
        user_text = " ".join(m["content"] for m in record if m["role"] == "user")
        task = _classify_task_type(user_text)
        train_records[i] = inject_task_system_prompt(record, task)
    
    # Same for eval records
    for i, record in enumerate(eval_records):
        user_text = " ".join(m["content"] for m in record if m["role"] == "user")
        task = _classify_task_type(user_text)
        eval_records[i] = inject_task_system_prompt(record, task)
    
    print(f"  ✓ Task-aware system prompts injected")

Cell 11 (EvalRewardCallback) — no change needed:

System prompts were injected in Cell 8, so eval data already has the right prompts.

Cell 13 (Validation) — use task-aware selection:

Replace:

system_msg = {"role": "system", "content": SYSTEM_PT}

With task-aware selection inside the loop:

# REMOVE the fixed system_msg line above the loop

# Inside the loop, before generating:
    task = _classify_task_type(prompt["content"])
    system_msg = {"role": "system", "content": get_system_prompt(task)}
    messages = [system_msg, prompt]

---

Summary

Cell	What changes	Lines affected
Cell 3	Replace SYSTEM_PT with 4 task prompts + get_system_prompt() + THINK_BUDGETS	~50 lines added
Cell 6	Add reward_think_efficiency(), modify commerce_reward_fn dispatch	~35 lines added, ~10 modified
Cell 7	Add inject_task_system_prompt(), use in calibration loop	~15 lines added
Cell 8	Inject task-aware system prompts into train/eval records	~10 lines added
Cell 13	Use get_system_prompt(task) instead of fixed SYSTEM_PT	~3 lines modified

Expected impact

Task	Current think tokens	Expected after patch	Mechanism
Extraction	2000-3000 (100% ceiling)	~300-800 (-60-70%)	"Não pense em excesso" + think penalty reward
Push	1000-2000	~100-300 (-70-80%)	"Responda diretamente" + think penalty reward
SQL Q&A	1500-2500	~400-800 (-50%)	"Seja conciso no raciocínio" + think budget reward
Insights	2000-3200 (ceiling)	~800-1500 (-30-40%)	"Use no máximo 500 tokens" + higher think budget