# Copyright 2026 Octen. Licensed under the Apache License, Version 2.0 (see LICENSE-CODE).
"""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                                   # <|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 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  # noqa: F401
                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)                         # row number = position in keys
        ok0 = np.zeros(V, bool)
        ok0[np.unique(key // V)] = True
        self.ok0 = torch.from_numpy(ok0).to(device)                          # valid level-0 codes
        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):                                         # left padding: the last position is each sequence's last token
            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()                  # [B, V]
            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)                                                            # [B, k0]
            # Level 1: one sequence per (query, level-0 candidate) = prompt + that SID token
            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()          # [B*k0, V]
            cc = c0.reshape(-1)
            lo = torch.searchsorted(self.keys, cc * V)
            n = torch.searchsorted(self.keys, (cc + 1) * V) - lo                                    # number of valid children per prefix
            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()