train_basil_v2.py

# train_basil_v2.py
# Dual-objective training script for Basil with LoRA adapter support.
#
# Supports multiple modes:
#   --mode mixed    : (DEFAULT) Alternating world/basil with LoRA adapters
#   --mode world    : Train base model (trunk) on full transcripts only
#   --mode basil    : Train LoRA adapter on Basil turns only (trunk frozen)
#   --mode session  : Legacy full-session mode (aliases to world)
#   --mode graded   : Legacy graded turns mode
#   --mode legacy   : Legacy conversation logs mode [legacy]
#   --mode combined : Legacy combined mode [legacy]
#
# Trunk masking (--mask-mode):
#   none            : No masking — trunk sees everything
#   basil_only      : Mask only Basil's output tokens
#   basil_and_after : (DEFAULT) Mask Basil + Sophie post-grade + Tutor answer
#
# LoRA epoch cap (--lora-max-epochs):
#   If unset, scales linearly with age_band: 0 at band 0, max_epochs at band 7.
#   Manual override available for experiments.

import os
import json
import time
import math
import argparse
import torch
from datetime import datetime
from transformers import AutoTokenizer, AutoModelForCausalLM
from torch.utils.data import Dataset, DataLoader, random_split
from torch.optim.lr_scheduler import LambdaLR
from torch.cuda.amp import autocast, GradScaler
from peft import LoraConfig, get_peft_model, PeftModel

from config import (
    LOG_DIR, MODELS_DIR,
    TRAIN_MIN_SCORE, MIN_TRAINING_WEIGHT,
    score_to_weight, score_to_weight_basil_policy, get_basil_model_name,
    RECENCY_HALF_LIFE_TRAIN_RUNS_GRADED, RECENCY_HALF_LIFE_TRAIN_RUNS_LEGACY,
    MIN_RECENCY_WEIGHT, LEGACY_WEIGHT_MULT, recency_weight,
    LORA_RANK, LORA_ALPHA, LORA_DROPOUT, LORA_TARGET_MODULES,
    LORA_ADAPTER_NAME, LORA_ADAPTER_SUBDIR,
    LORA_ACTIVATION_AGE_BAND,
    BASIL_POLICY_SCORE_WEIGHTS_TABLE, TRUNK_WEIGHT_DIVISOR,
    lora_max_epochs_for_age_band,
)
from utils.scoring import average_entropy
from memory_manager import update_metrics_after_training, load_basil_assessment, load_metrics
from metrics_manager import record_training_attempt, load_rolling_metrics

# === CONFIG ===
MAX_AGE_DAYS = 365
ENTROPY_THRESHOLD = 8
# Block size fixed at model's n_positions limit (1024 for GPT-2)
# Full sessions (~3-5K tokens) use sliding window with stride = block_size // 2
BLOCK_SIZE_BY_AGE_BAND = {
    0: 1024, 1: 1024, 2: 1024, 3: 1024,
    4: 1024, 5: 1024, 6: 1024, 7: 1024,
}
# STRIDE removed - computed as block_size // 2 inside dataset classes

# --- Batching (target: 32-64 sequences per optimizer step = 32k-64k tokens) ---
# A100 40GB: batch_size=2, grad_accum=16 -> effective 32 sequences
# A100 80GB / H100: batch_size=4, grad_accum=16 -> effective 64 sequences
BATCH_SIZE = 4  # Microbatch size (adjust down if OOM)
ACCUMULATION_STEPS = 16  # Effective batch = BATCH_SIZE * ACCUMULATION_STEPS

# --- Learning Rate ---
PEAK_LR = 1e-4  # For training from scratch (random init)
CONTINUE_PEAK_LR = 3e-5  # For continued training (model already partially trained)
LORA_PEAK_LR = 3e-5  # LoRA adapter LR — lower than trunk because alpha/rank amplifies updates
LORA_CONTINUE_LR = 1e-5  # LoRA LR for continued training
MIN_LR = 1e-5  # 10% of PEAK_LR for cosine floor
WARMUP_RATIO = 0.03  # 3% of total steps for linear warmup
GRAD_CLIP_NORM = 1.0  # Gradient clipping

# --- Training limits ---
MAX_TRAIN_TIME = 12 * 3600  # 12 hours in seconds
MAX_EPOCHS = 100

# --- Validation-based early stopping ---
VAL_SPLIT_RATIO = 0.05  # 5% held out for validation
EVAL_EVERY_STEPS = 500  # Evaluate validation loss every N optimizer steps
EARLY_STOP_MIN_DELTA = 0.005  # Stop if val loss doesn't improve by at least this
EARLY_STOP_PATIENCE = 8  # Number of evals without improvement before stopping

# --- Val loss floor (global safety net) ---
# Single global floor to catch extreme memorization. Patience-based early
# stopping (EARLY_STOP_PATIENCE) is the primary guard against overfitting;
# this floor is a last-resort safety net only.
VAL_LOSS_FLOOR_BY_AGE_BAND = {band: 0.5 for band in range(8)}

# Legacy (deprecated, kept for reference)
_LEGACY_MIN_LOSS_DELTA = 0.1
_LEGACY_EARLY_LOSS_THRESHOLD = 2.4


# =============================================================================
# DATA LOADING
# =============================================================================

def _parse_session_key_from_filename(fname: str) -> str:
    """
    Extract a sortable session key from a log filename.
    
    Supports formats:
      - log_YYYYMMDD_HHMMSS.jsonl -> "YYYYMMDD_HHMMSS"
      - session_YYYYMMDD_HHMMSS_graded.jsonl -> "YYYYMMDD_HHMMSS"
      - day_NNN.jsonl -> "day_NNN" (numeric sort handled separately)
    
    Returns the session key string for sorting.
    """
    if fname.startswith("log_"):
        # log_YYYYMMDD_HHMMSS.jsonl or log_YYYYMMDD_HHMMSS_graded.jsonl
        parts = fname.replace(".jsonl", "").replace("_graded", "").replace("_structured", "")
        return parts.replace("log_", "")
    elif fname.startswith("session_"):
        # session_YYYYMMDD_HHMMSS_graded.jsonl
        parts = fname.replace(".jsonl", "").replace("_graded", "").replace("_structured", "")
        return parts.replace("session_", "")
    elif fname.startswith("day_"):
        # day_NNN.jsonl - pad for lexicographic sorting
        day_num = fname.replace("day_", "").replace(".jsonl", "")
        return f"day_{int(day_num):06d}"
    else:
        # Fallback: use filename as-is
        return fname


def _parse_turn(turn: dict, session_key: str, min_score: float) -> dict:
    """
    Parse a single graded turn record into a training example.
    
    Returns a dict suitable for training, or None if the turn should be skipped.
    """
    # Skip identity probe entries (logged for analysis only)
    if turn.get("type") == "identity_probe":
        return None
    
    # Skip wrapup turns (not for training)
    if turn.get("type") == "wrapup":
        return None
    
    # Extract score and weight
    grade = turn.get("grade", {})
    score = grade.get("score", 0)
    weight = turn.get("weight", score_to_weight(score))
    
    if score < min_score:
        return None
    
    # Build training example
    # Context: what Basil saw before responding
    context_parts = []
    
    # Phase B.1: Tutor teaching paragraph
    if turn.get("tutor_teaching"):
        context_parts.append(f"Tutor: {turn['tutor_teaching']}")
    
    # Phase B.2: Sophie reaction to teaching
    if turn.get("sophie_reaction"):
        context_parts.append(f"Sophie: {turn['sophie_reaction']}")
    
    # Phase B.3: Task delivery (naturalized form)
    task = turn.get("task_spec", {})
    task_text = task.get("task_text", "")
    asker = task.get("asker", "tutor")
    if task_text:
        if asker == "sophie":
            context_parts.append(f"Sophie: Basil, {task_text}")
        else:
            context_parts.append(f"Tutor: Basil, {task_text}")
    
    context = "\n".join(context_parts)
    basil_output = turn.get("basil", "")
    
    # Phase E: Sophie's post-grade response (encouragement + question)
    sophie_post_grade = turn.get("sophie", "")
    
    # Phase F: Tutor answers Sophie's question
    tutor_answer = turn.get("tutor_answer", "")
    
    if not basil_output.strip():
        return None
    
    return {
        "context": context,
        "basil_output": basil_output,
        "sophie_post_grade": sophie_post_grade,
        "tutor_answer": tutor_answer,
        "weight": weight,
        "task_category": turn.get("task_category", "unknown"),
        "score": score,
        "session_key": session_key,
    }


def load_graded_turns(min_score: float = None) -> list:
    """
    Load graded turns from Bootstrap Basil log files.
    
    Supports two file formats:
      - Legacy per-session: session_YYYYMMDD_HHMMSS_graded.jsonl
      - Batch-level:        batch_YYYYMMDD_HHMMSS_graded.jsonl
        (each record has an embedded "session_id" field)
    
    Returns list of dicts with keys: context, basil_output, sophie_post_grade,
    tutor_answer, weight, task_category, score, session_key (for recency ordering)
    """
    min_score = min_score if min_score is not None else TRAIN_MIN_SCORE
    graded_turns = []
    
    if not os.path.exists(LOG_DIR):
        return graded_turns
    
    for fname in sorted(os.listdir(LOG_DIR)):
        if not fname.endswith("_graded.jsonl"):
            continue
        
        filepath = os.path.join(LOG_DIR, fname)
        
        if fname.startswith("batch_"):
            # Batch-level file: session_id is embedded in each record
            try:
                with open(filepath, "r") as f:
                    for line in f:
                        turn = json.loads(line)
                        # Derive session_key from embedded session_id
                        session_key = turn.get("session_id", _parse_session_key_from_filename(fname))
                        parsed = _parse_turn(turn, session_key, min_score)
                        if parsed:
                            graded_turns.append(parsed)
            except Exception as e:
                print(f"[!] Error loading {fname}: {e}")
                continue
        elif fname.startswith("session_"):
            # Legacy per-session file: session_key from filename
            session_key = _parse_session_key_from_filename(fname)
            try:
                with open(filepath, "r") as f:
                    for line in f:
                        turn = json.loads(line)
                        parsed = _parse_turn(turn, session_key, min_score)
                        if parsed:
                            graded_turns.append(parsed)
            except Exception as e:
                print(f"[!] Error loading {fname}: {e}")
                continue
    
    return graded_turns


def load_legacy_logs() -> list:
    """
    Load traditional conversation logs with entropy filtering.
    
    Returns list of (filename, lines, session_key) tuples.
    """
    selected_logs = []
    
    if not os.path.exists(LOG_DIR):
        return selected_logs
    
    today = datetime.now().timetuple().tm_yday
    
    for fname in sorted(os.listdir(LOG_DIR)):
        # Skip graded files and non-jsonl
        if fname.endswith("_graded.jsonl") or fname.endswith("_structured.jsonl"):
            continue
        if not fname.endswith(".jsonl"):
            continue
        
        # Parse date from filename
        try:
            if fname.startswith("day_"):
                day_num = int(fname.split("_")[1].split(".")[0])
                age = today - day_num
            elif fname.startswith("log_"):
                log_date_str = fname.split("_")[1]
                log_date = datetime.strptime(log_date_str, "%Y%m%d")
                today_date = datetime.today()
                age = (today_date - log_date).days
            elif fname.startswith("session_"):
                # New session format: session_YYYYMMDD_HHMMSS_structured.jsonl
                continue  # Skip, we use graded files for these
            else:
                print(f"[!] Skipping unrecognized file format: {fname}")
                continue
        except Exception as e:
            print(f"[!] Error parsing filename {fname}: {e}")
            continue
        
        filepath = os.path.join(LOG_DIR, fname)
        with open(filepath, "r") as f:
            lines = [json.loads(line) for line in f]
        
        # Filter out identity_probe entries (logged for analysis only, not training)
        lines = [line for line in lines if line.get("type") != "identity_probe"]
        
        # Filter by entropy
        basil_lines = [line["text"] for line in lines if line.get("speaker") == "Basil"]
        entropy = average_entropy(basil_lines)
        
        session_key = _parse_session_key_from_filename(fname)
        
        if age == 0:
            print(f"[+] Including today's log: {fname}")
            selected_logs.append((fname, lines, session_key))
        elif age <= MAX_AGE_DAYS and entropy < ENTROPY_THRESHOLD:
            print(f"[+] Including {fname} (entropy {entropy:.2f})")
            selected_logs.append((fname, lines, session_key))
        else:
            print(f"[-] Skipping {fname} (entropy {entropy:.2f}, age {age})")
    
    return selected_logs


def load_session_transcripts() -> list:
    """
    Load full session transcripts from *_sessions.jsonl files.
    
    Returns list of session dicts with:
        - session_id: str
        - text: str (full session transcript)
        - basil_spans: list of {"start_char": int, "end_char": int, "score": float}
        - session_key: str (for recency weighting)
    """
    sessions = []
    
    if not os.path.exists(LOG_DIR):
        return sessions
    
    # Group records by session_id
    session_data = {}  # session_id -> {transcript_lines: [], episodes: []}
    
    for fname in sorted(os.listdir(LOG_DIR)):
        if not fname.endswith("_sessions.jsonl"):
            continue
        
        filepath = os.path.join(LOG_DIR, fname)
        try:
            with open(filepath, "r") as f:
                for line in f:
                    record = json.loads(line)
                    session_id = record.get("session_id")
                    if not session_id:
                        continue
                    
                    if session_id not in session_data:
                        session_data[session_id] = {
                            "transcript_lines": [],
                            "episodes": [],
                            "session_key": _parse_session_key_from_filename(fname) if fname.startswith("batch_") else session_id
                        }
                    
                    record_type = record.get("record_type")
                    if record_type == "transcript":
                        # Skip only identity_probe entries (include wrapup)
                        if record.get("type") == "identity_probe":
                            continue
                        session_data[session_id]["transcript_lines"].append(record)
                    elif record_type == "episode":
                        session_data[session_id]["episodes"].append(record)
        except Exception as e:
            print(f"[!] Error loading {fname}: {e}")
            continue
    
    # Build session text and identify Basil spans
    for session_id, data in session_data.items():
        transcript_lines = data["transcript_lines"]
        episodes = data["episodes"]
        
        if not transcript_lines:
            continue
        
        # Build full session text
        text_parts = []
        for line in transcript_lines:
            speaker = line.get("speaker", "")
            text = line.get("text", "")
            if speaker and text:
                # Strip [WRAPUP] prefix if present
                if text.startswith("[WRAPUP]"):
                    text = text.replace("[WRAPUP]", "", 1).strip()
                text_parts.append(f"{speaker}: {text}\n")
        
        full_text = "".join(text_parts)
        if not full_text.strip():
            continue
        
        # Build Basil span map: find each Basil response in text and map to score.
        # Also build episode-local context for LoRA training to avoid contamination
        # from previous Basil garbage in multi-episode sessions.
        basil_spans = []
        for episode in episodes:
            basil_output = episode.get("basil_output", "")
            grade = episode.get("grade", {})
            score = grade.get("score", 0)
            
            if not basil_output:
                continue
            
            # Build episode-local context from messages (phases before "C")
            # This gives the LoRA clean context without prior Basil garbage.
            # Apply same cleaning as inference (_get_episode_transcript) so context matches.
            episode_messages = episode.get("messages", [])
            episode_context = None
            if episode_messages:
                from utils.transcript import clean_transcript_for_training
                ctx_parts = []
                for msg in episode_messages:
                    if msg.get("phase") == "C":
                        # Add speaker prefix but stop before Basil's reply
                        ctx_parts.append("Basil: ")
                        break
                    content = clean_transcript_for_training(msg.get("content", ""))
                    if content:
                        ctx_parts.append(f"{msg['role']}: {content}\n")
                if ctx_parts:
                    episode_context = "".join(ctx_parts)
            
            # Find Basil response in transcript text
            # Look for "Basil: {basil_output}" pattern
            search_pattern = f"Basil: {basil_output}"
            start_char = full_text.find(search_pattern)
            
            if start_char != -1:
                # Found it - mark the span
                # Start after "Basil: " prefix
                basil_start = start_char + len("Basil: ")
                basil_end = basil_start + len(basil_output)
                
                basil_spans.append({
                    "start_char": basil_start,
                    "end_char": basil_end,
                    "score": score,
                    "episode_context": episode_context,
                    "basil_reply": basil_output,
                })
            else:
                # Try without "Basil: " prefix (might be in middle of text)
                start_char = full_text.find(basil_output)
                if start_char != -1:
                    basil_spans.append({
                        "start_char": start_char,
                        "end_char": start_char + len(basil_output),
                        "score": score,
                        "episode_context": episode_context,
                        "basil_reply": basil_output,
                    })
        
        sessions.append({
            "session_id": session_id,
            "text": full_text,
            "basil_spans": basil_spans,
            "session_key": data["session_key"]
        })
    
    print(f"[SessionLoader] Loaded {len(sessions)} sessions from *_sessions.jsonl files")
    return sessions


# =============================================================================
# LORA SETUP
# =============================================================================

def setup_lora_model(base_model):
    """Wrap a base GPT-2 model with LoRA adapters for Basil-policy training.
    
    Returns the PeftModel wrapper. The base model weights are preserved inside;
    LoRA adds small trainable low-rank matrices to attention layers.
    """
    lora_config = LoraConfig(
        r=LORA_RANK,
        lora_alpha=LORA_ALPHA,
        lora_dropout=LORA_DROPOUT,
        target_modules=LORA_TARGET_MODULES,
        bias="none",
        task_type="CAUSAL_LM",
    )
    peft_model = get_peft_model(base_model, lora_config)
    
    trainable_params = sum(p.numel() for p in peft_model.parameters() if p.requires_grad)
    all_params = sum(p.numel() for p in peft_model.parameters())
    print(f"[LoRA] Adapter configured: rank={LORA_RANK}, alpha={LORA_ALPHA}, dropout={LORA_DROPOUT}")
    print(f"[LoRA] Target modules: {LORA_TARGET_MODULES}")
    print(f"[LoRA] Trainable params: {trainable_params:,} / {all_params:,} "
          f"({100 * trainable_params / all_params:.2f}%)")
    return peft_model


def freeze_for_world_training(model):
    """Configure model for World/trunk training: disable LoRA, train base params only."""
    if hasattr(model, 'disable_adapter_layers'):
        model.disable_adapter_layers()
    
    for name, param in model.named_parameters():
        if 'lora_' in name:
            param.requires_grad = False
        else:
            param.requires_grad = True
    
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"[World Mode] Trunk trainable: {trainable:,} params, LoRA frozen/disabled")


def freeze_for_basil_training(model):
    """Configure model for Basil-policy training: enable LoRA, freeze base params."""
    if hasattr(model, 'enable_adapter_layers'):
        model.enable_adapter_layers()
    
    # Freeze everything first
    for param in model.parameters():
        param.requires_grad = False
    
    # Unfreeze only LoRA params
    for name, param in model.named_parameters():
        if 'lora_' in name:
            param.requires_grad = True
    
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"[Basil Mode] LoRA trainable: {trainable:,} params, trunk frozen")


def load_model_with_lora(model_path, device=None):
    """Load base model and optionally attach LoRA adapter for generation.
    
    Returns (model, tokenizer). If an adapter exists at
    {model_path}/{LORA_ADAPTER_SUBDIR}, it is loaded and enabled.
    """
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    base_model = AutoModelForCausalLM.from_pretrained(model_path)
    
    adapter_path = os.path.join(model_path, LORA_ADAPTER_SUBDIR)
    if os.path.exists(adapter_path):
        print(f"[LoRA] Loading adapter from {adapter_path}")
        model = PeftModel.from_pretrained(base_model, adapter_path)
        # Ensure adapter is active for generation
        if hasattr(model, 'enable_adapter_layers'):
            model.enable_adapter_layers()
    else:
        model = base_model
    
    model.eval()
    if device:
        model.to(device)
    
    return model, tokenizer


# =============================================================================
# DATASETS
# =============================================================================

def _get_training_boundaries() -> list:
    """
    Get chronologically sorted training-run boundary session IDs from metrics.
    
    Each boundary is the session ID (YYYYMMDD_HHMMSS) of the last session that
    was generated by the *previous* model version before a training run produced
    a new version. This is recorded by the orchestrator in metrics.json via
    record_training_attempt().
    
    All sessions with keys <= boundary[i] were generated by model version i.
    Sessions with keys > boundary[-1] are from the current (latest) model.
    
    Returns:
        Sorted list of boundary session IDs. Empty list on a fresh repo (no
        training has occurred yet), in which case all sessions belong to run 0
        and get identical recency weight.
    """
    metrics = load_metrics()
    boundaries = metrics.get("training_run_boundaries", [])
    return sorted(boundaries)


def _compute_training_run_for_sessions(session_keys: list) -> dict:
    """
    Map each session key to its training run index.
    
    All sessions produced by the same model version (between two training
    boundaries) get the same run index. The current (newest) run gets the
    highest index.
    
    Boundary semantics: boundary[i] is the *last session generated by model i*.
    So:
      Run 0: sessions with key <= boundary[0]   (original model)
      Run 1: sessions with key <= boundary[1] and > boundary[0]
      ...
      Run N: sessions with key > boundary[-1]   (current model)
    
    Returns:
        dict mapping session_key -> run_index (0 = oldest run, max = current run)
    """
    boundaries = _get_training_boundaries()
    
    # If no training has happened yet, all sessions are in run 0
    if not boundaries:
        return {key: 0 for key in set(session_keys)}
    
    mapping = {}
    for key in set(session_keys):
        run = 0
        for boundary in boundaries:
            if key > boundary:
                run += 1
            else:
                break
        mapping[key] = run
    
    return mapping


class WeightedGradedDataset(Dataset):
    """
    Dataset for graded turns with per-example weights.
    
    Each example is a (context + Basil response) with an associated weight.
    Weights include both score-based and recency-based components.
    """
    
    def __init__(self, graded_turns: list, tokenizer, block_size: int = 1024):
        self.examples = []
        self.weights = []
        
        # Compute training-run mapping for recency weighting.
        # All sessions from the same model version get the same weight.
        session_keys = [turn["session_key"] for turn in graded_turns]
        session_to_run = _compute_training_run_for_sessions(session_keys)
        newest_run = max(session_to_run.values()) if session_to_run else 0
        num_runs = len(set(session_to_run.values()))
        
        # Track recency weights for debug logging
        recency_weights_debug = []
        
        # Pre-compute marker tokens for Basil output boundary detection
        # These are used to find where Basil's output starts/ends in tokenized sequences
        basil_marker_tokens = tokenizer.encode("\nBasil:", add_special_tokens=False)
        sophie_marker_tokens = tokenizer.encode("\nSophie:", add_special_tokens=False)
        tutor_marker_tokens = tokenizer.encode("\nTutor:", add_special_tokens=False)
        
        n_basil_masked = 0  # Track how many examples had Basil output masked
        
        for turn in graded_turns:
            # Build full episode text in chronological order:
            #   B.1 Tutor teaching + B.2 Sophie reaction + B.3 Task delivery
            #   C   Basil's attempt
            #   E   Sophie's post-grade response
            #   F   Tutor answers Sophie's question
            context = turn["context"]
            basil_output = turn["basil_output"]
            sophie_post_grade = turn.get("sophie_post_grade", "")
            tutor_answer = turn.get("tutor_answer", "")
            score = turn["score"]
            score_weight = turn["weight"]  # Original score-based weight
            
            # Compute recency weight (per training run, not per session)
            session_key = turn["session_key"]
            run_index = session_to_run.get(session_key, newest_run)
            runs_ago = newest_run - run_index
            rec_w = recency_weight(runs_ago, RECENCY_HALF_LIFE_TRAIN_RUNS_GRADED)
            recency_weights_debug.append(rec_w)
            
            # Final weight = score_weight * recency_weight
            final_weight = score_weight * rec_w
            
            # Format as full episode training example
            full_text = f"{context}\nBasil: {basil_output}"
            if sophie_post_grade:
                full_text += f"\nSophie: {sophie_post_grade}"
            if tutor_answer:
                full_text += f"\nTutor: {tutor_answer}"
            
            # Tokenize
            tokenized = tokenizer(
                full_text,
                return_tensors="pt",
                truncation=True,
                max_length=block_size,
                padding="max_length",
            )
            
            input_ids = tokenized["input_ids"].squeeze()
            attention_mask = tokenized["attention_mask"].squeeze()
            
            # Create labels: train on full conversation with reward weighting
            # The example-level weight handles reinforcement:
            #   - Score 5 -> weight 1.0 (full signal from entire episode)
            #   - Score 0 -> weight 0.05 (learns language patterns from context)
            labels = input_ids.clone()
            
            # CRITICAL: Mask padding tokens (they shouldn't contribute to loss)
            labels[attention_mask == 0] = -100
            
            # BOOTSTRAPPING: When score is 0, mask Basil's output tokens
            # This prevents reinforcing gibberish while still learning language
            # patterns from the surrounding Tutor/Sophie context
            basil_masked = False
            if score == 0:
                input_list = input_ids.tolist()
                content_end = int((attention_mask == 1).sum().item())
                marker_len = len(basil_marker_tokens)
                
                # Find "\nBasil:" marker in token sequence
                basil_start = -1
                for i in range(content_end - marker_len + 1):
                    if input_list[i:i + marker_len] == basil_marker_tokens:
                        basil_start = i + marker_len
                        break
                
                if basil_start != -1:
                    # Find where Basil's output ends (next speaker or end of content)
                    basil_end = content_end
                    for i in range(basil_start, content_end):
                        # Check for Sophie marker
                        if (i + len(sophie_marker_tokens) <= content_end and
                                input_list[i:i + len(sophie_marker_tokens)] == sophie_marker_tokens):
                            basil_end = i
                            break
                        # Check for Tutor marker
                        if (i + len(tutor_marker_tokens) <= content_end and
                                input_list[i:i + len(tutor_marker_tokens)] == tutor_marker_tokens):
                            basil_end = i
                            break
                    
                    labels[basil_start:basil_end] = -100
                    basil_masked = True
                    n_basil_masked += 1
                else:
                    # Fallback: couldn't find marker, don't mask (still learns from context)
                    print(f"[GradedDataset] WARNING: '\\nBasil:' marker not found in example "
                          f"{len(self.examples) + 1}, score=0 output NOT masked")
            
            # Enhanced diagnostics: log first 3 examples
            if len(self.examples) < 3:
                n_supervised = (labels != -100).sum().item()
                n_masked = (labels == -100).sum().item()
                n_padding = (attention_mask == 0).sum().item()
                n_basil_tokens_masked = n_masked - n_padding
                
                supervised_tokens = tokenizer.decode(labels[labels != -100][:50])
                print(f"[GradedDataset] Example {len(self.examples) + 1} diagnostics:")
                print(f"  Score: {score}, Basil masked: {basil_masked}")
                print(f"  Full episode phases: B+C{'+E' if sophie_post_grade else ''}{'+F' if tutor_answer else ''}")
                print(f"  Supervised tokens: {n_supervised} (context + post-Basil phases)")
                print(f"  Masked tokens: {n_masked} (padding={n_padding}, basil_output={n_basil_tokens_masked})")
                print(f"  Supervised text preview: {supervised_tokens[:200]}...")
                print(f"  Example weight: {final_weight:.3f}")
            
            self.examples.append({
                "input_ids": input_ids,
                "attention_mask": attention_mask,
                "labels": labels,
            })
            self.weights.append(final_weight)
        
        print(f"[GradedDataset] Basil output masked in {n_basil_masked}/{len(self.examples)} examples (score=0)")
        
        # Debug logging for recency stats
        num_sessions = len(set(session_keys))
        print(f"[GradedDataset] Created {len(self.examples)} weighted examples")
        print(f"[GradedDataset] Recency ordering: per training run ({num_runs} runs, {num_sessions} sessions)")
        print(f"[GradedDataset] Current run index: {newest_run}")
        if recency_weights_debug:
            min_rec = min(recency_weights_debug)
            max_rec = max(recency_weights_debug)
            mean_rec = sum(recency_weights_debug) / len(recency_weights_debug)
            print(f"[GradedDataset] Recency weights: min={min_rec:.3f}, max={max_rec:.3f}, mean={mean_rec:.3f}")
        if self.weights:
            avg_weight = sum(self.weights) / len(self.weights)
            print(f"[GradedDataset] Final weights (score*recency): avg={avg_weight:.3f}")
    
    def __len__(self):
        return len(self.examples)
    
    def __getitem__(self, i):
        example = self.examples[i].copy()
        example["weight"] = torch.tensor(self.weights[i], dtype=torch.float32)
        return example


class WorldDataset(Dataset):
    """
    Dataset for the WORLD/TRUNK training objective.
    
    Standard next-token language model on full session transcripts (Tutor + Sophie +
    Story + Basil). Trains the base model to learn language and story structure.
    
    Sliding window with stride = block_size // 2.
    
    mask_mode controls how Basil-related tokens are handled:
        "none"           : No masking — trunk sees everything at weight 1.0
        "basil_only"     : Mask only Basil's output tokens (labels=-100)
        "basil_and_after": (DEFAULT) Score-weighted masking for Basil output
                           (Zone A) and Sophie's immediate reaction (Zone B).
                           Weight = LoRA weight / TRUNK_WEIGHT_DIVISOR.
                           If weight is 0.0 (low-score garbage), fully masked
                           (labels=-100). If weight > 0.0, token_weights is set
                           to the fractional weight and labels are kept.
                           Everything after Zone B (popquiz, wrap-up, next
                           episode teaching) stays at full weight (1.0).
    
    age_band is used to look up the score-to-weight mapping from
    BASIL_POLICY_SCORE_WEIGHTS_TABLE, ensuring the trunk's rising quality
    floor matches the LoRA's.
    """
    
    def __init__(self, sessions: list, tokenizer, block_size: int = 1024,
                 mask_mode: str = "basil_and_after", age_band: int = 0):
        self.examples = []
        self.token_weights_list = []  # Per-chunk token weight tensors
        self.mask_mode = mask_mode
        self.age_band = max(0, min(int(age_band), 7))
        
        # Compute training-run mapping for recency weighting
        session_keys = [s["session_key"] for s in sessions]
        session_to_run = _compute_training_run_for_sessions(session_keys)
        newest_run = max(session_to_run.values()) if session_to_run else 0
        num_runs = len(set(session_to_run.values()))
        
        stride = block_size // 2  # 50% overlap between chunks
        
        # Track diagnostics
        total_tokens = 0
        total_chunks = 0
        total_tokens_masked = 0
        total_tokens_weighted = 0
        
        for session in sessions:
            session_id = session["session_id"]
            text = session["text"]
            session_key = session["session_key"]
            basil_spans = session.get("basil_spans", [])
            
            # Compute recency weight for this session
            run_index = session_to_run.get(session_key, newest_run)
            runs_ago = newest_run - run_index
            recency_w = recency_weight(runs_ago, RECENCY_HALF_LIFE_TRAIN_RUNS_GRADED)
            
            # Tokenize full session (request offset_mapping for masking)
            encoding = tokenizer(
                text,
                add_special_tokens=False,
                return_attention_mask=False,
                return_offsets_mapping=True,
            )
            input_ids = encoding["input_ids"]
            offsets = encoding.get("offset_mapping", None)
            
            if len(input_ids) == 0:
                continue
            
            total_tokens += len(input_ids)
            
            # Start all tokens at weight 1.0; Basil zones will be adjusted below
            token_weights = torch.ones(len(input_ids), dtype=torch.float32)
            labels = torch.tensor(input_ids, dtype=torch.long)
            
            # --- Mask / weight tokens based on mask_mode ---
            if mask_mode != "none" and basil_spans and offsets:
                sorted_spans = sorted(basil_spans, key=lambda s: s["start_char"])
                
                for si, span in enumerate(sorted_spans):
                    score = int(span.get("score", 0))
                    
                    if mask_mode == "basil_only":
                        # Simple: mask only Basil's output tokens
                        zone_a_start = span["start_char"]
                        zone_a_end = span["end_char"]
                        for tok_idx, (tok_start, tok_end) in enumerate(offsets):
                            if tok_end > zone_a_start and tok_start < zone_a_end:
                                labels[tok_idx] = -100
                                total_tokens_masked += 1
                    else:
                        # basil_and_after: three-zone weighted masking
                        # Zone A = Basil's output
                        # Zone B = Sophie's immediate reaction (1st speaker turn after Basil)
                        # Zone C = everything after Zone B (full weight, not masked)
                        
                        zone_a_start = span["start_char"]
                        zone_a_end = span["end_char"]
                        
                        # Scan forward from Basil's end to find speaker boundaries
                        speaker_markers = ["\nSophie:", "\nTutor:", "\nBasil:"]
                        pos = zone_a_end
                        speakers_found = 0
                        zone_b_end = len(text)  # default: to end of text
                        
                        while pos < len(text):
                            found_marker = False
                            for m in speaker_markers:
                                if text[pos:pos + len(m)] == m:
                                    speakers_found += 1
                                    if speakers_found == 2:
                                        zone_b_end = pos
                                        found_marker = True
                                        break
                                    pos += len(m)
                                    found_marker = True
                                    break
                            if zone_b_end < len(text):
                                break
                            if not found_marker:
                                pos += 1
                        
                        # Compute trunk weight from LoRA table
                        score_row = BASIL_POLICY_SCORE_WEIGHTS_TABLE.get(
                            score, BASIL_POLICY_SCORE_WEIGHTS_TABLE[0]
                        )
                        lora_w = score_row.get(self.age_band, 0.0)
                        trunk_w = lora_w / TRUNK_WEIGHT_DIVISOR
                        
                        # Apply to tokens in Zone A + Zone B
                        for tok_idx, (tok_start, tok_end) in enumerate(offsets):
                            if tok_end > zone_a_start and tok_start < zone_b_end:
                                if trunk_w == 0.0:
                                    labels[tok_idx] = -100
                                    total_tokens_masked += 1
                                else:
                                    token_weights[tok_idx] = trunk_w
                                    total_tokens_weighted += 1
            
            elif mask_mode != "none" and basil_spans and offsets is None:
                # Fallback: tokenizer doesn't support offset_mapping; use token-marker search
                basil_marker = tokenizer.encode("\nBasil:", add_special_tokens=False)
                sophie_marker = tokenizer.encode("\nSophie:", add_special_tokens=False)
                tutor_marker = tokenizer.encode("\nTutor:", add_special_tokens=False)
                all_markers = [sophie_marker, tutor_marker, basil_marker]
                input_list = input_ids
                marker_len = len(basil_marker)
                
                sorted_fb_spans = sorted(basil_spans, key=lambda s: s["start_char"])
                span_idx = 0
                
                i = 0
                while i < len(input_list) - marker_len + 1:
                    if input_list[i:i + marker_len] == basil_marker:
                        zone_a_start = i + marker_len  # after "Basil:" marker
                        
                        if mask_mode == "basil_only":
                            # Find end of Basil's output (next speaker marker)
                            zone_a_end_fb = len(input_list)
                            for j in range(zone_a_start, len(input_list)):
                                for marker in all_markers:
                                    if (j + len(marker) <= len(input_list) and
                                            input_list[j:j + len(marker)] == marker):
                                        zone_a_end_fb = j
                                        break
                                if zone_a_end_fb != len(input_list):
                                    break
                            labels[zone_a_start:zone_a_end_fb] = -100
                            total_tokens_masked += (zone_a_end_fb - zone_a_start)
                            i = zone_a_end_fb
                        else:
                            # basil_and_after: three-zone weighted masking
                            # Find Zone A end (1st speaker) and Zone B end (2nd speaker)
                            speakers_found = 0
                            zone_b_end_fb = len(input_list)
                            for j in range(zone_a_start, len(input_list)):
                                for marker in all_markers:
                                    if (j + len(marker) <= len(input_list) and
                                            input_list[j:j + len(marker)] == marker):
                                        speakers_found += 1
                                        if speakers_found == 2:
                                            zone_b_end_fb = j
                                        break
                                if zone_b_end_fb != len(input_list):
                                    break
                            
                            # Look up score from basil_spans (matched by order)
                            fb_score = 0
                            if span_idx < len(sorted_fb_spans):
                                fb_score = int(sorted_fb_spans[span_idx].get("score", 0))
                                span_idx += 1
                            
                            score_row = BASIL_POLICY_SCORE_WEIGHTS_TABLE.get(
                                fb_score, BASIL_POLICY_SCORE_WEIGHTS_TABLE[0]
                            )
                            lora_w = score_row.get(self.age_band, 0.0)
                            trunk_w = lora_w / TRUNK_WEIGHT_DIVISOR
                            
                            if trunk_w == 0.0:
                                labels[zone_a_start:zone_b_end_fb] = -100
                                total_tokens_masked += (zone_b_end_fb - zone_a_start)
                            else:
                                for k in range(zone_a_start, zone_b_end_fb):
                                    token_weights[k] = trunk_w
                                total_tokens_weighted += (zone_b_end_fb - zone_a_start)
                            
                            i = zone_b_end_fb
                    else:
                        i += 1
            
            # Apply recency weight as multiplier
            token_weights = token_weights * recency_w
            
            # Sliding window: create chunks with stride
            for start_idx in range(0, len(input_ids), stride):
                end_idx = min(start_idx + block_size, len(input_ids))
                
                if end_idx - start_idx < block_size // 4:
                    break
                
                chunk_ids = input_ids[start_idx:end_idx]
                chunk_labels = labels[start_idx:end_idx]
                chunk_weights = token_weights[start_idx:end_idx]
                
                # Pad if needed
                if len(chunk_ids) < block_size:
                    padding_len = block_size - len(chunk_ids)
                    padded_ids = chunk_ids + [tokenizer.pad_token_id] * padding_len
                    padded_labels = torch.cat([
                        chunk_labels,
                        torch.full((padding_len,), -100, dtype=torch.long)
                    ])
                    padded_weights = torch.cat([
                        chunk_weights,
                        torch.zeros(padding_len, dtype=torch.float32)
                    ])
                    attention_mask = torch.cat([
                        torch.ones(len(chunk_ids), dtype=torch.long),
                        torch.zeros(padding_len, dtype=torch.long)
                    ])
                else:
                    padded_ids = chunk_ids
                    padded_labels = chunk_labels
                    padded_weights = chunk_weights
                    attention_mask = torch.ones(len(chunk_ids), dtype=torch.long)
                
                self.examples.append({
                    "input_ids": torch.tensor(padded_ids, dtype=torch.long),
                    "attention_mask": attention_mask,
                    "labels": padded_labels,
                })
                self.token_weights_list.append(padded_weights)
                total_chunks += 1
        
        # Diagnostics
        n_sessions = len(sessions) or 1
        print(f"[WorldDataset] Created {len(self.examples)} chunks from {len(sessions)} sessions")
        print(f"[WorldDataset] Total tokens: {total_tokens:,} ({total_tokens / n_sessions:.0f} per session)")
        print(f"[WorldDataset] Chunks per session: {total_chunks / n_sessions:.1f}")
        if mask_mode != "none":
            print(f"[WorldDataset] Mask mode: {mask_mode} (age_band={self.age_band})")
            print(f"[WorldDataset]   Fully masked (labels=-100): {total_tokens_masked:,} tokens")
            print(f"[WorldDataset]   Partially weighted (Basil+reaction): {total_tokens_weighted:,} tokens")
            print(f"[WorldDataset]   Full weight: {total_tokens - total_tokens_masked - total_tokens_weighted:,} tokens")
        else:
            print(f"[WorldDataset] Mask mode: none (trunk sees full conversation)")
        print(f"[WorldDataset] Recency ordering: per training run ({num_runs} runs, {len(sessions)} sessions)")
        print(f"[WorldDataset] Current run index: {newest_run}")
    
    def __len__(self):
        return len(self.examples)
    
    def __getitem__(self, i):
        example = self.examples[i].copy()
        example["token_weights"] = self.token_weights_list[i]
        return example


# Backward-compatible alias
FullSessionDataset = WorldDataset


class BasilDataset(Dataset):
    """
    Dataset for the BASIL-POLICY training objective.
    
    Trains ONLY on Basil's reply tokens. For each Basil turn:
    - Input: episode-local context up to "Basil:" + Basil's reply
    - Labels: -100 for all context tokens, real token IDs for Basil's reply tokens
    - Per-example weight from score_to_weight_basil_policy(score, age_band)
    
    Episode-local context: for multi-episode classroom sessions, each Basil
    training example sees only the current episode's dialogue (teaching,
    quiz, task) — NOT previous episodes' Basil outputs. This prevents
    contamination from Basil's earlier garbage polluting the LoRA's
    conditioning context. Single-episode sessions (storytime, howitworks)
    fall back to full session context since there are no prior Basil turns.
    
    This trains the LoRA adapter to learn "how Basil should respond" with
    score-weighted reinforcement. Score 0 gets zero weight.  Weights are
    age-band-aware: as Basil matures, lower scores receive less weight.
    """
    
    def __init__(self, sessions: list, tokenizer, block_size: int = 1024,
                 age_band: int = 0):
        self.examples = []
        self.example_weights = []
        self.example_scores = []
        self.age_band = age_band
        
        # Compute training-run mapping for recency weighting
        session_keys = [s["session_key"] for s in sessions]
        session_to_run = _compute_training_run_for_sessions(session_keys)
        newest_run = max(session_to_run.values()) if session_to_run else 0
        num_runs = len(set(session_to_run.values()))
        
        # Track diagnostics
        total_basil_turns = 0
        total_reply_tokens = 0
        score_distribution = {}
        n_skipped_empty = 0
        n_skipped_zero_weight = 0
        n_episode_context = 0  # Used episode-local context (classroom)
        n_fallback_context = 0  # Used full session (storytime, howitworks, or missing messages)
        
        for session in sessions:
            text = session["text"]
            basil_spans = session["basil_spans"]
            session_key = session["session_key"]
            
            # Compute recency weight
            run_index = session_to_run.get(session_key, newest_run)
            runs_ago = newest_run - run_index
            recency_w = recency_weight(runs_ago, RECENCY_HALF_LIFE_TRAIN_RUNS_GRADED)
            
            for span in basil_spans:
                start_char = span["start_char"]
                end_char = span["end_char"]
                score = span["score"]
                
                # Compute weight for this Basil turn (age-band-aware)
                basil_weight = score_to_weight_basil_policy(score, self.age_band)
                final_weight = basil_weight * recency_w
                
                # Track score distribution
                score_distribution[score] = score_distribution.get(score, 0) + 1
                
                # Skip if zero weight (score 0)
                if basil_weight == 0.0:
                    n_skipped_zero_weight += 1
                    continue
                
                # Use episode-local context for LoRA training when available.
                # This avoids contaminating the context with Basil's previous
                # garbage outputs from earlier episodes in multi-episode sessions.
                # Falls back to full session context for single-episode sessions
                # (storytime, howitworks) where episode_context is None.
                episode_context = span.get("episode_context")
                if episode_context is not None:
                    context_text = episode_context
                    basil_reply_text = span.get("basil_reply", text[start_char:end_char])
                    n_episode_context += 1
                else:
                    context_text = text[:start_char]
                    basil_reply_text = text[start_char:end_char]
                    n_fallback_context += 1
                
                if not basil_reply_text.strip():
                    n_skipped_empty += 1
                    continue
                
                # Tokenize context and reply separately to know the boundary
                context_tokens = tokenizer.encode(context_text, add_special_tokens=False)
                reply_tokens = tokenizer.encode(basil_reply_text, add_special_tokens=False)
                
                if not reply_tokens:
                    n_skipped_empty += 1
                    continue
                
                # Truncate context from the left if needed to fit within block_size
                max_context_len = block_size - len(reply_tokens)
                if max_context_len <= 0:
                    # Reply alone exceeds block_size — truncate reply
                    reply_tokens = reply_tokens[:block_size // 2]
                    max_context_len = block_size - len(reply_tokens)
                
                if len(context_tokens) > max_context_len:
                    # Keep the END of context (most relevant recent turns)
                    context_tokens = context_tokens[-max_context_len:]
                
                n_context = len(context_tokens)
                n_reply = len(reply_tokens)
                all_tokens = context_tokens + reply_tokens
                
                # Labels: -100 for context (don't train on context), real IDs for reply
                labels = [-100] * n_context + reply_tokens
                
                # Pad if needed
                seq_len = len(all_tokens)
                if seq_len < block_size:
                    padding_len = block_size - seq_len
                    all_tokens = all_tokens + [tokenizer.pad_token_id] * padding_len
                    labels = labels + [-100] * padding_len
                    attention_mask = [1] * seq_len + [0] * padding_len
                else:
                    attention_mask = [1] * seq_len
                
                self.examples.append({
                    "input_ids": torch.tensor(all_tokens, dtype=torch.long),
                    "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
                    "labels": torch.tensor(labels, dtype=torch.long),
                })
                self.example_weights.append(final_weight)
                self.example_scores.append(score)
                
                total_basil_turns += 1
                total_reply_tokens += n_reply
        
        # Diagnostics
        print(f"[BasilDataset] Created {len(self.examples)} examples from {len(sessions)} sessions")
        print(f"[BasilDataset] Context source: {n_episode_context} episode-local (classroom), "
              f"{n_fallback_context} full-session fallback (storytime/howitworks or missing messages)")
        print(f"[BasilDataset] Total Basil turns: {total_basil_turns}")
        print(f"[BasilDataset] Average reply tokens: {total_reply_tokens / max(1, total_basil_turns):.1f}")
        print(f"[BasilDataset] Skipped: {n_skipped_zero_weight} zero-weight (score=0), "
              f"{n_skipped_empty} empty replies")
        if score_distribution:
            print(f"[BasilDataset] Score distribution: {dict(sorted(score_distribution.items()))}")
        if self.example_weights:
            avg_w = sum(self.example_weights) / len(self.example_weights)
            print(f"[BasilDataset] Average weight: {avg_w:.3f}")
        self.avg_score = (sum(self.example_scores) / len(self.example_scores)) if self.example_scores else 0.0
        print(f"[BasilDataset] Average score: {self.avg_score:.2f}")
        print(f"[BasilDataset] Recency ordering: per training run ({num_runs} runs, {len(sessions)} sessions)")
    
    def __len__(self):
        return len(self.examples)
    
    def __getitem__(self, i):
        example = self.examples[i].copy()
        example["weight"] = torch.tensor(self.example_weights[i], dtype=torch.float32)
        return example


class LegacyDataset(Dataset):
    """
    Dataset for conversation logs with recency-based weighting.
    
    Weights include recency decay and optional legacy multiplier.
    """
    
    def __init__(self, logs: list, tokenizer, block_size: int = 1024, stride: int = 768,
                 apply_legacy_mult: bool = False):
        """
        Args:
            logs: List of (filename, lines, session_key) tuples
            tokenizer: HuggingFace tokenizer
            block_size: Max sequence length
            stride: Sliding window stride
            apply_legacy_mult: If True, apply LEGACY_WEIGHT_MULT (for combined mode)
        """
        self.examples = []
        self.weights = []
        
        # Compute training-run mapping for recency weighting.
        # All sessions from the same model version get the same weight.
        session_keys = [log[2] for log in logs]  # (fname, lines, session_key)
        session_to_run = _compute_training_run_for_sessions(session_keys)
        newest_run = max(session_to_run.values()) if session_to_run else 0
        num_runs = len(set(session_to_run.values()))
        
        # Track recency weights for debug logging
        recency_weights_debug = []
        
        for fname, log, session_key in logs:
            # Compute recency weight (per training run, not per session)
            run_index = session_to_run.get(session_key, newest_run)
            runs_ago = newest_run - run_index
            rec_w = recency_weight(runs_ago, RECENCY_HALF_LIFE_TRAIN_RUNS_LEGACY)
            
            # Apply legacy multiplier if in combined mode
            base_weight = LEGACY_WEIGHT_MULT if apply_legacy_mult else 1.0
            final_weight = base_weight * rec_w
            
            text = ""
            for entry in log:
                # Skip identity_probe entries (logged for analysis only)
                if entry.get("type") == "identity_probe":
                    continue
                if entry.get("speaker") in ["Tutor", "Sophie", "Basil"]:
                    text += f"{entry['speaker']}: {entry['text']}\n"
            
            tokenized = tokenizer(
                text,
                return_attention_mask=True,
                return_tensors="pt",
                padding=True,
                truncation=True,
            )
            input_ids = tokenized["input_ids"].squeeze()
            attention_mask = tokenized["attention_mask"].squeeze()
            
            # Handle 1D tensor case
            if input_ids.dim() == 0:
                continue
            
            # Track recency weight (once per session, applied to all chunks)
            recency_weights_debug.append(rec_w)
            
            # Sliding window with stride
            for start in range(0, max(1, len(input_ids) - block_size + 1), stride):
                end = start + block_size
                if end > len(input_ids):
                    # Pad the last chunk
                    padded_ids = torch.full((block_size,), tokenizer.pad_token_id)
                    padded_mask = torch.zeros((block_size,), dtype=torch.long)
                    chunk_len = len(input_ids) - start
                    padded_ids[:chunk_len] = input_ids[start:]
                    padded_mask[:chunk_len] = attention_mask[start:]
                    
                    # Create labels and mask padding
                    chunk_labels = padded_ids.clone()
                    chunk_labels[padded_mask == 0] = -100
                    
                    self.examples.append({
                        "input_ids": padded_ids,
                        "attention_mask": padded_mask,
                        "labels": chunk_labels,
                    })
                else:
                    chunk_ids = input_ids[start:end]
                    chunk_mask = attention_mask[start:end]
                    
                    # Create labels and mask padding
                    chunk_labels = chunk_ids.clone()
                    chunk_labels[chunk_mask == 0] = -100
                    
                    self.examples.append({
                        "input_ids": chunk_ids,
                        "attention_mask": chunk_mask,
                        "labels": chunk_labels,
                    })
                self.weights.append(final_weight)
        
        # Enhanced diagnostics: log first 3 examples
        if len(self.examples) > 0:
            for i in range(min(3, len(self.examples))):
                example = self.examples[i]
                labels = example["labels"]
                attention_mask = example["attention_mask"]
                n_supervised = (labels != -100).sum().item()
                n_masked = (labels == -100).sum().item()
                n_padding = (attention_mask == 0).sum().item()
                
                supervised_tokens = tokenizer.decode(labels[labels != -100][:30])
                print(f"[LegacyDataset] Example {i+1} diagnostics:")
                print(f"  Supervised tokens: {n_supervised}")
                print(f"  Masked tokens: {n_masked} (padding)")
                print(f"  Padding tokens: {n_padding}")
                print(f"  Supervised text: {supervised_tokens[:100]}...")
        
        # Debug logging for recency stats
        num_sessions = len(set(session_keys))
        print(f"[LegacyDataset] Created {len(self.examples)} legacy examples")
        print(f"[LegacyDataset] Recency ordering: per training run ({num_runs} runs, {num_sessions} sessions)")
        print(f"[LegacyDataset] Current run index: {newest_run}")
        if recency_weights_debug:
            min_rec = min(recency_weights_debug)
            max_rec = max(recency_weights_debug)
            mean_rec = sum(recency_weights_debug) / len(recency_weights_debug)
            print(f"[LegacyDataset] Recency weights: min={min_rec:.3f}, max={max_rec:.3f}, mean={mean_rec:.3f}")
        if self.weights:
            avg_weight = sum(self.weights) / len(self.weights)
            mult_str = f" (with LEGACY_WEIGHT_MULT={LEGACY_WEIGHT_MULT})" if apply_legacy_mult else ""
            print(f"[LegacyDataset] Final weights{mult_str}: avg={avg_weight:.3f}")
    
    def __len__(self):
        return len(self.examples)
    
    def __getitem__(self, i):
        example = self.examples[i].copy()
        example["weight"] = torch.tensor(self.weights[i], dtype=torch.float32)
        return example


class CombinedDataset(Dataset):
    """
    Combines graded and legacy datasets.
    """
    
    def __init__(self, graded_dataset, legacy_dataset):
        self.graded = graded_dataset
        self.legacy = legacy_dataset
        self.graded_len = len(graded_dataset) if graded_dataset else 0
        self.legacy_len = len(legacy_dataset) if legacy_dataset else 0
    
    def __len__(self):
        return self.graded_len + self.legacy_len
    
    def __getitem__(self, i):
        if i < self.graded_len:
            return self.graded[i]
        else:
            return self.legacy[i - self.graded_len]


# =============================================================================
# TRAINING
# =============================================================================

def compute_weighted_loss(model, batch, device):
    """
    Compute weighted cross-entropy loss.
    
    Supports two weighting modes (auto-detected from batch keys):
    - Per-token weights (WorldDataset): batch["token_weights"] shape [batch_size, seq_len]
    - Per-example weights (BasilDataset / legacy): batch["weight"] shape [batch_size]
    """
    input_ids = batch["input_ids"].to(device)
    attention_mask = batch["attention_mask"].to(device)
    
    if "labels" in batch:
        labels = batch["labels"].to(device)
    else:
        labels = input_ids
    
    outputs = model(
        input_ids=input_ids,
        attention_mask=attention_mask,
        labels=labels,
    )
    
    logits = outputs.logits
    
    # Shift for causal LM
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()
    
    # Per-token loss
    loss_fct = torch.nn.CrossEntropyLoss(reduction='none', ignore_index=-100)
    per_token_loss = loss_fct(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1)
    )
    per_token_loss = per_token_loss.view(shift_labels.size())
    
    mask = (shift_labels != -100).float()
    
    if "token_weights" in batch:
        # Per-token weighting (WorldDataset)
        token_weights = batch["token_weights"].to(device)
        shift_weights = token_weights[..., 1:]
        
        weighted_token_loss = per_token_loss * shift_weights
        weighted_loss = (weighted_token_loss * mask).sum() / mask.sum().clamp(min=1)
        raw_loss = (per_token_loss * mask).sum() / mask.sum().clamp(min=1)
    elif "weight" in batch:
        # Per-example weighting (BasilDataset / legacy)
        weights = batch["weight"].to(device)
        
        per_example_loss = (per_token_loss * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
        raw_loss = per_example_loss.mean()
        weighted_loss = (per_example_loss * weights).mean()
    else:
        # No weighting (plain loss)
        weighted_loss = (per_token_loss * mask).sum() / mask.sum().clamp(min=1)
        raw_loss = weighted_loss
    
    return weighted_loss, raw_loss


def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, min_lr_ratio=0.1):
    """
    Create a schedule with linear warmup followed by cosine decay to min_lr.
    
    Args:
        optimizer: The optimizer
        num_warmup_steps: Number of warmup steps
        num_training_steps: Total number of training steps
        min_lr_ratio: Minimum LR as fraction of peak (default 0.1 = 10%)
    """
    def lr_lambda(current_step):
        # Warmup phase: linear from 0 to 1
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        
        # Cosine decay phase
        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
        cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress))
        
        # Scale between min_lr_ratio and 1.0
        return min_lr_ratio + (1.0 - min_lr_ratio) * cosine_decay
    
    return LambdaLR(optimizer, lr_lambda)


def evaluate_validation(model, val_loader, device, use_bf16=False):
    """
    Compute average validation loss.
    Returns (weighted_loss, raw_loss) tuple for diagnostics.
    """
    model.eval()
    total_weighted_loss = 0.0
    total_raw_loss = 0.0
    num_batches = 0
    
    with torch.no_grad():
        for batch in val_loader:
            if use_bf16:
                with autocast(dtype=torch.bfloat16):
                    weighted_loss, raw_loss = compute_weighted_loss(model, batch, device)
            else:
                with autocast():
                    weighted_loss, raw_loss = compute_weighted_loss(model, batch, device)
            total_weighted_loss += weighted_loss.item()
            total_raw_loss += raw_loss.item()
            num_batches += 1
    
    model.train()
    avg_weighted = total_weighted_loss / max(1, num_batches)
    avg_raw = total_raw_loss / max(1, num_batches)
    return avg_weighted, avg_raw


def _run_single_objective_epoch(
    model, train_loader, optimizer, scheduler, scaler,
    device, grad_accum, use_bf16, global_step, start_time,
    val_loader=None, best_val_loss=float('inf'), best_model_state=None,
    patience_counter=0, phase_name="train", val_loss_floor=None,
):
    """
    Run one training epoch (one pass through a DataLoader).
    
    Args:
        val_loss_floor: If set and this is a WORLD phase, stop training when
            val_loss drops below this value (prevents over-training at early
            developmental stages).
    
    Returns (epoch_loss, epoch_raw_loss, global_step, best_val_loss,
             best_model_state, patience_counter, early_stopped)
    """
    epoch_loss = 0.0
    epoch_raw_loss = 0.0
    early_stopped = False
    model.train()
    
    optimizer.zero_grad()
    for step, batch in enumerate(train_loader):
        # Forward pass with mixed precision
        if use_bf16:
            with autocast(dtype=torch.bfloat16):
                weighted_loss, raw_loss = compute_weighted_loss(model, batch, device)
                scaled_loss = weighted_loss / grad_accum
            scaled_loss.backward()
        else:
            with autocast():
                weighted_loss, raw_loss = compute_weighted_loss(model, batch, device)
                scaled_loss = weighted_loss / grad_accum
            scaler.scale(scaled_loss).backward()
        
        epoch_loss += weighted_loss.item()
        epoch_raw_loss += raw_loss.item()
        
        # Optimizer step with gradient accumulation
        if (step + 1) % grad_accum == 0 or (step + 1) == len(train_loader):
            if use_bf16:
                torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP_NORM)
                optimizer.step()
            else:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP_NORM)
                scaler.step(optimizer)
                scaler.update()
            
            optimizer.zero_grad()
            scheduler.step()
            global_step += 1
            
            # Periodic validation
            if val_loader and global_step % EVAL_EVERY_STEPS == 0:
                val_w, val_r = evaluate_validation(model, val_loader, device, use_bf16)
                elapsed = time.time() - start_time
                improvement = best_val_loss - val_w
                improved_str = f"(+{improvement:.4f})" if improvement > EARLY_STOP_MIN_DELTA else ""
                
                avg_w = epoch_loss / max(1, step + 1)
                avg_r = epoch_raw_loss / max(1, step + 1)
                print(f"[{phase_name} Step {global_step}] train={avg_w:.4f} (raw={avg_r:.4f}) | "
                      f"val={val_w:.4f} (raw={val_r:.4f}) {improved_str} | time={elapsed:.0f}s")
                
                if improvement > EARLY_STOP_MIN_DELTA:
                    best_val_loss = val_w
                    best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
                    patience_counter = 0
                    print(f"   New best val_loss: {best_val_loss:.4f}")
                else:
                    patience_counter += 1
                    print(f"   No improvement ({patience_counter}/{EARLY_STOP_PATIENCE})")
                    if patience_counter >= EARLY_STOP_PATIENCE:
                        print(f"\n   Early stopping: no improvement for {EARLY_STOP_PATIENCE} evals")
                        early_stopped = True
                        break
                
                # Val loss floor: prevent over-training at early developmental stages
                if val_loss_floor is not None and val_w <= val_loss_floor:
                    print(f"\n   Val loss floor reached: {val_w:.4f} <= {val_loss_floor} — stopping to prevent memorization")
                    # Save current state as best if it improved
                    if best_model_state is None or val_w < best_val_loss:
                        best_val_loss = val_w
                        best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
                    early_stopped = True
                    break
        
        if time.time() - start_time > MAX_TRAIN_TIME:
            print(f"\n   Max time exceeded")
            early_stopped = True
            break
    
    return (epoch_loss, epoch_raw_loss, global_step, best_val_loss,
            best_model_state, patience_counter, early_stopped)


def train(
    mode: str = "mixed",
    min_score: float = None,
    max_epochs: int = None,
    peak_lr: float = None,
    batch_size: int = None,
    grad_accum: int = None,
    mask_mode: str = "basil_and_after",
    lora_max_epochs: int = None,
    max_graded_turns: int = None,
):
    """
    Main training function supporting dual-objective (LoRA) and legacy modes.
    
    Modes:
        "mixed"    : (DEFAULT) Alternating world/basil training with LoRA adapters
        "world"    : Train base model on full transcripts only (no LoRA)
        "basil"    : Train LoRA adapter on Basil turns only (trunk frozen)
        "session"  : Legacy full-session mode (alias for world)
        "graded"   : Legacy graded turns mode
        "legacy"   : Legacy conversation logs mode
        "combined" : Legacy combined mode
    
    mask_mode:
        "none"           : No masking — trunk sees everything
        "basil_only"     : Mask only Basil's output tokens
        "basil_and_after": (DEFAULT) Mask Basil + Sophie post-grade + Tutor answer
    
    lora_max_epochs:
        Manual override for maximum LoRA (BASIL) epochs in mixed mode.
        If None, scales linearly with age_band (0 at band 0, max_epochs at
        band 7) via lora_max_epochs_for_age_band().
    
    max_graded_turns:
        If set, limit training to at most N graded turns (Basil spans).
        Sessions are included chronologically until the limit is reached.
        Useful for quick directional tests on a subset of the data.
    """
    # Map legacy alias
    if mode == "session":
        print("[train] 'session' mode mapped to 'world' (use 'mixed' for dual-objective)")
        mode = "world"
    
    is_dual_objective = mode in ["mixed", "world", "basil"]
    
    max_epochs = max_epochs or MAX_EPOCHS
    batch_size = batch_size or BATCH_SIZE
    grad_accum = grad_accum or ACCUMULATION_STEPS
    
    # Resolve block_size from current age_band
    assessment = load_basil_assessment()
    current_age_band = assessment.get("age_band", 0)
    block_size = BLOCK_SIZE_BY_AGE_BAND.get(current_age_band, 1024)
    
    # --- Staged LoRA activation ---
    # Below LORA_ACTIVATION_AGE_BAND, skip LoRA entirely (absorption phase).
    # Downgrade "mixed" → "world" so the trunk trains alone.
    lora_skipped = False
    if current_age_band < LORA_ACTIVATION_AGE_BAND and mode in ["mixed", "basil"]:
        original_mode = mode
        if mode == "mixed":
            mode = "world"
        elif mode == "basil":
            print(f"\n   LoRA skipped: age_band={current_age_band} < LORA_ACTIVATION_AGE_BAND={LORA_ACTIVATION_AGE_BAND}")
            print(f"   Nothing to train in basil-only mode during absorption phase.")
            return
        lora_skipped = True
        is_dual_objective = mode in ["mixed", "world", "basil"]
        print(f"\n   LoRA skipped: age_band={current_age_band} < LORA_ACTIVATION_AGE_BAND={LORA_ACTIVATION_AGE_BAND}")
        print(f"   Downgrading mode: {original_mode} → {mode} (absorption phase, trunk-only)")
    
    effective_batch = batch_size * grad_accum
    tokens_per_step = effective_batch * block_size
    
    print(f"\n{'='*60}")
    print(f"Bootstrap Basil Training {'(Dual-Objective)' if is_dual_objective else '(Legacy)'}")
    print(f"Mode: {mode}")
    print(f"{'='*60}\n")
    
    # === Device & Precision Setup ===
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    use_bf16 = False
    use_scaler = True
    if torch.cuda.is_available():
        capability = torch.cuda.get_device_capability()
        if capability[0] >= 8:
            use_bf16 = torch.cuda.is_bf16_supported()
            if use_bf16:
                use_scaler = False
                print(f"   Using bf16 (Ampere+ detected: sm_{capability[0]}{capability[1]})")
    
    # === Load Model ===
    model_path = get_basil_model_name()
    print(f"   Loading model from: {model_path}")
    
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    model = AutoModelForCausalLM.from_pretrained(model_path)
    model.to(device)
    print(f"   Device: {device}")
    
    # Enable gradient checkpointing for memory efficiency
    if hasattr(model, 'gradient_checkpointing_enable'):
        model.gradient_checkpointing_enable()
        print(f"   Gradient checkpointing: enabled")
    
    # === LoRA Setup (for dual-objective modes) ===
    has_lora = False
    is_continuation = False
    if is_dual_objective and mode in ["mixed", "basil"]:
        # Check if a pre-existing LoRA adapter exists to resume from
        existing_adapter_path = os.path.join(model_path, LORA_ADAPTER_SUBDIR)
        if os.path.exists(existing_adapter_path):
            is_continuation = True
            print(f"   Loading existing LoRA adapter from {existing_adapter_path}")
            model = PeftModel.from_pretrained(model, existing_adapter_path)
            # Ensure gradient checkpointing stays enabled on the wrapped model
            if hasattr(model, 'gradient_checkpointing_enable'):
                model.gradient_checkpointing_enable()
        else:
            print(f"   Initializing new LoRA adapter")
            model = setup_lora_model(model)
        has_lora = True
    
    # === Resolve Learning Rate ===
    # Use lower LR for continued training unless caller explicitly set peak_lr
    if peak_lr is not None:
        pass  # Caller explicitly set LR, respect it
    elif is_continuation:
        peak_lr = CONTINUE_PEAK_LR
        print(f"   Continuing from trained model — using CONTINUE_PEAK_LR={CONTINUE_PEAK_LR}")
    else:
        peak_lr = PEAK_LR
    
    # LoRA gets a separate (lower) LR — alpha/rank amplifies updates
    lora_lr = LORA_CONTINUE_LR if is_continuation else LORA_PEAK_LR
    
    # === Load Data ===
    print(f"\n   Loading training data...")
    
    if is_dual_objective:
        sessions = load_session_transcripts()
        if not sessions:
            print("   No session data found!")
            return
        
        # Subset sessions by max_graded_turns if specified
        if max_graded_turns is not None:
            total_turns = 0
            subset_sessions = []
            for s in sessions:
                n_spans = len(s.get("basil_spans", []))
                if total_turns + n_spans <= max_graded_turns:
                    subset_sessions.append(s)
                    total_turns += n_spans
                else:
                    remaining = max_graded_turns - total_turns
                    if remaining > 0:
                        truncated = dict(s)
                        truncated["basil_spans"] = s["basil_spans"][:remaining]
                        subset_sessions.append(truncated)
                        total_turns += remaining
                    break
            original_total = sum(len(s.get("basil_spans", [])) for s in sessions)
            print(f"   Subset: {total_turns}/{original_total} graded turns "
                  f"from {len(subset_sessions)}/{len(sessions)} sessions")
            sessions = subset_sessions
        
        world_dataset = None
        basil_dataset = None
        
        if mode in ["mixed", "world"]:
            world_dataset = WorldDataset(sessions, tokenizer, block_size,
                                          mask_mode=mask_mode,
                                          age_band=current_age_band)
            print(f"   World dataset: {len(world_dataset)} chunks")
        
        if mode in ["mixed", "basil"]:
            basil_dataset = BasilDataset(sessions, tokenizer, block_size,
                                         age_band=current_age_band)
            print(f"   Basil dataset: {len(basil_dataset)} examples (age_band={current_age_band})")
        
        if (mode == "world" and (world_dataset is None or len(world_dataset) == 0)):
            print("   No world training data!")
            return
        if (mode == "basil" and (basil_dataset is None or len(basil_dataset) == 0)):
            print("   No basil training data!")
            return
        if (mode == "mixed" and
                (world_dataset is None or len(world_dataset) == 0 or
                 basil_dataset is None or len(basil_dataset) == 0)):
            print("   Need both world and basil data for mixed mode!")
            return
        
        # === Compute adaptive LoRA epoch cap (mixed mode) ===
        adaptive_lora_cap = None
        if mode == "mixed" and basil_dataset is not None:
            if lora_max_epochs is not None:
                adaptive_lora_cap = lora_max_epochs
                print(f"   LoRA epoch cap: {adaptive_lora_cap} (manual override)")
            else:
                adaptive_lora_cap = lora_max_epochs_for_age_band(current_age_band, max_epochs)
                print(f"   LoRA epoch cap: {adaptive_lora_cap} (age_band={current_age_band}, max_epochs={max_epochs})")
        
        # === Split datasets ===
        def _split_dataset(ds, name):
            total = len(ds)
            val_sz = max(3, int(total * VAL_SPLIT_RATIO))
            train_sz = total - val_sz
            train_ds, val_ds = random_split(ds, [train_sz, val_sz],
                                            generator=torch.Generator().manual_seed(42))
            print(f"   {name}: {train_sz} train, {val_sz} val")
            return train_ds, val_ds
        
        if mode in ["mixed", "world"]:
            world_train, world_val = _split_dataset(world_dataset, "World")
            world_train_loader = DataLoader(world_train, batch_size=batch_size, shuffle=True)
            world_val_loader = DataLoader(world_val, batch_size=batch_size, shuffle=False)
        
        if mode in ["mixed", "basil"]:
            basil_train, basil_val = _split_dataset(basil_dataset, "Basil")
            basil_train_loader = DataLoader(basil_train, batch_size=batch_size, shuffle=True)
            basil_val_loader = DataLoader(basil_val, batch_size=batch_size, shuffle=False)
        
        # Determine primary validation loader (for early stopping)
        # World val has far more data (hundreds of chunks vs ~3 Basil examples),
        # giving a stable signal. Basil val is logged separately for monitoring.
        if mode == "world":
            primary_val_loader = world_val_loader
        elif mode == "mixed":
            primary_val_loader = world_val_loader
        else:
            primary_val_loader = basil_val_loader
        
        # === Calculate Steps ===
        if mode == "mixed":
            # In mixed mode, we alternate: odd epochs = world, even epochs = basil
            world_steps_per_epoch = max(1, len(world_train_loader) // grad_accum)
            basil_steps_per_epoch = max(1, len(basil_train_loader) // grad_accum)
            total_steps = (world_steps_per_epoch + basil_steps_per_epoch) * (max_epochs // 2 + 1)
        elif mode == "world":
            total_steps = max(1, len(world_train_loader) // grad_accum) * max_epochs
        else:  # basil
            total_steps = max(1, len(basil_train_loader) // grad_accum) * max_epochs
        
        warmup_steps = int(total_steps * WARMUP_RATIO)
        
        # === Optimizers & Schedulers ===
        if mode == "world":
            # Train all base model params
            optimizer = torch.optim.AdamW(
                [p for p in model.parameters() if p.requires_grad], lr=peak_lr)
            scheduler = get_cosine_schedule_with_warmup(
                optimizer, warmup_steps, total_steps, MIN_LR / peak_lr)
        elif mode == "basil":
            # Train only LoRA params
            freeze_for_basil_training(model)
            lora_params = [p for p in model.parameters() if p.requires_grad]
            optimizer = torch.optim.AdamW(lora_params, lr=lora_lr)
            scheduler = get_cosine_schedule_with_warmup(
                optimizer, warmup_steps, total_steps, MIN_LR / lora_lr)
        else:  # mixed
            # Two separate optimizers
            # World optimizer: trunk params
            freeze_for_world_training(model)
            world_params = [p for p in model.parameters() if p.requires_grad]
            world_optimizer = torch.optim.AdamW(world_params, lr=peak_lr)
            world_scheduler = get_cosine_schedule_with_warmup(
                world_optimizer, warmup_steps // 2, total_steps // 2, MIN_LR / peak_lr)
            
            # Basil optimizer: LoRA params (lower LR — alpha/rank amplifies updates)
            freeze_for_basil_training(model)
            lora_params = [p for p in model.parameters() if p.requires_grad]
            basil_optimizer = torch.optim.AdamW(lora_params, lr=lora_lr)
            basil_scheduler = get_cosine_schedule_with_warmup(
                basil_optimizer, warmup_steps // 2, total_steps // 2, MIN_LR / lora_lr)
        
        scaler = GradScaler() if use_scaler else None
        
        # === Print Config ===
        print(f"\n   Training configuration:")
        print(f"   Microbatch: {batch_size}, Grad accum: {grad_accum}")
        print(f"   Effective batch: {effective_batch} sequences ({tokens_per_step:,} tokens/step)")
        print(f"   Max epochs: {max_epochs}")
        print(f"   Total steps (est): {total_steps}")
        print(f"   Warmup steps: {warmup_steps}")
        print(f"   Trunk Peak LR: {peak_lr}, LoRA Peak LR: {lora_lr}, Min LR: {MIN_LR}")
        print(f"   Precision: {'bf16' if use_bf16 else 'fp16 + GradScaler'}")
        print(f"   LoRA active: {has_lora}")
        if mask_mode != "none":
            print(f"   Mask mode: {mask_mode}")
        if adaptive_lora_cap is not None:
            print(f"   Adaptive LoRA cap: {adaptive_lora_cap} epochs")
        
        # === Training Loop (Dual-Objective) ===
        best_model_state = None
        global_step = 0
        basil_epoch_count = 0  # Track LoRA epochs for adaptive cap
        
        # Per-phase early stopping: WORLD and BASIL train different params
        # (trunk vs LoRA) so they need independent convergence tracking.
        if mode == "mixed":
            world_best_val = float('inf')
            world_patience = 0
            world_stopped = False
            basil_best_val = float('inf')
            basil_patience = 0
            basil_stopped = False
        else:
            best_val_loss = float('inf')
            patience_counter = 0
        
        print(f"\n   Beginning {'alternating ' if mode == 'mixed' else ''}training loop...")
        start_time = time.time()
        
        for epoch in range(1, max_epochs + 1):
            if mode == "mixed":
                is_world_phase = (epoch % 2 == 1)  # Odd = world, even = basil
                
                basil_capped = (adaptive_lora_cap is not None
                                and basil_epoch_count >= adaptive_lora_cap)
                # Skip BASIL if capped or converged -> run WORLD instead
                if not is_world_phase and (basil_stopped or basil_capped):
                    if basil_capped and not basil_stopped:
                        print(f"   [LoRA cap] BASIL skipped (reached cap of {adaptive_lora_cap}), running WORLD")
                    else:
                        print(f"   [BASIL converged] Skipping BASIL, running WORLD")
                    is_world_phase = True
                
                # Skip WORLD if converged -> run BASIL instead (or stop if both done)
                if is_world_phase and world_stopped:
                    basil_capped = (adaptive_lora_cap is not None
                                    and basil_epoch_count >= adaptive_lora_cap)
                    if basil_stopped or basil_capped:
                        print(f"\n   Both phases completed — stopping training")
                        break
                    print(f"   [WORLD converged] Skipping WORLD, running BASIL")
                    is_world_phase = False
            elif mode == "world":
                is_world_phase = True
            else:
                is_world_phase = False
            
            if is_world_phase:
                phase_name = "WORLD"
                freeze_for_world_training(model)
                current_loader = world_train_loader
                current_optimizer = world_optimizer if mode == "mixed" else optimizer
                current_scheduler = world_scheduler if mode == "mixed" else scheduler
            else:
                phase_name = "BASIL"
                freeze_for_basil_training(model)
                basil_epoch_count += 1
                current_loader = basil_train_loader
                current_optimizer = basil_optimizer if mode == "mixed" else optimizer
                current_scheduler = basil_scheduler if mode == "mixed" else scheduler
            
            print(f"\n[Epoch {epoch}/{max_epochs}] Phase: {phase_name}")
            
            # Val loss floor applies only during WORLD phase
            epoch_val_loss_floor = VAL_LOSS_FLOOR_BY_AGE_BAND.get(current_age_band, 1.0) if is_world_phase else None
            
            # Select phase-appropriate validation loader and state
            if mode == "mixed":
                if is_world_phase:
                    phase_val_loader = world_val_loader
                    phase_best_val = world_best_val
                    phase_patience = world_patience
                else:
                    phase_val_loader = basil_val_loader
                    phase_best_val = basil_best_val
                    phase_patience = basil_patience
            else:
                phase_val_loader = primary_val_loader
                phase_best_val = best_val_loss
                phase_patience = patience_counter
            
            (epoch_loss, epoch_raw_loss, global_step, phase_best_val,
             best_model_state, phase_patience, early_stopped) = _run_single_objective_epoch(
                model=model,
                train_loader=current_loader,
                optimizer=current_optimizer,
                scheduler=current_scheduler,
                scaler=scaler,
                device=device,
                grad_accum=grad_accum,
                use_bf16=use_bf16,
                global_step=global_step,
                start_time=start_time,
                val_loader=phase_val_loader,
                best_val_loss=phase_best_val,
                best_model_state=best_model_state,
                patience_counter=phase_patience,
                phase_name=phase_name,
                val_loss_floor=epoch_val_loss_floor,
            )
            
            # Epoch summary
            n_batches = len(current_loader)
            avg_loss = epoch_loss / max(1, n_batches)
            avg_raw = epoch_raw_loss / max(1, n_batches)
            
            # End-of-epoch validation on phase-specific loader
            val_w, val_r = evaluate_validation(model, phase_val_loader, device, use_bf16)
            elapsed = time.time() - start_time
            
            # Cross-validation: log the other phase's val for monitoring.
            # Next epoch's freeze_for_*() restores the correct LoRA state.
            cross_val_str = ""
            if mode == "mixed":
                if is_world_phase and basil_val_loader is not None:
                    if has_lora and hasattr(model, 'enable_adapter_layers'):
                        model.enable_adapter_layers()
                    b_val_w, _ = evaluate_validation(model, basil_val_loader, device, use_bf16)
                    cross_val_str = f" | basil_val={b_val_w:.4f}"
                elif not is_world_phase and world_val_loader is not None:
                    if has_lora and hasattr(model, 'disable_adapter_layers'):
                        model.disable_adapter_layers()
                    w_val_w, _ = evaluate_validation(model, world_val_loader, device, use_bf16)
                    cross_val_str = f" | world_val={w_val_w:.4f}"
            
            # Epoch-level early stopping (phase-specific)
            epoch_improvement = phase_best_val - val_w
            if epoch_improvement > EARLY_STOP_MIN_DELTA:
                phase_best_val = val_w
                best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
                phase_patience = 0
                es_str = "  New best"
            else:
                phase_patience += 1
                es_str = f"  ({phase_patience}/{EARLY_STOP_PATIENCE})"
            
            print(f"[Epoch {epoch}] {phase_name} train={avg_loss:.4f} (raw={avg_raw:.4f}) | "
                  f"val={val_w:.4f} (raw={val_r:.4f}){cross_val_str} | steps={global_step} | "
                  f"time={elapsed:.0f}s{es_str}")
            
            # Write back phase-specific state and check stopping
            if mode == "mixed":
                if is_world_phase:
                    world_best_val = phase_best_val
                    world_patience = phase_patience
                    if early_stopped or world_patience >= EARLY_STOP_PATIENCE:
                        if not world_stopped:
                            print(f"   WORLD converged (patience={world_patience}/{EARLY_STOP_PATIENCE})")
                        world_stopped = True
                else:
                    basil_best_val = phase_best_val
                    basil_patience = phase_patience
                    if early_stopped or basil_patience >= EARLY_STOP_PATIENCE:
                        if not basil_stopped:
                            print(f"   BASIL converged (patience={basil_patience}/{EARLY_STOP_PATIENCE})")
                        basil_stopped = True
                
                basil_capped = (adaptive_lora_cap is not None
                                and basil_epoch_count >= adaptive_lora_cap)
                if world_stopped and (basil_stopped or basil_capped):
                    print(f"\n   All phases completed — stopping at epoch {epoch}")
                    break
            else:
                best_val_loss = phase_best_val
                patience_counter = phase_patience
                if early_stopped or patience_counter >= EARLY_STOP_PATIENCE:
                    print(f"\n   Early stopping at epoch {epoch}")
                    break
            
            # Val loss floor (WORLD phase only): prevent over-training at early stages
            val_loss_floor = VAL_LOSS_FLOOR_BY_AGE_BAND.get(current_age_band, 1.0)
            if is_world_phase and val_w <= val_loss_floor:
                print(f"\n   Val loss floor reached at epoch {epoch}: {val_w:.4f} <= {val_loss_floor} (age_band={current_age_band})")
                if best_model_state is None or val_w < (world_best_val if mode == "mixed" else best_val_loss):
                    if mode == "mixed":
                        world_best_val = val_w
                    else:
                        best_val_loss = val_w
                    best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
                break
            
            if time.time() - start_time > MAX_TRAIN_TIME:
                print(f"\n   Max time exceeded")
                break
    
    else:
        # =====================================================================
        # LEGACY MODES (graded, legacy, combined)
        # =====================================================================
        graded_dataset = None
        legacy_dataset = None
        
        if mode in ["graded", "combined"]:
            graded_turns = load_graded_turns(min_score)
            if graded_turns:
                graded_dataset = WeightedGradedDataset(graded_turns, tokenizer, block_size)
                print(f"   Graded: {len(graded_dataset)} examples")
        
        if mode in ["legacy", "combined"]:
            legacy_logs = load_legacy_logs()
            if legacy_logs:
                apply_legacy_mult = (mode == "combined")
                stride = block_size // 2
                legacy_dataset = LegacyDataset(
                    legacy_logs, tokenizer, block_size, stride,
                    apply_legacy_mult=apply_legacy_mult
                )
                print(f"   Legacy: {len(legacy_dataset)} examples")
        
        if mode == "graded":
            full_dataset = graded_dataset
        elif mode == "legacy":
            full_dataset = legacy_dataset
        else:
            full_dataset = CombinedDataset(graded_dataset, legacy_dataset)
        
        if full_dataset is None or len(full_dataset) == 0:
            print("   No training data found!")
            return
        
        total_size = len(full_dataset)
        val_size = max(5, int(total_size * VAL_SPLIT_RATIO))
        train_size = total_size - val_size
        
        train_dataset, val_dataset = random_split(
            full_dataset, [train_size, val_size],
            generator=torch.Generator().manual_seed(42)
        )
        
        train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
        val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
        
        if len(train_loader) < grad_accum:
            grad_accum = max(1, len(train_loader))
            effective_batch = batch_size * grad_accum
        
        steps_per_epoch = max(1, len(train_loader) // grad_accum)
        total_steps = steps_per_epoch * max_epochs
        warmup_steps = int(total_steps * WARMUP_RATIO)
        
        optimizer = torch.optim.AdamW(model.parameters(), lr=peak_lr)
        scheduler = get_cosine_schedule_with_warmup(
            optimizer, warmup_steps, total_steps, MIN_LR / peak_lr)
        scaler = GradScaler() if use_scaler else None
        
        print(f"   {mode} dataset: {train_size} train, {val_size} val")
        
        best_val_loss = float('inf')
        best_model_state = None
        patience_counter = 0
        global_step = 0
        start_time = time.time()
        
        legacy_val_loss_floor = VAL_LOSS_FLOOR_BY_AGE_BAND.get(current_age_band, 1.0)
        
        for epoch in range(1, max_epochs + 1):
            (epoch_loss, epoch_raw_loss, global_step, best_val_loss,
             best_model_state, patience_counter, early_stopped) = _run_single_objective_epoch(
                model=model,
                train_loader=train_loader,
                optimizer=optimizer,
                scheduler=scheduler,
                scaler=scaler,
                device=device,
                grad_accum=grad_accum,
                use_bf16=use_bf16,
                global_step=global_step,
                start_time=start_time,
                val_loader=val_loader,
                best_val_loss=best_val_loss,
                best_model_state=best_model_state,
                patience_counter=patience_counter,
                phase_name=mode.upper(),
                val_loss_floor=legacy_val_loss_floor,
            )
            
            n_batches = len(train_loader)
            avg_loss = epoch_loss / max(1, n_batches)
            val_w, val_r = evaluate_validation(model, val_loader, device, use_bf16)
            elapsed = time.time() - start_time
            
            epoch_improvement = best_val_loss - val_w
            if epoch_improvement > EARLY_STOP_MIN_DELTA:
                best_val_loss = val_w
                best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
                patience_counter = 0
                es_str = "  New best"
            else:
                patience_counter += 1
                es_str = f"  ({patience_counter}/{EARLY_STOP_PATIENCE})"
            
            print(f"[Epoch {epoch}] train={avg_loss:.4f} | val={val_w:.4f} | "
                  f"steps={global_step} | time={elapsed:.0f}s{es_str}")
            
            if early_stopped or patience_counter >= EARLY_STOP_PATIENCE:
                break
            
            # Val loss floor for legacy modes
            if val_w <= legacy_val_loss_floor:
                print(f"\n   Val loss floor reached at epoch {epoch}: {val_w:.4f} <= {legacy_val_loss_floor} (age_band={current_age_band})")
                if best_model_state is None or val_w < best_val_loss:
                    best_val_loss = val_w
                    best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
                break
            
            if time.time() - start_time > MAX_TRAIN_TIME:
                break
    
    # =========================================================================
    # COMMON: Restore best, save, update metrics
    # =========================================================================
    if best_model_state is not None:
        if mode == "mixed":
            print(f"\n   Restoring best checkpoint (world_val={world_best_val:.4f}, basil_val={basil_best_val:.4f})")
        else:
            print(f"\n   Restoring best checkpoint (val_loss={best_val_loss:.4f})")
        model.load_state_dict({k: v.to(device) for k, v in best_model_state.items()})
    else:
        print(f"\n   No checkpoint saved (using final model state)")
    
    # === Save Model ===
    existing_versions = sorted([
        int(name.split("_v")[-1])
        for name in os.listdir(MODELS_DIR)
        if name.startswith("basil_v") and name.split("_v")[-1].isdigit()
    ])
    next_version = max(existing_versions, default=1) + 1
    new_model_dir = os.path.join(MODELS_DIR, f"basil_v{next_version:03}")
    os.makedirs(new_model_dir, exist_ok=True)
    
    if has_lora:
        # Save LoRA adapter first (PeftModel.save_pretrained saves only adapter weights)
        adapter_dir = os.path.join(new_model_dir, LORA_ADAPTER_SUBDIR)
        model.save_pretrained(adapter_dir)
        print(f"   LoRA adapter saved to {adapter_dir}")
        
        # Save clean base model (trunk) without LoRA artifacts.
        # PEFT replaces Linear layers with LoRA wrappers that rename keys:
        #   .weight -> .base_layer.weight, .bias -> .base_layer.bias
        # We must reverse this mapping to produce a valid GPT-2 checkpoint.
        clean_state = {}
        for key, param in model.base_model.model.state_dict().items():
            if 'lora_' in key:
                continue  # Skip LoRA adapter weights (saved separately above)
            clean_key = key.replace('.base_layer', '')
            clean_state[clean_key] = param
        
        # Create a vanilla model from config and load the extracted trunk weights
        base_config = model.base_model.model.config
        clean_model = AutoModelForCausalLM.from_config(base_config)
        missing, unexpected = clean_model.load_state_dict(clean_state, strict=False)
        if unexpected:
            print(f"   WARNING: Unexpected keys in clean base state: {unexpected[:5]}...")
        if missing:
            print(f"   WARNING: Missing keys in clean base state: {missing[:5]}...")
        clean_model.save_pretrained(new_model_dir)
        tokenizer.save_pretrained(new_model_dir)
        print(f"   Base model saved to {new_model_dir}")
    else:
        model.save_pretrained(new_model_dir)
        tokenizer.save_pretrained(new_model_dir)
        print(f"   Model saved to {new_model_dir}")
    
    # === Update Metrics ===
    update_metrics_after_training()
    
    from metrics_manager import get_recent_session_metrics
    recent_sessions = get_recent_session_metrics(1, training_phase="normal")
    trigger_session_id = recent_sessions[0].get("session_id") if recent_sessions else None
    record_training_attempt("kept", session_id=trigger_session_id)
    
    elapsed = time.time() - start_time
    print(f"\n{'='*60}")
    print(f"Training complete!")
    print(f"  Mode: {mode}")
    if mode == "mixed":
        print(f"  World best val: {world_best_val:.4f}")
        print(f"  Basil best val: {basil_best_val:.4f}")
        print(f"  BASIL epochs: {basil_epoch_count}")
    else:
        print(f"  Best val_loss: {best_val_loss:.4f}")
    print(f"  Total steps: {global_step}")
    print(f"  Epochs: {epoch}")
    print(f"  Time: {elapsed:.0f}s")
    print(f"  LoRA: {'yes' if has_lora else 'no'}")
    print(f"  New model: {new_model_dir}")
    print(f"{'='*60}\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Train Basil with dual-objective LoRA architecture",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Default: alternating world/basil with LoRA (recommended)
  python train_basil_v2.py

  # World objective only (train trunk, no LoRA)
  python train_basil_v2.py --mode world

  # Basil-policy only (train LoRA, freeze trunk)
  python train_basil_v2.py --mode basil

  # Legacy modes (backward compatibility)
  python train_basil_v2.py --mode graded
  python train_basil_v2.py --mode combined

  # A100 40GB (smaller microbatch)
  python train_basil_v2.py --batch-size 2 --grad-accum 16
"""
    )
    parser.add_argument(
        "--mode",
        choices=["mixed", "world", "basil", "session", "graded", "legacy", "combined"],
        default="mixed",
        help="Training mode (default: mixed = alternating world/basil with LoRA)"
    )
    parser.add_argument(
        "--min-score",
        type=float,
        default=None,
        help="Minimum score for graded examples (legacy modes only)"
    )
    parser.add_argument(
        "--epochs",
        type=int,
        default=None,
        help=f"Override max epochs (default: {MAX_EPOCHS})"
    )
    parser.add_argument(
        "--peak-lr",
        type=float,
        default=None,
        help=f"Peak learning rate (default: {PEAK_LR})"
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=None,
        help=f"Microbatch size (default: {BATCH_SIZE})"
    )
    parser.add_argument(
        "--grad-accum",
        type=int,
        default=None,
        help=f"Gradient accumulation steps (default: {ACCUMULATION_STEPS})"
    )
    parser.add_argument(
        "--mask-mode",
        choices=["none", "basil_only", "basil_and_after"],
        default="basil_and_after",
        help="Trunk masking mode (default: basil_and_after)"
    )
    parser.add_argument(
        "--lora-max-epochs",
        type=int,
        default=None,
        help="Manual override for max LoRA epochs in mixed mode (default: scales with age_band)"
    )
    parser.add_argument(
        "--max-graded-turns",
        type=int,
        default=None,
        help="Limit training to N graded turns (for quick directional tests)"
    )
    args = parser.parse_args()
    
    train(
        mode=args.mode,
        min_score=args.min_score,
        max_epochs=args.epochs,
        peak_lr=args.peak_lr,
        batch_size=args.batch_size,
        grad_accum=args.grad_accum,
        mask_mode=args.mask_mode,
        lora_max_epochs=args.lora_max_epochs,
        max_graded_turns=args.max_graded_turns,
    )