esh-pfi-infrastructure/scripts/r49-corpus/gen_beats.py

"""Skaldsong's actual shape: one beat sentence in, one paragraph out, stitchable.

That is a narrower job than anything tested so far, and it fails in ways free-form
continuation does not:

  * DRIFT off the beat breaks the stitch -- the next paragraph no longer follows.
  * RUN-ON breaks it too. The deliverable is a paragraph, not 400 tokens that wander
    into the following scene, because the next beat owns that scene.
  * FRAMING ("I told it briefly") renders nothing at all -- measured across 6 seeds
    on the handoff prompt.
  * RENAMING is now a product blocker rather than a curiosity: the D2/D3 rename pool
    taught the adapter that character names come from it, so a caller's own name can
    be rewritten mid-passage and the stitched story loses its protagonist.

So each format is scored on all four, not eyeballed. Run-on is measured by whether a
paragraph break arrived before the token budget ran out -- the text is truncated at
the first blank line for display, and whether truncation was NEEDED is the signal.

FORMATS, in ascending order of how much structure they impose. The few-shot one is
the interesting entry: a completion model's native instruction channel is a worked
example, and none of the earlier prompts gave it one.
"""
from __future__ import annotations
import argparse, json, re, time
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# A worked example for the few-shot formats. Written by hand in the target register,
# deliberately on a subject unrelated to dogs and kittens so it cannot leak content
# into the answer -- only shape.
EX_BEAT = "The carrier's cart broke its axle at the ford."
EX_PARA = ("The cart came to a standstill in the middle of the water, canted over like a "
           "ship gone aground, and the carrier stood in the shallows with his hand on the "
           "shaft, saying nothing at all. I watched from the bank. The river ran brown and "
           "quick about his boots; a hamper had gone over and was turning slowly downstream, "
           "and he let it go. It was not the loss that held him, I think, but the hour: he "
           "had been due at the mill before noon, and it was past two.")

FORMATS = {
    "bare": lambda b: b + " ",
    "para-break": lambda b: b + "\n\n",
    "labelled": lambda b: f"The passage I wrote from this beat:\n\nBEAT: {b}\n\nPASSAGE: ",
    "epigraph": lambda b: f"_{b}_\n\n",
    "fewshot": lambda b: (f"BEAT: {EX_BEAT}\nPASSAGE: {EX_PARA}\n\nBEAT: {b}\nPASSAGE: "),
    "fewshot-bare": lambda b: (f"{EX_BEAT}\n\n{EX_PARA}\n\n{b}\n\n"),
    # ---- round two. Round one failed everywhere, so before calling that a property
    # of the adapter these four give the strongest untested patterns a fair run.
    # THREE examples, not one: one-shot is thin, and a format dismissed on one
    # example has not been tested, it has been under-fed.
    "fewshot3": lambda b: ("".join(f"BEAT: {eb}\nPASSAGE: {ep}\n\n"
                                   for eb, ep in EXTRA_EXAMPLES)
                           + f"BEAT: {b}\nPASSAGE: "),
    # The letter prompt's winning move was "label the artifact AND begin it". These
    # apply it to a beat: state the beat, then open the paragraph with a phrase that
    # COMMITS to elaborating what was just said, so moving on is off the table.
    "elaborate": lambda b: f"{b} It happened in this way. ",
    "recount": lambda b: f"{b} I remember the whole of it, and will set it down. ",
    # Label + begin, with the paragraph seeded by the beat's own opening words so the
    # first thing it writes is already inside the beat rather than after it.
    "label-begin": lambda b: (f"BEAT: {b}\nPASSAGE: " + " ".join(b.split()[:3]) + " "),
}
EXTRA_EXAMPLES = [
    (EX_BEAT, EX_PARA),
    ("The housekeeper refused to give up the key.",
     "She stood with her hand closed over it and her chin down, and said that the room "
     "had been shut since March and would stay shut. I asked her whose order it was. She "
     "said it was nobody's order, it was sense; and then, seeing I meant to press her, she "
     "put the key into her apron pocket and held the pocket. There was no arguing with the "
     "gesture. I went back along the passage and heard her breathing behind me the whole way."),
    ("A letter came for the master and was burned unopened.",
     "It lay on the salver a quarter of an hour, and I saw the hand on it -- a small, "
     "sloped, foreign hand -- before he came in. He turned it over once, read the "
     "postmark, and put it on the fire without breaking the seal. The wax ran first and "
     "then the paper caught. He watched it to the end, which is what I remember: not the "
     "burning, but that he stayed to see it finished."),
]

ap = argparse.ArgumentParser()
ap.add_argument("--base", required=True)
ap.add_argument("--adapter", default=None)
ap.add_argument("--beats", required=True, help="json list of {id, beat}")
ap.add_argument("--formats", nargs="+", default=list(FORMATS))
ap.add_argument("--out", required=True)
ap.add_argument("--seeds", type=int, nargs="+", default=[1234, 5678])
ap.add_argument("--max-new-tokens", type=int, default=300)
ap.add_argument("--temperature", type=float, default=0.9)
ap.add_argument("--top-p", type=float, default=0.95)
a = ap.parse_args()

beats = json.loads(Path(a.beats).read_text())
tok = AutoTokenizer.from_pretrained(a.base)
model = AutoModelForCausalLM.from_pretrained(a.base, dtype=torch.bfloat16,
                                             attn_implementation="sdpa").to("cuda")
if a.adapter:
    from peft import PeftModel
    model = PeftModel.from_pretrained(model, a.adapter)
    deltas = [float(m.lora_B["default"].weight.abs().sum())
              for m in model.modules() if hasattr(m, "lora_B")]
    nz = sum(1 for d in deltas if d > 0)
    print(f"[gen] adapter bound: {nz}/{len(deltas)} lora_B tensors non-zero", flush=True)
    if nz == 0:
        raise SystemExit("REFUSING: adapter applied but every lora_B is zero")
model.eval()

STOP = re.compile(r"\n\s*\n")


def keywords(beat):
    """Content words worth checking for, to score staying ON the beat."""
    drop = {"the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "of", "his",
            "her", "he", "she", "it", "was", "were", "had", "would", "not", "be",
            "by", "as", "with", "for", "from", "that", "this", "up", "down", "she"}
    return [w for w in re.findall(r"[a-z']+", beat.lower()) if w not in drop and len(w) > 3]


out = Path(a.out); out.parent.mkdir(parents=True, exist_ok=True)
t0 = time.time()
with out.open("w", encoding="utf-8") as fh:
    for fmt in a.formats:
        build = FORMATS[fmt]
        for b in beats:
            for seed in a.seeds:
                torch.manual_seed(seed)
                prompt = build(b["beat"])
                ids = tok(prompt, return_tensors="pt").to("cuda")
                with torch.no_grad():
                    g = model.generate(**ids, do_sample=True, temperature=a.temperature,
                                       top_p=a.top_p, max_new_tokens=a.max_new_tokens,
                                       pad_token_id=tok.eos_token_id)
                raw = tok.decode(g[0][ids["input_ids"].shape[1]:], skip_special_tokens=True)
                m = STOP.search(raw.strip())
                para = (raw.strip()[:m.start()] if m else raw.strip()).strip()
                kws = keywords(b["beat"])
                hit = sum(1 for k in kws if k[:5] in para.lower())
                fh.write(json.dumps({
                    "format": fmt, "id": b["id"], "beat": b["beat"], "seed": seed,
                    "prompt": prompt, "paragraph": para, "raw_tail": raw.strip()[m.end():][:200] if m else "",
                    # ran_on: the model never closed a paragraph inside the budget, so
                    # a stitcher would have to cut it mid-thought.
                    "ran_on": m is None,
                    "words": len(para.split()),
                    "beat_keywords": kws, "keyword_hits": hit,
                }) + "\n")
                print(f"  {fmt:14} {b['id']:>8} seed={seed} {len(para.split()):>4}w "
                      f"kw {hit}/{len(kws)} {'RAN-ON' if m is None else ''}", flush=True)
print(f"[gen] -> {out} in {time.time()-t0:.0f}s", flush=True)