online_learner.py

"""
online_learner.py — Phase 5: Online LoRA Learning Demo

Simulates 1000 tasks (50 unique questions, cycled). For each task:
  1. Encode canonical prompt [BOS,USER]+clean(q)+[ASST] → 512-dim embedding
  2. Retrieve top-3 similar adapters from memory → merge
  3. Generate answer with merged adapter (greedy)
  4. Score against ground truth (token F1)
  5. If score < ATTEMPT_THRESHOLD (not yet memorized): Adam-overfit a per-task
     adapter; store it (keyed by the unperturbed task embedding) only if its
     OWN post-update greedy decode reproduces the answer (>= STORE_THRESHOLD)
  6. Every EVAL_EVERY tasks: run full 50-task benchmark

Target: benchmark F1 increases monotonically as recurring questions are
memorized and recalled — proof of online EPISODIC learning (recall of seen
tasks). The 42M base model cannot answer these questions zero-shot, so this
is NOT a generalization claim.

Runs on server CPU. No GPU required.
Usage:
  python3 online_learner.py [--ckpt path/to/checkpoint.pt] [--tasks 1000]
"""
import sys, os, time, json, random, argparse
from pathlib import Path
from collections import deque

import torch
import torch.nn.functional as F
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download

from lora import LoRAAdapter, KaizenWithLoRA
from task_memory import TaskMemory
from eval_benchmark import (evaluate, BENCHMARK_TASKS, token_f1, generate,
                             build_prompt_ids, clean_ids)

# ── Config ──────────────────────────────────────────────────────────────────
HF_TOKEN           = os.environ.get('HF_TOKEN', '')
HF_MODEL_REPO      = 'qoa/kaizen-42m'
HF_TOK_REPO        = 'qoa/kaizen-tokenizer'
STORE_DIR          = os.path.join(os.path.expanduser('~'), '.kaizen', 'memory')
LOG_PATH           = os.path.join(os.path.expanduser('~'), '.kaizen', 'logs', 'online.log')
RESULTS_PATH       = os.path.join(os.path.expanduser('~'), '.kaizen', 'logs', 'results.json')

LORA_RANK          = 4
LORA_ALPHA         = 32.0      # MUST equal TaskMemory.ALPHA — scale=alpha/rank
                                # is a runtime attr, not persisted in state_dict.
ONLINE_LR          = 1e-2      # Adam lr. Proven (in-memory) to memorize a single
                                # (Q,A) on this rank-4/alpha-32 adapter within
                                # ~20 steps (loss 0.09-0.5, F1=1.0 on Paris/Bell/221).
                                # SGD @ 3 steps cannot overfit at all (loss 4-7,
                                # empty output, F1=0) — replaced.
ONLINE_STEPS       = 25        # max Adam steps per task.
LOSS_EARLY_STOP    = 0.1       # stop early once teacher-forced loss < this —
                                # 20 steps was the sweet spot in proof; 40 steps
                                # over/under-shot some multi-token answers.
ATTEMPT_THRESHOLD  = 0.99      # skip the update entirely if the (retrieved-
                                # adapter) greedy answer already scores >= this —
                                # the task is already memorized, no need to retrain.
STORE_THRESHOLD    = 0.5       # only store the updated adapter if its OWN
                                # post-update greedy decode reproduces the answer
                                # at >= this F1. Guarantees every stored adapter
                                # is useful on recall (clean monotonic curve).
EVAL_EVERY         = 100       # tasks between benchmark evals
TOP_K              = 3         # adapters to retrieve and merge
BLOCK_SIZE         = 1024
MAX_GEN            = 80
EOS_ID             = 3

# ── Logging ──────────────────────────────────────────────────────────────────
Path(LOG_PATH).parent.mkdir(parents=True, exist_ok=True)

def log(msg: str):
    line = f'[{time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}] {msg}'
    print(line, flush=True)
    with open(LOG_PATH, 'a') as f:
        f.write(line + '\n')


# ── Task stream: rotating 3 domains from benchmark ──────────────────────────
def task_stream(n_tasks: int, seed: int = 99):
    """
    Yield (question, answer, task_type) for n_tasks rounds.
    Extends BENCHMARK_TASKS with repetition + minor perturbations for longer runs.
    """
    rng   = random.Random(seed)
    pool  = list(BENCHMARK_TASKS)
    for i in range(n_tasks):
        task = pool[i % len(pool)]
        yield task


# ── Build training sequence for online update ────────────────────────────────
def build_update_seq(tokenizer, question: str, answer: str):
    """
    Build (x, labels) for online update loss, using the SAME canonical prompt
    (build_prompt_ids) as generation/embedding. Format:
      x      = [BOS, USER] + clean(Q) + [ASST] + clean(A) + [EOS]
      labels = [-100]*(len(prompt)-1) + clean(A) + [EOS]   (predict answer+EOS only)
    """
    prompt_ids = build_prompt_ids(tokenizer, question)[:200]
    answer_ids = clean_ids(tokenizer, answer)[:150] + [EOS_ID]

    full   = prompt_ids + answer_ids          # length L
    x      = full[:-1]                        # length L-1
    labels = [-100] * (len(prompt_ids) - 1) + answer_ids   # length L-1, predicts x shifted by 1

    x      = x[:BLOCK_SIZE]
    labels = labels[:BLOCK_SIZE]
    return (torch.tensor(x, dtype=torch.long).unsqueeze(0),
            torch.tensor(labels, dtype=torch.long).unsqueeze(0))


# ── Online update ────────────────────────────────────────────────────────────
def online_update(model: KaizenWithLoRA, adapter: LoRAAdapter,
                  x: torch.Tensor, labels: torch.Tensor) -> float:
    """
    Adam update on adapter params only. Base model frozen.
    Stops early once teacher-forced loss < LOSS_EARLY_STOP (proven sweet spot:
    ~20 steps memorizes a single (Q,A) on this rank-4/alpha-32 adapter without
    over/under-shooting multi-token answers — see proof in plan).
    Returns final loss value.
    """
    for p in model.parameters():
        p.requires_grad_(False)
    for p in adapter.parameters():
        p.requires_grad_(True)

    optimizer = torch.optim.Adam(adapter.parameters(), lr=ONLINE_LR)
    last_loss = float('inf')

    for _ in range(ONLINE_STEPS):
        optimizer.zero_grad()
        _, loss = model(x, targets=labels, adapter=adapter)
        loss.backward()
        optimizer.step()
        last_loss = loss.item()
        if last_loss < LOSS_EARLY_STOP:
            break

    return last_loss


# ── Main loop ─────────────────────────────────────────────────────────────────
def run(ckpt_path: str, n_tasks: int = 1000):
    log(f'=== Phase 5 Online Learning — {n_tasks} tasks ===')
    log(f'Checkpoint: {ckpt_path}')
    log(f'Store: {STORE_DIR}')

    # Load tokenizer
    tok_file = hf_hub_download(HF_TOK_REPO, 'tokenizer.json',
                               token=HF_TOKEN, cache_dir=None)
    tokenizer = Tokenizer.from_file(tok_file)
    log(f'Tokenizer loaded: vocab={tokenizer.get_vocab_size()}')

    # Load model (frozen base)
    model = KaizenWithLoRA()
    model.load_base(ckpt_path)
    model.eval()
    total_base  = sum(p.numel() for p in model.parameters())
    log(f'Base model: {total_base:,} params (all frozen)')

    # Task memory
    memory = TaskMemory(STORE_DIR, top_k=TOP_K)
    log(f'Memory: {len(memory)} existing tasks loaded')

    # Baseline eval (no adapters)
    log('Running baseline benchmark (task=0, no adapters)...')
    baseline = evaluate(model, tokenizer, memory=None)
    log(f'Baseline: F1={baseline["overall_f1"]:.4f} | '
        f'factual={baseline["factual_f1"]:.4f} | '
        f'math={baseline["math_f1"]:.4f} | '
        f'commonsense={baseline["commonsense_f1"]:.4f}')

    eval_checkpoints = [0]
    eval_results     = [baseline]
    task_scores      = deque(maxlen=100)   # rolling window
    n_updated        = 0

    t0 = time.time()

    for task_idx, (question, answer, task_type) in enumerate(task_stream(n_tasks)):
        # 1. Embed task — canonical prompt P = [BOS,USER]+clean(q)+[ASST],
        #    SAME format used by build_update_seq and eval_benchmark.evaluate.
        prompt_ids = build_prompt_ids(tokenizer, question)[:BLOCK_SIZE - MAX_GEN]
        x_prompt   = torch.tensor([prompt_ids], dtype=torch.long)
        ref_ids    = clean_ids(tokenizer, answer)

        with torch.no_grad():
            task_emb = model.embed_task(x_prompt, adapter=None)

        # 2. Retrieve + merge adapters
        adapter = memory.retrieve_merged(task_emb)

        # 3. Generate
        with torch.no_grad():
            gen_ids = generate(model, tokenizer, prompt_ids, adapter=adapter)

        score = token_f1(gen_ids, ref_ids)
        task_scores.append(score)

        if (task_idx + 1) % 10 == 0:
            avg10 = sum(list(task_scores)[-10:]) / min(10, len(task_scores))
            elapsed = time.time() - t0
            log(f'task {task_idx+1:>5}/{n_tasks} | '
                f'score {score:.3f} | avg10 {avg10:.3f} | '
                f'memory {len(memory)} | updates {n_updated} | '
                f'{elapsed:.0f}s')

        # 4. Online update — skip if already memorized (score >= ATTEMPT_THRESHOLD).
        #    Otherwise overfit a per-task adapter, then store it ONLY if its own
        #    post-update greedy decode reproduces the answer (>= STORE_THRESHOLD).
        #    Retrieval key = unperturbed task_emb (same as query-time embedding).
        if score < ATTEMPT_THRESHOLD:
            new_adapter = LoRAAdapter(model.N_LAYERS, model.D_MODEL,
                                      LORA_RANK, LORA_ALPHA)
            if adapter is not None:
                new_adapter.load_state_dict(adapter.state_dict())

            x_upd, y_upd = build_update_seq(tokenizer, question, answer)
            update_loss  = online_update(model, new_adapter, x_upd, y_upd)

            with torch.no_grad():
                post_gen_ids = generate(model, tokenizer, prompt_ids, adapter=new_adapter)
            post_score = token_f1(post_gen_ids, ref_ids)

            if post_score >= STORE_THRESHOLD:
                memory.add(task_emb, new_adapter, {
                    'task_type':     task_type,
                    'question':      question[:100],
                    'pre_score':     score,
                    'post_score':    post_score,
                    'update_loss':   update_loss,
                    'task_idx':      task_idx,
                })
                n_updated += 1

        # 5. Benchmark eval
        if (task_idx + 1) % EVAL_EVERY == 0:
            log(f'--- Benchmark at task {task_idx+1} ---')
            metrics = evaluate(model, tokenizer, memory=memory)
            log(f'  F1={metrics["overall_f1"]:.4f} | '
                f'factual={metrics["factual_f1"]:.4f} | '
                f'math={metrics["math_f1"]:.4f} | '
                f'commonsense={metrics["commonsense_f1"]:.4f}')

            prev_f1 = eval_results[-1]['overall_f1']
            delta   = metrics['overall_f1'] - prev_f1
            trend   = '↑' if delta > 0 else ('↓' if delta < 0 else '→')
            log(f'  Δ={delta:+.4f} {trend}')

            eval_checkpoints.append(task_idx + 1)
            eval_results.append(metrics)

    # Final summary
    total_time = time.time() - t0
    log(f'\n=== Phase 5 complete ===')
    log(f'Tasks: {n_tasks} | Updates: {n_updated} | Time: {total_time/60:.1f}min')

    f1_values = [r['overall_f1'] for r in eval_results]
    monotonic = all(f1_values[i] <= f1_values[i+1] for i in range(len(f1_values)-1))
    log(f'F1 trajectory: {" → ".join(f"{v:.4f}" for v in f1_values)}')
    log(f'Monotonically increasing: {monotonic}')

    memory.flush()   # persist any remaining unsaved tasks

    results = {
        'eval_checkpoints': eval_checkpoints,
        'eval_results':     eval_results,
        'n_tasks':          n_tasks,
        'n_updated':        n_updated,
        'total_time_s':     total_time,
        'monotonic':        monotonic,
        'ckpt_path':        ckpt_path,
    }
    Path(RESULTS_PATH).write_text(json.dumps(results, indent=2))
    log(f'Results saved to {RESULTS_PATH}')
    return results


# ── CLI ──────────────────────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--ckpt', type=str, default=None,
                        help='Local path to checkpoint (.pt). '
                             'If not given, downloads phase4_latest.pt (or phase2_latest.pt).')
    parser.add_argument('--tasks', type=int, default=1000)
    args = parser.parse_args()

    if args.ckpt:
        ckpt_path = args.ckpt
    else:
        # Try phase4 first, fall back to phase3, then phase2
        for ckpt_name in ('phase4_latest.pt', 'phase3_latest.pt', 'phase2_latest.pt'):
            try:
                ckpt_path = hf_hub_download(
                    HF_MODEL_REPO, ckpt_name,
                    token=HF_TOKEN, cache_dir=None,
                )
                log(f'Using checkpoint: {ckpt_name}')
                break
            except Exception:
                continue
        else:
            raise RuntimeError('No checkpoint found on HF (phase2/3/4_latest.pt)')

    run(ckpt_path, n_tasks=args.tasks)


if __name__ == '__main__':
    main()