# Copyright 2026 Octen. Licensed under the Apache License, Version 2.0 (see LICENSE-CODE).
"""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                                   # <|endoftext|>, used only for left padding (masked out)


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()                  # the 65,536 rows of level 0
        self.W1 = W[BASE_V + V:BASE_V + 2 * V].to(ldt).contiguous()          # the 65,536 rows of level 1
        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)))  # valid children of prefix c = rows [start[c], start[c+1])
        self.ok0 = self.start[1:] > self.start[:-1]                            # valid level-0 codes
        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):                                            # left padding: the last position is each sequence's last token
            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)                          # [B, k0]
        # Level 1: copy the KV cache per beam (the k0 beams of query b are adjacent); each beam feeds only its SID token
        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)               # position of the new token = number of real tokens
        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()                              # [B*k0, V]; only the valid children's columns are used
        cc = c0.reshape(-1)
        lo = self.start[cc]
        n = self.start[cc + 1] - lo                                           # number of valid children per beam
        r = torch.repeat_interleave(torch.arange(B * k0), n)                  # beam of each child
        j = torch.arange(int(n.sum())) - torch.repeat_interleave(n.cumsum(0) - n, n) + lo[r]   # table row of each child
        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]                                   # sum of the two log-probabilities
        off = [0] + n.view(B, k0).sum(1).cumsum(0).tolist()                    # children of one query are adjacent in sc
        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()