Spaces:
Running on Zero
Running on Zero
| """Future Oracle inference (self-contained port of nla.future_lens.eval's readout path). | |
| A Future Oracle is a LoRA adapter on a Qwen3 base model that reads a single residual-stream | |
| vector (block-`layer` output at token t) and verbalises the tokens the model will produce next. | |
| Recipe (identical to training/eval in EasyNLA `nla/future_lens`): | |
| 1. activation h = hidden_states[layer + 1][t] of the bare model on the token ids (no prompt) | |
| 2. ground truth = the bare model's greedy continuation of ids[:t+1] (the oracle's label convention) | |
| 3. decoder prompt = template.format(layer, k, MARKER) + "\\n" (prompt_format "plain") | |
| 4. the marker token's *input embedding* is replaced by alpha_layer * alpha_mult * h / ||h|| | |
| (alpha_layer = p75 activation norm of that layer, from the collection sidecar) | |
| 5. greedy decode with the adapter on, cut at EOS, truncated to k tokens. | |
| Registry: every `ckpts/<run>/iter_*/adapter_config.json` in the HF dataset repo is an oracle | |
| (latest iter per run; sweep/ablation runs skipped). Alpha per layer + template come from the | |
| matching `data*/train.parquet.nla_meta.yaml` sidecar (same base model, layers covered). | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import re | |
| import threading | |
| from dataclasses import dataclass, field | |
| import torch | |
| import yaml | |
| from huggingface_hub import HfApi, hf_hub_download, snapshot_download | |
| REPO = os.environ.get("FL_HF_REPO", "syvb/rl-future-lens-qwen3-8b") | |
| TOKEN = os.environ.get("HF_TOKEN") | |
| SKIP_RUN = re.compile(r"^(sweep_|ablation_|sft_filt|sft_unf)") | |
| # base checkpoints served (the obsolete chat-model/text-label oracle would add a second 16 GB 8B) | |
| BASES = [x for x in os.environ.get("FO_BASES", "Qwen/Qwen3-0.6B-Base,Qwen/Qwen3-1.7B-Base,Qwen/Qwen3-8B-Base").split(",") if x] | |
| class Sidecar: | |
| data_dir: str | |
| base_model: str | |
| template: str | |
| prompt_format: str | |
| marker: str | |
| marker_id: int | |
| left_id: int | |
| right_id: int | |
| alpha_by_layer: dict[int, float] | |
| class Oracle: | |
| run: str | |
| iter_dir: str # repo path ckpts/<run>/iter_XXXXXXX | |
| step: int | |
| base_model: str | |
| layers: list[int] | |
| alpha_mult: float | |
| injection: str | |
| label: str | |
| sidecar: Sidecar | |
| local_dir: str | None = None | |
| extra: dict = field(default_factory=dict) | |
| def adapter_name(self) -> str: # torch module names can't contain "." | |
| return self.run.replace(".", "_") | |
| def size(self) -> str: | |
| m = re.search(r"Qwen3-([\d.]+B)", self.base_model) | |
| return m.group(1) if m else self.base_model | |
| def name(self) -> str: | |
| kind = "RL" if self.run.startswith("rl_") else "SFT" | |
| lay = f"layer {self.layers[0]}" if len(self.layers) == 1 else f"layers {','.join(map(str, self.layers))}" | |
| return f"Future Oracle 路 Qwen3-{self.size} 路 {lay} 路 {kind} {self.run} @ {self.step}" | |
| def _dl(path: str, local_dir: str) -> str: | |
| return hf_hub_download(REPO, path, repo_type="dataset", token=TOKEN, local_dir=local_dir) | |
| def load_sidecar(data_dir: str, cache: str) -> Sidecar | None: | |
| for split in ("train", "eval"): | |
| try: | |
| p = _dl(f"{data_dir}/{split}.parquet.nla_meta.yaml", cache) | |
| except Exception: | |
| continue | |
| m = yaml.safe_load(open(p)) | |
| fl, tk = m.get("future_lens", {}), m.get("tokens", {}) | |
| if not fl or not tk: | |
| continue | |
| pf = (fl.get("extra") or {}).get("prompt_format") or ("plain" if "Base" in m["extraction"]["base_model"] else "chat") | |
| return Sidecar( | |
| data_dir=data_dir, base_model=m["extraction"]["base_model"], template=fl["template"], prompt_format=pf, | |
| marker=tk["injection_char"], marker_id=int(tk["injection_token_id"]), | |
| left_id=int(tk["injection_left_neighbor_id"]), right_id=int(tk["injection_right_neighbor_id"]), | |
| alpha_by_layer={int(k): float(v) for k, v in fl["injection_scale_by_layer"].items()}, | |
| ) | |
| return None | |
| def scan_registry(cache: str = "/tmp/fo_meta") -> list[Oracle]: | |
| """Every finished oracle adapter in the repo, newest run first.""" | |
| api = HfApi(token=TOKEN) | |
| files = [s.rfilename for s in api.dataset_info(REPO).siblings] | |
| sidecars: list[Sidecar] = [] | |
| for d in sorted({f.split("/")[0] for f in files if f.startswith("data") and f.endswith(".nla_meta.yaml")}): | |
| if d == "data_unf": # filter-ablation split: not a training split of any kept oracle | |
| continue | |
| sc = load_sidecar(d, cache) | |
| if sc: | |
| sidecars.append(sc) | |
| iters: dict[str, list[tuple[int, str]]] = {} | |
| for f in files: | |
| m = re.match(r"ckpts/([^/]+)/iter_(\d+)/adapter_config\.json$", f) | |
| if m and not SKIP_RUN.match(m.group(1)) and f"ckpts/{m.group(1)}/iter_{m.group(2)}/future_lens.json" in files: | |
| iters.setdefault(m.group(1), []).append((int(m.group(2)), f"ckpts/{m.group(1)}/iter_{m.group(2)}")) | |
| oracles: list[Oracle] = [] | |
| for run, its in iters.items(): | |
| step, idir = max(its) | |
| cfg = json.load(open(_dl(f"{idir}/adapter_config.json", cache))) | |
| fl = json.load(open(_dl(f"{idir}/future_lens.json", cache))) | |
| base = cfg["base_model_name_or_path"] | |
| if base not in BASES: | |
| continue | |
| layers = [int(x) for x in str(fl.get("layers", "")).split(",") if x] | |
| cands = [s for s in sidecars if s.base_model == base and set(layers) <= set(s.alpha_by_layer)] | |
| if not cands: | |
| continue | |
| sfx = re.search(r"_(\d+(?:\.\d+)?b)_L(\d+)$", run) | |
| pref = f"data_{sfx.group(1)}_L{sfx.group(2)}" if sfx else None | |
| sc = next((s for s in cands if s.data_dir == pref), None) or sorted(cands, key=lambda s: s.data_dir)[-1] | |
| oracles.append(Oracle(run=run, iter_dir=idir, step=step, base_model=base, layers=layers, | |
| alpha_mult=float(fl.get("alpha_mult", 1.0)), injection=fl.get("injection", "replace_embed"), | |
| label=fl.get("label", "greedy"), sidecar=sc, extra=fl)) | |
| # small models first, then by run name | |
| def key(o): | |
| m = re.search(r"([\d.]+)B", o.size) | |
| return (float(m.group(1)) if m else 99, o.run) | |
| return sorted(oracles, key=key) | |
| # ---------------------------------------------------------------------------- | |
| # models | |
| # ---------------------------------------------------------------------------- | |
| class Bank: | |
| """Base models (bf16) with all oracle adapters attached as named PEFT adapters. | |
| ZeroGPU: everything is assembled on CPU at import time and moved to CUDA once per base model | |
| (`to_device`), which ZeroGPU packs for its workers. Attaching an adapter later (Refresh) moves | |
| that base back to CPU, attaches, and moves it to CUDA again.""" | |
| def __init__(self, device: str = "cuda", dtype=torch.bfloat16): | |
| self.device, self.dtype = device, dtype | |
| self.tok = None | |
| self.models: dict[str, torch.nn.Module] = {} # base id -> PeftModel (or bare base) | |
| self.on_device: set[str] = set() | |
| self.loaded: set[str] = set() # oracle runs attached | |
| self.lock = threading.Lock() | |
| def to_device(self): | |
| for b, m in self.models.items(): | |
| if b not in self.on_device: | |
| m.to(self.device) | |
| self.on_device.add(b) | |
| def tokenizer(self, base: str): | |
| if self.tok is None: | |
| from transformers import AutoTokenizer | |
| self.tok = AutoTokenizer.from_pretrained(base) # Qwen3 sizes share one tokenizer | |
| return self.tok | |
| def base(self, base: str): | |
| if base not in self.models: | |
| self.tokenizer(base) | |
| from transformers import AutoModelForCausalLM | |
| m = AutoModelForCausalLM.from_pretrained(base, dtype=self.dtype, attn_implementation="sdpa") | |
| self.models[base] = m.eval() # CPU until to_device() | |
| return self.models[base] | |
| def attach(self, o: Oracle): | |
| if o.run in self.loaded: | |
| return | |
| with self.lock: | |
| if o.run in self.loaded: | |
| return | |
| if o.local_dir is None: | |
| o.local_dir = os.path.join(snapshot_download(REPO, repo_type="dataset", token=TOKEN, | |
| allow_patterns=[f"{o.iter_dir}/*"], | |
| local_dir="/tmp/fo_ckpts"), o.iter_dir) | |
| from peft import PeftModel | |
| m = self.base(o.base_model) | |
| back = o.base_model in self.on_device | |
| if back: # late attach: assemble on CPU, then re-pack | |
| m.to("cpu"); self.on_device.discard(o.base_model) | |
| # torch_device="cpu": PEFT otherwise infers "cuda" from ZeroGPU's patched availability flag | |
| # and maps the safetensors there, which is forbidden outside a GPU call | |
| if isinstance(m, PeftModel): | |
| m.load_adapter(o.local_dir, adapter_name=o.adapter_name, torch_device="cpu") | |
| else: | |
| self.models[o.base_model] = PeftModel.from_pretrained(m, o.local_dir, adapter_name=o.adapter_name, | |
| torch_device="cpu").eval() | |
| self.loaded.add(o.run) | |
| if back: | |
| self.to_device() | |
| def model(self, o: Oracle): | |
| self.attach(o) | |
| m = self.models[o.base_model] | |
| m.set_adapter(o.adapter_name) | |
| return m | |
| def stop_ids(tok, model) -> set[int]: | |
| ids = {tok.eos_token_id} | |
| gc = getattr(getattr(model, "generation_config", None), "eos_token_id", None) | |
| if gc is not None: | |
| ids.update(gc if isinstance(gc, (list, tuple)) else [gc]) | |
| ids.discard(None) | |
| return ids | |
| def read_future(bank: Bank, o: Oracle, layer: int, ids: list[int], t: int, k: int, | |
| control: str = "real", alpha_mult: float | None = None, shuffle_t: int | None = None) -> dict: | |
| """One oracle readout at token t of `ids`. control: real | none | shuffled (vector from | |
| position shuffle_t of the same text) | wrong_layer (layer given, adapter untrained on it).""" | |
| assert 0 <= t < len(ids) | |
| dev = bank.device | |
| model = bank.model(o) | |
| tok = bank.tokenizer(o.base_model) | |
| sc = o.sidecar | |
| x = torch.tensor([ids], dtype=torch.long, device=dev) | |
| # ---- target model: activation + its own greedy continuation (adapter OFF) ---- | |
| with model.disable_adapter(): | |
| out = model(input_ids=x, output_hidden_states=True, use_cache=False) | |
| src_t = shuffle_t if control == "shuffled" else t | |
| h = out.hidden_states[layer + 1][0, src_t].float() | |
| top1 = int(out.logits[0, t].argmax()) | |
| gen = model.generate(input_ids=x[:, : t + 1], attention_mask=torch.ones_like(x[:, : t + 1]), | |
| max_new_tokens=k, do_sample=False, pad_token_id=tok.pad_token_id or tok.eos_token_id) | |
| greedy = gen[0, t + 1:].tolist() | |
| p_next = torch.log_softmax(out.logits[0, t].float(), -1) | |
| # ---- decoder prompt with the marker ---- | |
| prompt = sc.template.format(layer=int(layer), k=int(k), injection_char=sc.marker) | |
| if sc.prompt_format == "plain": | |
| text = prompt + "\n" | |
| else: | |
| text = tok.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False, | |
| add_generation_prompt=True, enable_thinking=False) | |
| p_ids = tok.encode(text, add_special_tokens=False) | |
| pos = [i for i in range(1, len(p_ids) - 1) | |
| if p_ids[i] == sc.marker_id and p_ids[i - 1] == sc.left_id and p_ids[i + 1] == sc.right_id] | |
| assert len(pos) == 1, f"marker not found exactly once in the decoder prompt: {pos}" | |
| alpha = sc.alpha_by_layer[int(layer)] * (o.alpha_mult if alpha_mult is None else alpha_mult) | |
| vec = None if control == "none" else (h / h.norm().clamp_min(1e-12)) * alpha | |
| embed = model.get_input_embeddings() | |
| def hook(module, args, kwargs, output): | |
| if vec is None: | |
| return output | |
| inp = kwargs.get("input") if kwargs else None | |
| if inp is None and args: | |
| inp = args[0] | |
| if inp is None or inp.dim() != 2 or inp.shape[1] < 2: | |
| return output # KV-cached decode steps | |
| out2 = output.clone() | |
| out2[0, pos[0]] = vec.to(out2.dtype) | |
| return out2 | |
| hd = embed.register_forward_hook(hook, with_kwargs=True) | |
| try: | |
| px = torch.tensor([p_ids], dtype=torch.long, device=dev) | |
| g = model.generate(input_ids=px, attention_mask=torch.ones_like(px), max_new_tokens=k + 3, do_sample=False, | |
| pad_token_id=tok.pad_token_id or tok.eos_token_id) | |
| finally: | |
| hd.remove() | |
| resp = g[0, len(p_ids):].tolist() | |
| eos = stop_ids(tok, model) | |
| cut = next((i for i, x_ in enumerate(resp) if x_ in eos), len(resp)) | |
| readout = resp[:cut][:k] | |
| actual = ids[t + 1: t + 1 + k] | |
| return { | |
| "oracle": o.name, "run": o.run, "size": o.size, "layer": int(layer), "t": t, "k": k, "control": control, | |
| "alpha": alpha, "act_norm": float(h.norm()), | |
| "readout": readout, "readout_text": tok.decode(readout), | |
| "greedy": greedy, "greedy_text": tok.decode(greedy), | |
| "actual": actual, "actual_text": tok.decode(actual), | |
| "hits": [i < len(readout) and i < len(greedy) and readout[i] == greedy[i] for i in range(k)], | |
| "top1_correct": (top1 == ids[t + 1]) if t + 1 < len(ids) else None, # None: no next token to compare | |
| "p_actual_next": float(p_next[ids[t + 1]].exp()) if t + 1 < len(ids) else None, | |
| } | |