| |
| """Reference GPU inference for Octen-WebGR: query -> top-K URLs by constrained beam search over two semantic-ID tokens. |
| |
| Inference protocol (identical to the evaluation in the technical report): |
| 1. Normalise the query: collapse Unicode whitespace to one space, strip, lowercase. |
| 2. Keep at most the first 32 tokens of the query (decoded back to a string before building the prompt). |
| 3. Prompt = "Query: {q}\\nURL:", tokenised without special tokens. |
| 4. Level t (t = 0, 1) of the SID is vocabulary token 151936 + t * 65536 + code; only those 65,536 rows of |
| lm_head are scored at each level. |
| 5. Constrained beam search: only prefixes present in sid_table/ are allowed, and log_softmax is taken over the |
| valid children only. K beams are kept at each level (default 20). |
| 6. Score = sum of the two log-probabilities (exp(score) = the model's probability of the URL among valid SIDs). |
| (sid0, sid1) is looked up in sid_table/ to get the URL. |
| |
| This implementation recomputes the prompt for every beam at level 1 and pads each batch to its longest prompt; |
| it favours clarity over speed. infer_cpu.py reuses the prompt's KV cache instead. |
| |
| python infer_gpu.py --model Octen/Octen-WebGR-4B --query "youtube" "boeing internship" |
| python infer_gpu.py --model /path/to/local/copy --check # compare against smoke_test.json |
| """ |
| 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, L, MAX_Q = 151936, 65536, 2, 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 GR: |
| def __init__(self, model, device="cuda"): |
| self.root = resolve(model) |
| self.dev = device |
| self.dev_type = torch.device(device).type |
| self.tok = AutoTokenizer.from_pretrained(self.root) |
| attn = "sdpa" |
| if self.dev_type == "cuda": |
| try: |
| import flash_attn |
| attn = "flash_attention_2" |
| except ImportError: |
| pass |
| self.lm = AutoModelForCausalLM.from_pretrained(self.root, dtype=torch.bfloat16, attn_implementation=attn).to(device).eval() |
| assert self.lm.config.vocab_size == BASE_V + L * V, "unexpected vocabulary size: should be 151936 + 2 x 65536" |
| self.W = self.lm.lm_head.weight |
| 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]) |
| key = t.column("sid0").to_numpy().astype(np.int64) * V + t.column("sid1").to_numpy().astype(np.int64) |
| assert (np.diff(key) > 0).all(), "sid_table must be sorted by key in strictly ascending order (one URL per SID)" |
| self.keys = torch.from_numpy(key).to(device) |
| ok0 = np.zeros(V, bool) |
| ok0[np.unique(key // V)] = True |
| self.ok0 = torch.from_numpy(ok0).to(device) |
| self.url = URLs(files) |
| self.n = len(key) |
|
|
| 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"] |
|
|
| def last_hidden(self, seqs): |
| P = max(len(s) for s in seqs) |
| ids = torch.full((len(seqs), P), PAD, dtype=torch.long) |
| am = torch.zeros((len(seqs), P), dtype=torch.long) |
| for i, s in enumerate(seqs): |
| ids[i, P - len(s):] = torch.tensor(s) |
| am[i, P - len(s):] = 1 |
| ids, am = ids.to(self.dev), am.to(self.dev) |
| pos = (am.cumsum(1) - 1).clamp(min=0) |
| return self.lm.model(input_ids=ids, attention_mask=am, position_ids=pos).last_hidden_state[:, -1] |
|
|
| @torch.no_grad() |
| def search(self, qs, K=20): |
| """-> one list per query of (url, log_prob, table_row), sorted by score.""" |
| with torch.autocast(self.dev_type, dtype=torch.bfloat16): |
| P = self.prompts(qs) |
| B = len(P) |
| ninf = float("-inf") |
| lg0 = (self.last_hidden(P) @ self.W[BASE_V:BASE_V + V].t()).float() |
| lp0 = lg0.masked_fill(~self.ok0, ninf).log_softmax(-1).masked_fill(~self.ok0, ninf) |
| k0 = min(K, int(self.ok0.sum())) |
| s0, c0 = lp0.topk(k0, dim=1) |
| |
| P1 = [p + [BASE_V + c] for p, cs in zip(P, c0.tolist()) for c in cs] |
| lg1 = (self.last_hidden(P1) @ self.W[BASE_V + V:BASE_V + 2 * V].t()).float() |
| cc = c0.reshape(-1) |
| lo = torch.searchsorted(self.keys, cc * V) |
| n = torch.searchsorted(self.keys, (cc + 1) * V) - lo |
| r = torch.repeat_interleave(torch.arange(len(cc), device=self.dev), n) |
| j = torch.arange(int(n.sum()), device=self.dev) - torch.repeat_interleave(n.cumsum(0) - n, n) + lo[r] |
| ok1 = torch.zeros_like(lg1, dtype=torch.bool) |
| ok1[r, self.keys[j] - cc[r] * V] = True |
| lp1 = lg1.masked_fill(~ok1, ninf).log_softmax(-1).masked_fill(~ok1, ninf) |
| cand = (s0.reshape(-1, 1) + lp1).view(B, k0 * V) |
| s, flat = cand.topk(min(K, k0 * V), dim=1) |
| key = c0.gather(1, flat // V) * V + flat % V |
| row = torch.searchsorted(self.keys, key) |
| s, row = s.float().cpu().numpy(), row.cpu().numpy() |
| out = [] |
| for b in range(B): |
| m = np.isfinite(s[b]) |
| rr = row[b][m] |
| out.append(list(zip(self.url.take(rr), s[b][m].tolist(), rr.tolist()))) |
| return out |
|
|
|
|
| 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 batching and 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("--device", default="cuda") |
| ap.add_argument("--check", action="store_true") |
| a = ap.parse_args() |
| t0 = time.time() |
| gr = GR(a.model, a.device) |
| print(f"Loaded {gr.n:,} candidate URLs in {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() |
|
|
