| """Shared utilities: device, seed, prompt format, reward, checkpoint load/save.""" |
| from __future__ import annotations |
|
|
| import os |
| import random |
| import re |
| from contextlib import nullcontext |
|
|
| import numpy as np |
| import torch |
|
|
| from tokenizer import ASSISTANT, EOS, USER |
|
|
| |
|
|
| def get_device() -> str: |
| if torch.cuda.is_available(): |
| return "cuda" |
| if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available(): |
| return "mps" |
| return "cpu" |
|
|
|
|
| def autocast_ctx(device: str): |
| """Mixed precision: bf16 where the GPU supports it (A100/L4), fp16 on T4, fp32 on CPU.""" |
| if device != "cuda": |
| return nullcontext(), None |
| dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 |
| return torch.autocast("cuda", dtype=dtype), dtype |
|
|
|
|
| def set_seed(seed: int) -> None: |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
|
|
|
|
| def ensure_dir(path: str) -> None: |
| d = os.path.dirname(path) |
| if d: |
| os.makedirs(d, exist_ok=True) |
|
|
|
|
| |
| PROMPT_TEMPLATE = "Write a short story that uses the words: {words}." |
|
|
|
|
| def build_prompt(words: list[str]) -> str: |
| return f"{USER}{PROMPT_TEMPLATE.format(words=', '.join(words))}{ASSISTANT}" |
|
|
|
|
| def build_example(words: list[str], story: str) -> tuple[str, str]: |
| """Return (prompt, response). The SFT loss is computed on the response only.""" |
| return build_prompt(words), story.strip() + EOS |
|
|
|
|
| _STOP = set(""" |
| a an the and or but so if then than that this these those there their they them he she his her him it its |
| was were is are be been being am do does did done have has had having to of in on at for with from by as |
| into onto over under up down out off very too also just not no yes all any some one two three every each |
| said says say went go goes going came come comes get got make made can could would should will shall may |
| might must when where what which who whom why how because while after before again once day time little |
| big lot lots very much many more most other another then now away back here there with |
| upon named lived felt wanted liked loved looked found know knew think thought told asked started saw |
| """.split()) |
|
|
|
|
| def pick_keywords(story: str, k: int, rng: random.Random) -> list[str] | None: |
| """Pick k content words from a story — used as the SFT/DPO instruction.""" |
| cands = sorted({w for w in re.findall(r"[a-z]+", story.lower()) if len(w) >= 4 and w not in _STOP}) |
| if len(cands) < k: |
| return None |
| return rng.sample(cands, k) |
|
|
|
|
| |
| def keyword_reward(text: str, words: list[str], finished: bool) -> float: |
| """ |
| Rule-based, verifiable reward (in the spirit of 'verifiable rewards'): |
| - fraction of required words that appear (0..1) |
| - +0.2 if the model closes the story with EOS (not truncated) |
| """ |
| found = set(re.findall(r"[a-z]+", text.lower())) |
| hit = sum(1 for w in words if w in found) / max(len(words), 1) |
| return hit + (0.2 if finished else 0.0) |
|
|
|
|
| def keyword_hit_rate(text: str, words: list[str]) -> float: |
| found = set(re.findall(r"[a-z]+", text.lower())) |
| return sum(1 for w in words if w in found) / max(len(words), 1) |
|
|
|
|
| def keyword_mentions(text: str, words: list[str]) -> float: |
| """Average number of times each required word appears. ~1-2 is natural; much higher = stuffing.""" |
| toks = re.findall(r"[a-z]+", text.lower()) |
| return sum(toks.count(w) for w in words) / max(len(words), 1) |
|
|
|
|
| def repeated_trigram_rate(text: str) -> float: |
| """Fraction of word trigrams that already appeared earlier in the text (degenerate repetition).""" |
| toks = re.findall(r"[a-z']+", text.lower()) |
| grams = list(zip(toks, toks[1:], toks[2:])) |
| if not grams: |
| return 0.0 |
| return 1.0 - len(set(grams)) / len(grams) |
|
|
|
|
| def keyword_reward_v2(text: str, words: list[str], finished: bool) -> float: |
| """ |
| Anti-stuffing reward: v1 plus penalties for the failure modes seen in the first run. |
| - penalty 0.1 per mention beyond 2 per required word (keyword stuffing) |
| - penalty 1.0 x repeated-trigram rate (looping text) |
| """ |
| toks = re.findall(r"[a-z]+", text.lower()) |
| excess = sum(max(0, toks.count(w) - 2) for w in words) |
| return keyword_reward(text, words, finished) - 0.1 * excess - repeated_trigram_rate(text) |
|
|
|
|
| REWARDS = {"v1": keyword_reward, "v2": keyword_reward_v2} |
|
|
|
|
| |
| def save_checkpoint(path: str, model, model_cfg, extra: dict | None = None) -> None: |
| ensure_dir(path) |
| raw = getattr(model, "_orig_mod", model) |
| torch.save({"model": raw.state_dict(), "model_cfg": model_cfg.__dict__, **(extra or {})}, path) |
|
|
|
|
| def load_model(path: str, device: str): |
| from config import ModelConfig |
| from model import GPT |
|
|
| ckpt = torch.load(path, map_location=device, weights_only=False) |
| cfg = ModelConfig(**ckpt["model_cfg"]) |
| model = GPT(cfg).to(device) |
| model.load_state_dict(ckpt["model"]) |
| return model, cfg |
|
|
|
|
| def load_model_dir(path: str, device: str): |
| """Load an exported model folder (config.json + model.safetensors), e.g. from the Hugging Face Hub.""" |
| import json |
|
|
| from safetensors.torch import load_model as st_load |
|
|
| from config import ModelConfig |
| from model import GPT |
|
|
| with open(os.path.join(path, "config.json")) as f: |
| cfg = ModelConfig(**json.load(f)) |
| model = GPT(cfg) |
| st_load(model, os.path.join(path, "model.safetensors")) |
| return model.to(device).eval(), cfg |
|
|
