Add prior_run/train.py (API key redacted)

prior_run/train.py

@@ -0,0 +1,503 @@
+# NOTE: this is a copy of prior_run/train.py with a hard-coded OpenRouter API key
+# redacted before publication. The original file is unmodified on disk.
+"""
+GRPO + Diversity via LLM Judge — v2 (GDPO multi-reward)
+--------------------------------------------------------
+Three separate reward functions, decoupled-normalized then summed:
+  1. quality_reward       = quality_i + 0.5 * group_diversity
+  2. novelty_reward       = novelty_i if quality_i >= 5 else 0  (per GDPO §3.2)
+  3. completeness_reward  = deterministic, did the story actually end?
+
+The two judge-driven rewards share ONE judge call per group via a small cache,
+so we don't pay for the same scoring twice.
+
+Usage:
+  export OPENROUTER_API_KEY="sk-or-..."
+  python train.py
+"""
+
+import os
+import json
+import statistics
+import requests
+import wandb
+from datasets import load_dataset
+from transformers import TrainerCallback
+from trl import GRPOConfig, GRPOTrainer
+from peft import LoraConfig
+
+# ── Config ──────────────────────────────────────────────────────────
+BASE_MODEL = "Qwen/Qwen3-4B-Instruct-2507"
+JUDGE_MODEL = "google/gemini-3-flash-preview"
+OPENROUTER_API_KEY = "sk-or-v1-REDACTED-BEFORE-UPLOAD"
+
+NUM_GENERATIONS = 16          # group size — the diversity sample pool
+GROUP_DIV_BONUS = 0.5         # weight on group_diversity inside quality reward
+QUALITY_GATE = 5.0            # novelty only counts if quality >= this (GDPO conditioning)
+MAX_COMPLETION_TOKENS = 4096   # give stories room to actually end
+LEARNING_RATE = 5e-6
+BATCH_SIZE = 1                # prompts per micro-batch
+GRAD_ACCUM = 2                # → optimizer step every 2 prompts (~250 steps over run)
+NUM_TRAIN_EPOCHS = 2
+TRAIN_SIZE = 250
+EVAL_SIZE = 5
+OUTPUT_DIR = "./grpo-diverse"
+
+# LoRA
+LORA_R = 128
+LORA_ALPHA = 256
+LORA_TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj"]
+
+
+# ── Judge prompt + call ─────────────────────────────────────────────
+JUDGE_SYSTEM_PROMPT = """You are evaluating creative writing responses to the same writing prompt.
+
+## Scoring Instructions
+
+For EACH sample, score:
+
+### quality (1-10)
+- 1-3: Incoherent, nonsensical, or unfinished. Sentences contradict each other. Plot makes no sense. Abrupt cutoff mid-sentence.
+- 4-5: Readable but bland or generic. Makes sense but nothing memorable.
+- 6-7: Well-written, engaging, coherent from start to finish.
+- 8-10: Genuinely compelling. Strong voice, clear arc, emotional impact.
+
+IMPORTANT quality penalties:
+- If the story cuts off mid-sentence or clearly doesn't end, cap quality at 4.
+- If the story is nonsensical (events don't follow logically, characters act without motivation), cap quality at 3.
+- If the tone MISMATCHES the prompt (e.g., a clearly humorous/absurd prompt is answered with a completely serious, solemn story), subtract 2 points.
+
+### novelty (1-10)
+How DIFFERENT is this sample from all other samples in this set?
+
+Rate HIGH (7-10) if this sample uses a different:
+- Genre or tone (comedy vs horror vs literary fiction vs thriller vs satire)
+- Narrative structure (dialogue-driven vs prose vs epistolary vs journal vs screenplay)
+- Point of view (first person vs third person vs second person)
+- Emotional register (funny vs melancholic vs tense vs whimsical vs deadpan)
+
+Rate LOW (1-3) if this sample:
+- Uses the same genre/tone as most other samples
+- Opens with a similar sentence structure
+- Relies on the same rhetorical patterns as other samples
+
+IMPORTANT novelty penalties:
+- The pattern "It wasn't X, it was Y" or "He didn't X. He *Y*-ed" is a specific rhetorical device. If more than 2 samples use this negation-then-reframe pattern, score those samples 1-3 for novelty.
+- If all samples use heavy italic emphasis (*word*) for the same dramatic effect, this is NOT diversity. Note it and reduce novelty scores.
+- Surface-level changes (different character names, different settings) with the same tone and structure = LOW novelty.
+
+### group_diversity (1-10)
+Does this SET as a whole cover genuinely different creative territory?
+
+- 1-3: All samples share the same tone, genre, and emotional register. Even if plots differ, they all "feel" the same to read. All are dark literary fantasy, or all are solemn and atmospheric, or all are action thrillers.
+- 4-5: Some structural variety (different openings, different plot directions) but the tone/mood is largely uniform. A reader would notice they all sound like the same author in the same mood.
+- 6-7: Clear variety in at least TWO of: genre, tone, structure, POV. For example, one sample is humorous while others are serious, or one is dialogue-heavy while others are prose.
+- 8-10: Multiple genuinely distinct approaches. The set includes different genres (comedy AND drama AND thriller), different structures (prose AND dialogue AND epistolary), different emotional registers (funny AND melancholic AND tense). A reader would think these came from different authors.
+
+CRITICAL: A set of 8 dark atmospheric fantasy stories with different plots is a 2-3, not a 6. Diversity means TONAL and STRUCTURAL variety, not just plot variety.
+
+Respond with ONLY valid JSON, no markdown fences:
+{"scores": [{"quality": N, "novelty": N}, ...], "group_diversity": N}"""
+
+
+def call_judge(prompt: str, completions: list[str]) -> dict:
+    """Send all completions to the judge in one API call."""
+    samples = ""
+    for i, c in enumerate(completions):
+        samples += f"\n--- Sample {i+1} ---\n{c}\n"
+
+    user_prompt = f"""You are evaluating {len(completions)} creative writing responses to the same writing prompt.
+
+WRITING PROMPT:
+{prompt}
+
+SAMPLES:{samples}
+
+Respond with ONLY valid JSON, no markdown fences:
+{{"scores": [{{"quality": N, "novelty": N}}, ...], "group_diversity": N}}"""
+
+    resp = requests.post(
+        "https://openrouter.ai/api/v1/chat/completions",
+        headers={
+            "Authorization": f"Bearer {OPENROUTER_API_KEY}",
+            "Content-Type": "application/json",
+        },
+        json={
+            "model": JUDGE_MODEL,
+            "messages": [
+                {"role": "system", "content": JUDGE_SYSTEM_PROMPT},
+                {"role": "user", "content": user_prompt},
+            ],
+            "temperature": 0.2,
+        },
+        timeout=120,
+    )
+    resp.raise_for_status()
+
+    text = resp.json()["choices"][0]["message"]["content"].strip()
+    if text.startswith("```"):
+        text = text.split("\n", 1)[1].rsplit("```", 1)[0].strip()
+    return json.loads(text)
+
+
+# ── Judge cache: one call per group, shared across reward fns ──────
+_judge_cache: dict = {}
+_judge_cache_order: list = []
+_JUDGE_CACHE_MAX = 32
+
+
+def cached_judge(prompt: str, completions: list[str]) -> dict:
+    """Memoized judge call. quality_reward and novelty_reward both hit this
+    with the same (prompt, completions) tuple within a step → only one HTTP
+    request actually fires. Returns a neutral fallback on failure."""
+    key = hash((prompt, tuple(completions)))
+    if key in _judge_cache:
+        return _judge_cache[key]
+    try:
+        result = call_judge(prompt, completions)
+    except Exception as e:
+        print(f"  [judge error] {e}")
+        result = {
+            "scores": [{"quality": 5, "novelty": 5} for _ in completions],
+            "group_diversity": 5,
+            "_error": str(e),
+        }
+    _judge_cache[key] = result
+    _judge_cache_order.append(key)
+    while len(_judge_cache_order) > _JUDGE_CACHE_MAX:
+        _judge_cache.pop(_judge_cache_order.pop(0), None)
+    return result
+
+
+# ── Helpers ─────────────────────────────────────────────────────────
+def _extract_text(c) -> str:
+    if isinstance(c, list):
+        return c[-1]["content"] if c else ""
+    if isinstance(c, dict):
+        return c.get("content", str(c))
+    return str(c)
+
+
+def _normalize_prompt(p) -> str:
+    if isinstance(p, list):
+        return p[-1]["content"]
+    return str(p)
+
+
+def _group_by_prompt(prompts: list[str]) -> dict:
+    groups: dict = {}
+    for i, p in enumerate(prompts):
+        groups.setdefault(p, []).append(i)
+    return groups
+
+
+def _normalize_prompts_kwarg(completions, kwargs) -> list[str]:
+    raw = kwargs.get("prompts")
+    if raw is None:
+        return ["unknown"] * len(completions)
+    return [_normalize_prompt(p) for p in raw]
+
+
+# ── Reward functions ───────────────────────────────────────────────
+def quality_reward(completions, **kwargs) -> list[float]:
+    """quality_i + GROUP_DIV_BONUS * group_diversity (per sample)."""
+    texts = [_extract_text(c) for c in completions]
+    prompts = _normalize_prompts_kwarg(completions, kwargs)
+
+    rewards = [0.0] * len(texts)
+    qualities = []
+    group_divs = []
+
+    for prompt_str, indices in _group_by_prompt(prompts).items():
+        group_texts = [texts[i] for i in indices]
+        result = cached_judge(prompt_str, group_texts)
+        scores = result["scores"]
+        gd = float(result.get("group_diversity", 5))
+        group_divs.append(gd)
+
+        for j, idx in enumerate(indices):
+            q = float(scores[j].get("quality", 5))
+            qualities.append(q)
+            rewards[idx] = q + GROUP_DIV_BONUS * gd
+
+    if wandb.run is not None and qualities:
+        wandb.log({
+            "judge/avg_quality": statistics.fmean(qualities),
+            "judge/quality_std": statistics.pstdev(qualities) if len(qualities) > 1 else 0.0,
+            "judge/avg_group_diversity": statistics.fmean(group_divs),
+            "judge/group_diversity_std": statistics.pstdev(group_divs) if len(group_divs) > 1 else 0.0,
+        })
+
+    print(f"  [quality] mean={statistics.fmean(rewards):.2f} | gd_avg={statistics.fmean(group_divs):.2f}")
+    return rewards
+
+
+def novelty_reward(completions, **kwargs) -> list[float]:
+    """novelty_i if quality_i >= QUALITY_GATE else 0.
+
+    Conditioning kills the gaming strategy where the model produces
+    incoherent-but-different outputs to farm the novelty signal.
+    """
+    texts = [_extract_text(c) for c in completions]
+    prompts = _normalize_prompts_kwarg(completions, kwargs)
+
+    rewards = [0.0] * len(texts)
+    novelties = []
+    pass_gate = 0
+
+    for prompt_str, indices in _group_by_prompt(prompts).items():
+        group_texts = [texts[i] for i in indices]
+        result = cached_judge(prompt_str, group_texts)
+        scores = result["scores"]
+
+        for j, idx in enumerate(indices):
+            q = float(scores[j].get("quality", 5))
+            # Gemini sometimes truncates "novelty" to "novel"
+            n = float(scores[j].get("novelty") or scores[j].get("novel") or 5)
+            novelties.append(n)
+            if q >= QUALITY_GATE:
+                rewards[idx] = n
+                pass_gate += 1
+            else:
+                rewards[idx] = 0.0
+
+    if wandb.run is not None and novelties:
+        wandb.log({
+            "judge/avg_novelty": statistics.fmean(novelties),
+            "judge/novelty_std": statistics.pstdev(novelties) if len(novelties) > 1 else 0.0,
+            "judge/novelty_conditioned_pct": pass_gate / len(novelties),
+        })
+    return rewards
+
+
+def completeness_reward(completions, **kwargs) -> list[float]:
+    """Deterministic shaping: did the story actually end?
+
+    1.0 if it ends with a sentence terminator, 0.3 if it produced
+    substantial text but trails off, 0.0 otherwise. Cheap, zero
+    judge-noise variance, directly addresses truncation.
+    """
+    enders = ".!?\"'"
+    rewards = []
+    for c in completions:
+        text = _extract_text(c).strip()
+        if not text:
+            rewards.append(0.0)
+        elif text[-1] in enders:
+            rewards.append(1.0)
+        elif len(text) > 100:
+            rewards.append(0.3)
+        else:
+            rewards.append(0.0)
+
+    if wandb.run is not None and rewards:
+        complete_rate = sum(1 for r in rewards if r >= 1.0) / len(rewards)
+        wandb.log({
+            "judge/avg_completeness": statistics.fmean(rewards),
+            "judge/completeness_rate": complete_rate,
+        })
+    return rewards
+
+
+# ── Dataset ─────────────────────────────────────────────────────────
+def load_writing_prompts():
+    """
+    250 train + 5 eval prompts.
+
+    TODO (manual): audit these for tonal variety. Aim for:
+      - >=20% clearly comedic / absurd / lighthearted
+      - >=20% constrained ("in exactly 3 sentences", "as a diary entry", etc.)
+    The judge has more surface area to reward diversity when the prompt
+    pool itself spans tones and constraints.
+    """
+    ds = load_dataset("euclaise/writingprompts", split="train")
+
+    def format_prompt(example):
+        # NEVER truncate. The judge needs the full prompt for context, and
+        # the model needs to see exactly what we'll be scoring against.
+        text = (example.get("prompt") or example.get("text") or "").strip()
+        return {
+            "prompt": [
+                {"role": "user", "content": f"Write a short creative story based on this prompt:\n\n{text}"}
+            ]
+        }
+
+    total = TRAIN_SIZE + EVAL_SIZE
+    ds = ds.map(format_prompt, remove_columns=ds.column_names)
+    ds = ds.shuffle(seed=42).select(range(total))
+    eval_ds = ds.select(range(TRAIN_SIZE, total))
+    train_ds = ds.select(range(TRAIN_SIZE))
+    return train_ds, eval_ds
+
+
+# ── Eval-side metrics (compute, log, but don't reward on these) ────
+def _unique_5word_openers(samples: list[str]) -> int:
+    openers = set()
+    for s in samples:
+        words = s.strip().split()[:5]
+        if words:
+            openers.add(" ".join(w.lower() for w in words))
+    return len(openers)
+
+
+def _vocab_diversity(samples: list[str]) -> float:
+    """Type-token ratio across all samples in a group."""
+    tokens = []
+    for s in samples:
+        tokens.extend(w.lower().strip(".,!?\"'") for w in s.split())
+    if not tokens:
+        return 0.0
+    return len(set(tokens)) / len(tokens)
+
+
+# ── Eval callback ──────────────────────────────────────────────────
+class EvalSampleCallback(TrainerCallback):
+    """Every 50 steps: generate samples on held-out prompts, judge them
+    once per prompt, and dump everything (samples + raw judge response) to
+    disk for offline inspection."""
+
+    def __init__(self, eval_prompts, tokenizer, gen_count=8):
+        self.eval_prompts = eval_prompts
+        self.tokenizer = tokenizer
+        self.gen_count = gen_count
+
+    def on_step_end(self, args, state, control, model=None, **kwargs):
+        if state.global_step == 0 or state.global_step % 50 != 0:
+            return
+        if model is None:
+            return
+
+        import torch
+        model.eval()
+        step = state.global_step
+        print(f"\n{'='*60}\n  EVAL SAMPLES @ step {step}\n{'='*60}")
+
+        all_data = {}
+        all_unique_openers = []
+        all_vocab_div = []
+
+        for pi, prompt_msgs in enumerate(self.eval_prompts):
+            prompt_text = prompt_msgs[-1]["content"]
+            print(f"\n── Prompt {pi+1}: {prompt_text[:100]}...")
+
+            input_ids = self.tokenizer.apply_chat_template(
+                prompt_msgs, add_generation_prompt=True, return_tensors="pt"
+            ).to(model.device)
+
+            samples = []
+            for si in range(self.gen_count):
+                with torch.no_grad():
+                    out = model.generate(
+                        input_ids,
+                        max_new_tokens=MAX_COMPLETION_TOKENS,
+                        temperature=1.0,
+                        do_sample=True,
+                        top_p=0.95,
+                    )
+                text = self.tokenizer.decode(out[0][input_ids.shape[1]:], skip_special_tokens=True)
+                samples.append(text)
+                print(f"  [Sample {si+1}] {text[:200]}...")
+
+            # Judge the eval samples (fresh call, not from train cache)
+            try:
+                judge_result = call_judge(prompt_text, samples)
+            except Exception as e:
+                judge_result = {"_error": str(e)}
+
+            # Eval-only metrics: deterministic, no extra API calls
+            openers = _unique_5word_openers(samples)
+            vocab = _vocab_diversity(samples)
+            all_unique_openers.append(openers)
+            all_vocab_div.append(vocab)
+
+            gd = judge_result.get("group_diversity") if isinstance(judge_result, dict) else None
+            if isinstance(gd, (int, float)) and gd < 3:
+                print(f"  [LOW DIVERSITY] group_diversity={gd} on prompt {pi+1}")
+
+            all_data[prompt_text] = {
+                "samples": samples,
+                "judge": judge_result,
+                "unique_openers": openers,
+                "vocab_diversity": vocab,
+            }
+
+        # Aggregate eval metrics → wandb
+        if wandb.run is not None and all_unique_openers:
+            wandb.log({
+                "eval/unique_openers": statistics.fmean(all_unique_openers),
+                "eval/vocab_diversity": statistics.fmean(all_vocab_div),
+            })
+
+        save_dir = os.path.join(args.output_dir, "eval_samples")
+        os.makedirs(save_dir, exist_ok=True)
+        save_path = os.path.join(save_dir, f"step_{step}.json")
+        with open(save_path, "w") as f:
+            json.dump(all_data, f, indent=2)
+        print(f"\n  [saved] {save_path}\n{'='*60}\n")
+        model.train()
+
+
+# ── Training ────────────────────────────────────────────────────────
+def main():
+    print(f"Model:       {BASE_MODEL}")
+    print(f"Judge:       {JUDGE_MODEL}")
+    print(f"Group size:  {NUM_GENERATIONS}")
+    print(f"Quality gate: {QUALITY_GATE}")
+    print(f"Group div bonus: {GROUP_DIV_BONUS}")
+
+    dataset, eval_dataset = load_writing_prompts()
+    print(f"Dataset:     {len(dataset)} train, {len(eval_dataset)} eval")
+
+    lora_config = LoraConfig(
+        r=LORA_R,
+        lora_alpha=LORA_ALPHA,
+        target_modules=LORA_TARGET_MODULES,
+        task_type="CAUSAL_LM",
+    )
+
+    training_config = GRPOConfig(
+        output_dir=OUTPUT_DIR,
+        num_train_epochs=NUM_TRAIN_EPOCHS,
+        per_device_train_batch_size=BATCH_SIZE,
+        num_generations=NUM_GENERATIONS,
+        max_completion_length=MAX_COMPLETION_TOKENS,
+        learning_rate=LEARNING_RATE,
+        logging_steps=1,
+        save_steps=100,
+        bf16=True,
+        gradient_accumulation_steps=GRAD_ACCUM,
+        # vLLM
+        use_vllm=True,
+        vllm_mode="colocate",
+        vllm_gpu_memory_utilization=0.3,
+        vllm_max_model_length=4096,
+        # GRPO sampling
+        temperature=1.0,
+        # Dr. GRPO: kill KL penalty and reward std normalization
+        beta=0.0,
+        scale_rewards=False,
+        # GDPO: decoupled per-objective normalization before summing
+        multi_objective_aggregation="normalize_then_sum",
+        report_to="wandb",
+    )
+
+    trainer = GRPOTrainer(
+        model=BASE_MODEL,
+        reward_funcs=[quality_reward, novelty_reward, completeness_reward],
+        args=training_config,
+        train_dataset=dataset,
+        peft_config=lora_config,
+    )
+
+    eval_prompts = [ex["prompt"] for ex in eval_dataset]
+    eval_cb = EvalSampleCallback(eval_prompts, trainer.processing_class)
+    trainer.add_callback(eval_cb)
+
+    print("\n── Starting training ──")
+    trainer.train()
+    trainer.save_model(os.path.join(OUTPUT_DIR, "final"))
+    print("Done.")
+
+
+if __name__ == "__main__":
+    main()