""" 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()