| |
| """CPU inference for Octen-WebGR. Same protocol as infer_gpu.py (query normalisation, prompt, constrained beam search, |
| scoring); it only removes two algorithmic redundancies of the reference implementation: |
| |
| 1. Level 0 runs the prompt once and keeps its KV cache; the K beams of level 1 reuse it and each feed one token |
| (the reference recomputes "prompt + SID token" K times, roughly 7-8x the compute). |
| 2. Level-1 log_softmax is taken over the valid children only: sid_table/ is sorted by key = sid0 * 65536 + sid1, |
| so the valid children of prefix c are the contiguous rows [start[c], start[c+1]). |
| |
| Results match infer_gpu.py except for occasional swaps between candidates whose scores are nearly tied. |
| |
| python infer_cpu.py --model Octen/Octen-WebGR-0.6B --query "youtube" "boeing internship" --threads 16 |
| python infer_cpu.py --model /path/to/local/copy --check --threads 16 |
| """ |
| import argparse, glob, json, os, time |
| import numpy as np |
| import pyarrow as pa |
| import pyarrow.parquet as pq |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| BASE_V, V, MAX_Q = 151936, 65536, 32 |
| PAD = 151643 |
|
|
|
|
| def resolve(model): |
| """Local directory, or a Hugging Face repo id that is downloaded once into the HF cache.""" |
| if os.path.isdir(model): |
| return model |
| from huggingface_hub import snapshot_download |
| return snapshot_download(model) |
|
|
|
|
| def norm_query(q): |
| return " ".join(q.split()).lower() |
|
|
|
|
| class URLs: |
| """URL column kept as one array per table file (no combined copy, so peak memory is one table); rows are global.""" |
| def __init__(self, files): |
| self.parts = [pq.read_table(f, columns=["url"]).column("url").combine_chunks() for f in files] |
| self.start = np.cumsum([0] + [len(p) for p in self.parts]) |
|
|
| def take(self, rows): |
| f = np.searchsorted(self.start, rows, side="right") - 1 |
| return [self.parts[i][int(r - self.start[i])].as_py() for i, r in zip(f, rows)] |
|
|
|
|
| class GRCPU: |
| def __init__(self, model, dtype=torch.bfloat16, logits_fp32=False): |
| """dtype: precision of the forward pass (bf16 matches the GPU evaluation; fp32 is slower). |
| logits_fp32: compute the lm_head step with fp32 weights (bf16 logits are spaced 0.125 apart at |logit| >= 16, |
| so nearly tied candidates can swap depending on the kernel).""" |
| self.root = resolve(model) |
| self.tok = AutoTokenizer.from_pretrained(self.root) |
| self.lm = AutoModelForCausalLM.from_pretrained(self.root, dtype=dtype, attn_implementation="sdpa").eval() |
| assert self.lm.config.vocab_size == BASE_V + 2 * V, "unexpected vocabulary size: should be 151936 + 2 x 65536" |
| W = self.lm.lm_head.weight.detach() |
| ldt = torch.float32 if logits_fp32 else dtype |
| self.W0 = W[BASE_V:BASE_V + V].to(ldt).contiguous() |
| self.W1 = W[BASE_V + V:BASE_V + 2 * V].to(ldt).contiguous() |
| files = sorted(glob.glob(f"{self.root}/sid_table/part-*.parquet")) |
| t = pa.concat_tables([pq.read_table(f, columns=["sid0", "sid1"]) for f in files]) |
| s0 = t.column("sid0").to_numpy().astype(np.int64) |
| s1 = t.column("sid1").to_numpy().astype(np.int64) |
| assert (np.diff(s0 * V + s1) > 0).all(), "sid_table must be sorted by key in strictly ascending order (one URL per SID)" |
| self.start = torch.from_numpy(np.searchsorted(s0, np.arange(V + 1))) |
| self.ok0 = self.start[1:] > self.start[:-1] |
| self.sid1 = torch.from_numpy(s1) |
| self.url = URLs(files) |
| self.n = len(s0) |
|
|
| def prompts(self, qs): |
| qs = [norm_query(q) for q in qs] |
| qi = self.tok(qs, add_special_tokens=False)["input_ids"] |
| qs = [self.tok.decode(x[:MAX_Q]) if len(x) > MAX_Q else q for q, x in zip(qs, qi)] |
| return self.tok([f"Query: {q}\nURL:" for q in qs], add_special_tokens=False)["input_ids"] |
|
|
| @torch.inference_mode() |
| def search_rows(self, qs, K=20): |
| """-> (scores, rows) per query, sorted by score; rows index sid_table/.""" |
| P = self.prompts(qs) |
| B, T = len(P), max(len(p) for p in P) |
| ids = torch.full((B, T), PAD, dtype=torch.long) |
| am = torch.zeros((B, T), dtype=torch.long) |
| for i, p in enumerate(P): |
| ids[i, T - len(p):] = torch.tensor(p) |
| am[i, T - len(p):] = 1 |
| out = self.lm.model(input_ids=ids, attention_mask=am, position_ids=(am.cumsum(1) - 1).clamp(min=0), use_cache=True) |
| lg0 = (out.last_hidden_state[:, -1].to(self.W0.dtype) @ self.W0.t()).float().masked_fill(~self.ok0, float("-inf")) |
| k0 = min(K, int(self.ok0.sum())) |
| s0, c0 = lg0.log_softmax(-1).topk(k0, dim=1) |
| |
| cache = out.past_key_values |
| cache.batch_repeat_interleave(k0) |
| am1 = torch.cat([am.repeat_interleave(k0, 0), torch.ones((B * k0, 1), dtype=torch.long)], 1) |
| pos1 = am.sum(1, keepdim=True).repeat_interleave(k0, 0) |
| h1 = self.lm.model(input_ids=(BASE_V + c0).reshape(-1, 1), attention_mask=am1, position_ids=pos1, |
| past_key_values=cache, use_cache=True).last_hidden_state[:, -1] |
| lg1 = h1.to(self.W1.dtype) @ self.W1.t() |
| cc = c0.reshape(-1) |
| lo = self.start[cc] |
| n = self.start[cc + 1] - lo |
| r = torch.repeat_interleave(torch.arange(B * k0), n) |
| j = torch.arange(int(n.sum())) - torch.repeat_interleave(n.cumsum(0) - n, n) + lo[r] |
| x = lg1[r, self.sid1[j]].float() |
| mx = torch.full((B * k0,), float("-inf")).scatter_reduce(0, r, x, "amax") |
| lse = mx + torch.zeros(B * k0).index_add_(0, r, (x - mx[r]).exp()).log() |
| sc = s0.reshape(-1)[r] + x - lse[r] |
| off = [0] + n.view(B, k0).sum(1).cumsum(0).tolist() |
| res = [] |
| for b in range(B): |
| s, i = sc[off[b]:off[b + 1]].topk(min(K, off[b + 1] - off[b])) |
| res.append((s.numpy(), j[off[b] + i].numpy())) |
| return res |
|
|
| def search(self, qs, K=20): |
| """-> one list per query of (url, log_prob, table_row), sorted by score (same format as infer_gpu.GR.search).""" |
| return [list(zip(self.url.take(rr), s.tolist(), rr.tolist())) for s, rr in self.search_rows(qs, K)] |
|
|
|
|
| def check(gr, bs=32): |
| """Compare against smoke_test.json (hand-picked queries; expected results from the released weights and table).""" |
| ref = json.load(open(f"{gr.root}/smoke_test.json")) |
| assert gr.n == ref["table_rows"], f"sid_table has {gr.n:,} rows, expected {ref['table_rows']:,}" |
| qs, want = [x["query"] for x in ref["cases"]], [x["top20"] for x in ref["cases"]] |
| ov, top1, n = [], 0, len(qs) |
| for i in range(0, n, bs): |
| for res, w in zip(gr.search(qs[i:i + bs], 20), want[i:i + bs]): |
| got = [u for u, *_ in res] |
| ov.append(len(set(got) & set(w)) / max(len(w), 1)) |
| top1 += bool(got and w and got[0] == w[0]) |
| mo, t1 = float(np.mean(ov)), top1 / n |
| print(f"{n} queries: mean top-20 overlap {mo:.4f}, top-1 agreement {t1:.4f}, identical top-20 {np.mean(np.array(ov) == 1):.3f}") |
| ok = mo >= 0.9 and t1 >= 0.9 |
| print("PASS" if ok else "FAIL (small differences from kernels are expected, but overlap should be >= 0.9)") |
| return ok |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", default=os.path.dirname(os.path.abspath(__file__)), help="local directory or HF repo id") |
| ap.add_argument("--query", nargs="*", default=[]) |
| ap.add_argument("--k", type=int, default=20) |
| ap.add_argument("--threads", type=int, default=0, help="0 = torch default") |
| ap.add_argument("--fp32", action="store_true", help="fp32 forward pass (default bf16, as in the GPU evaluation)") |
| ap.add_argument("--check", action="store_true") |
| a = ap.parse_args() |
| if a.threads: |
| torch.set_num_threads(a.threads) |
| t0 = time.time() |
| gr = GRCPU(a.model, dtype=torch.float32 if a.fp32 else torch.bfloat16) |
| print(f"Loaded {gr.n:,} candidate URLs, {torch.get_num_threads()} threads, {time.time() - t0:.0f}s", flush=True) |
| for q in a.query: |
| t = time.perf_counter() |
| res = gr.search([q], a.k)[0] |
| print(f"\n{q} ({(time.perf_counter() - t) * 1000:.0f} ms)") |
| for i, (u, lp, _) in enumerate(res, 1): |
| print(f" {i:2d}. {np.exp(lp):.3f} {u}") |
| if a.check and not check(gr): |
| raise SystemExit(1) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
