sqljev 0.1.0__py3-none-any.whl

sqljev/finetune.py

@@ -0,0 +1,367 @@
+"""Fine-tune Laya on labelled SQL rows, on one GPU (a free Colab T4 is enough), and publish the checkpoint.
+
+    sqljev dataset "$DB_URL" "SELECT subject, body, team FROM tickets" --label team \\
+        --choice "which team should handle this?" --test-fraction 0.2 -o tickets.jsonl
+    sqljev finetune tickets.train.jsonl --out checkpoints/tickets --epochs 3
+    sqljev eval tickets.test.jsonl --model checkpoints/tickets
+    sqljev publish checkpoints/tickets --repo your-org/laya-tickets          # Hugging Face Hub
+
+Training data is what `sqljev dataset` writes: one row per line with the state (the row, as the model sees it at
+query time), the question built by core.laya_question and the expected answer, so the checkpoint learns exactly
+the questions sqljev will ask it.
+
+The training loop follows Laya's own fine-tuning notebook (Convai Innovations, Apache 2.0,
+https://github.com/NandhaKishorM/laya/tree/main/notebooks): policy-gradient updates against a proper scoring
+rule plus soft cross-entropy, then per-question-type temperature calibration on a held-out slice. It is
+single-GPU here, with gradient checkpointing and mixed precision, so it fits a 16 GB T4.
+"""
+import json
+import os
+import random
+import re
+import time
+
+from .core import JevError, check_question, laya_question, to_row_json
+
+LAYA_BASE = "convaiinnovations/laya"
+SUBFOLDERS = {"english": None, "multilingual": "multilingual", "typed-decisions": "typed-decisions"}
+
+
+# ---------------------------------------------------------------- examples
+
+def examples_from_rows(rows, question, kind, options=None, label=None, columns=None, drop_nulls=True):
+    """Labelled rows (dicts) -> training/eval records. The label column never reaches the state. For choice
+    questions without options, the options are the distinct labels."""
+    rows = [r for r in rows if r.get(label) is not None]
+    if kind == "choice" and not options:
+        options = sorted({str(r[label]) for r in rows})
+    kind, opts = check_question(kind, options)
+    q = {"q": laya_question(kind, question, opts)}
+    out = []
+    for r in rows:
+        lab = r[label]
+        if kind == "noul":
+            expected = _truthy(lab)
+        elif kind == "choice":
+            expected = str(lab)
+            if expected not in opts:
+                continue
+        else:
+            expected = opts.index(str(lab)) if str(lab) in opts else int(lab)
+        view = {k: v for k, v in r.items() if k != label and (not columns or k in columns)}
+        out.append({"state": json.loads(to_row_json(view, drop_nulls)), "questions": q, "expected": {"q": expected}})
+    return out
+
+
+def split(records, test_fraction=0.2, seed=0):
+    """Deterministic train/test split."""
+    idx = list(range(len(records)))
+    random.Random(seed).shuffle(idx)
+    n_test = int(len(records) * test_fraction)
+    test = set(idx[:n_test])
+    return [r for i, r in enumerate(records) if i not in test], [r for i, r in enumerate(records) if i in test]
+
+
+def write_jsonl(records, path):
+    with open(path, "w") as f:
+        for r in records:
+            f.write(json.dumps(r, ensure_ascii=False, default=str) + "\n")
+
+
+def read_jsonl(path):
+    with open(path) as f:
+        return [json.loads(line) for line in f if line.strip() and not line.startswith("#")]
+
+
+def _truthy(v):
+    if isinstance(v, bool):
+        return v
+    s = str(v).strip().lower()
+    if s in ("1", "true", "t", "yes", "y"):
+        return True
+    if s in ("0", "false", "f", "no", "n"):
+        return False
+    raise JevError("sqljev: label %r is not a boolean" % (v,))
+
+
+def accuracy(records, **settings):
+    """Share of expected answers the model gets right on `records` (a list or a JSONL path), with any engine
+    settings (model=<checkpoint>, device=..., backend=...). Returns {"accuracy", "decisions", "seconds"}."""
+    from .core import Jev
+    records = read_jsonl(records) if isinstance(records, str) else records
+    jev = Jev(**settings)
+    groups = {}
+    for rec in records:
+        for qid, q in rec["questions"].items():
+            crit = q.get("criteria")
+            opts = list(crit) if crit else None
+            groups.setdefault((q["type"], q["instructions"], json.dumps(opts)), []).append(
+                (rec["state"], rec["expected"][qid]))
+    t0, right, total = time.time(), 0, 0
+    for (kind, instr, opts_json), items in groups.items():
+        prefix = "Is it true that "
+        query = instr[len(prefix):-1] if kind == "noul" and instr.startswith(prefix) and instr.endswith("?") else instr
+        answers = jev.evaluate([s for s, _ in items], query, kind, json.loads(opts_json))
+        for (_, exp), a in zip(items, answers):
+            if kind == "noul":
+                right += (a["noul"] >= 0.5) == bool(exp)
+            elif kind == "choice":
+                right += a["choice"] == exp
+            else:
+                right += round(a["score"]) == int(exp)
+            total += 1
+    return {"accuracy": round(right / total, 4) if total else 0.0, "decisions": total,
+            "seconds": round(time.time() - t0, 1)}
+
+
+# ---------------------------------------------------------------- training
+
+def base_dir(base=LAYA_BASE):
+    """Local directory of a Laya checkpoint: a path, a named base checkpoint, or a Hub id."""
+    if os.path.isdir(base):
+        return base
+    from huggingface_hub import snapshot_download
+    sub = SUBFOLDERS.get(base, None)
+    repo = LAYA_BASE if base in SUBFOLDERS else base
+    d = snapshot_download(repo)
+    try:
+        from laya.agent import _fix_tokenizer_config
+        _fix_tokenizer_config(d)
+    except Exception:       # noqa: BLE001 -- older/newer laya: the files are already usable
+        pass
+    return os.path.join(d, sub) if sub else d
+
+
+def _items(records, tok, cfg):
+    from laya.common import QTYPES, build_sequence, render_options
+    items, skipped = [], 0
+    for rec in records:
+        for qid, q in rec["questions"].items():
+            t, crit = q["type"], q.get("criteria") or {}
+            exp = rec["expected"][qid]
+            if t == "noul":
+                target = [0.0, 1.0] if exp else [1.0, 0.0]
+            elif t == "choice":
+                keys = list(crit.keys()) if isinstance(crit, dict) else list(crit)
+                target = [float(k == exp) for k in keys]
+            else:
+                target = [float(i == int(exp)) for i in range(len(crit))]
+            seq, markers = build_sequence(tok, rec["state"], {"t": t, "ins": q["instructions"], "crit": crit},
+                                          cfg["max_len"], cfg["head_max_len"])
+            if len(markers) != len(render_options({"t": t, "crit": crit})) or sum(target) != 1:
+                skipped += 1
+                continue
+            items.append({"ids": seq, "markers": markers, "qtype": QTYPES[t], "target": target,
+                          "label": target.index(1.0)})
+    return items, skipped
+
+
+def _collate(items, pad_id):
+    import torch
+    n, L = len(items), max(len(it["ids"]) for it in items)
+    kmax = max(len(it["markers"]) for it in items)
+    ids = torch.full((n, L), pad_id, dtype=torch.long)
+    att = torch.zeros((n, L), dtype=torch.long)
+    mpos = torch.zeros((n, kmax), dtype=torch.long)
+    mmask = torch.zeros((n, kmax), dtype=torch.bool)
+    target = torch.zeros((n, kmax), dtype=torch.float32)
+    for i, it in enumerate(items):
+        ids[i, :len(it["ids"])] = torch.tensor(it["ids"])
+        att[i, :len(it["ids"])] = 1
+        k = len(it["markers"])
+        mpos[i, :k] = torch.tensor(it["markers"])
+        mmask[i, :k] = True
+        target[i, :k] = torch.tensor(it["target"])
+    return ids, att, mpos, mmask, target, torch.tensor([it["qtype"] for it in items])
+
+
+def _fit_temperature(pairs):
+    import torch
+    if len(pairs) < 10:
+        return 1.0
+    kmax = max(len(z) for z, _ in pairs)
+    Z = torch.full((len(pairs), kmax), -1e4)
+    T = torch.zeros((len(pairs), kmax))
+    for i, (z, t) in enumerate(pairs):
+        Z[i, :len(z)] = torch.tensor(z)
+        T[i, :len(t)] = torch.tensor(t)
+    log_t = torch.zeros(1, requires_grad=True)
+    opt = torch.optim.LBFGS([log_t], lr=0.1, max_iter=100)
+
+    def closure():
+        opt.zero_grad()
+        loss = -(T * torch.log_softmax(Z / log_t.exp(), -1)).sum(-1).mean()
+        loss.backward()
+        return loss
+    opt.step(closure)
+    return float(torch.clamp(log_t.exp(), 0.5, 5.0).item())
+
+
+def finetune(train, out_dir, base=LAYA_BASE, epochs=3, micro_batch=8, grad_accum=8, lr_encoder=2.5e-5,
+             lr_head=1e-4, group_size=4, train_layers=None, device=None, seed=0, log=print):
+    """Fine-tune a Laya checkpoint on `train` (records or a JSONL path) and save it to `out_dir` in Laya's
+    format, loadable with laya.load(out_dir) and usable as SQLJEV_MODEL=out_dir.
+
+    train_layers: train only the top N encoder layers plus the decision head, freezing the embeddings and the
+    layers below (low-memory mode for GPUs under ~10 GB free). None trains everything."""
+    import torch
+    from laya.common import build_model, proper_reward
+    from safetensors.torch import load_file, save_file
+    from transformers import AutoTokenizer
+
+    records = read_jsonl(train) if isinstance(train, str) else list(train)
+    if not records:
+        raise JevError("sqljev: no training records")
+    device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+    src = base_dir(base)
+    with open(os.path.join(src, "rl_agent_config.json")) as f:
+        cfg = json.load(f)
+    tok = AutoTokenizer.from_pretrained(os.path.join(src, "tokenizer"))
+    items, skipped = _items(records, tok, cfg)
+    if not items:
+        raise JevError("sqljev: none of the records could be encoded for this checkpoint")
+    rng = random.Random(seed)
+    rng.shuffle(items)
+    n_cal = min(400, len(items) // 10)
+    calib, items = items[:n_cal], items[n_cal:]
+    log("sqljev finetune: %d training items, %d held out for calibration%s, base %s, device %s"
+        % (len(items), len(calib), ", %d skipped" % skipped if skipped else "", base, device))
+
+    torch.manual_seed(seed)
+    model = build_model(cfg, encoder_dir=os.path.join(src, "encoder"))
+    model.load_state_dict(load_file(os.path.join(src, "model.safetensors")), strict=True)
+    if device.type == "cuda":
+        model.encoder.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+        model.head_checkpointing = True
+    model.to(device).train()
+    if train_layers is not None:
+        n_layers = 1 + max(int(m.group(1)) for n, _ in model.named_parameters()
+                           for m in [re.match(r"encoder\.layers\.(\d+)\.", n)] if m)
+        keep = n_layers - int(train_layers)
+        for n, p in model.named_parameters():
+            m = re.match(r"encoder\.layers\.(\d+)\.", n)
+            if n.startswith("encoder.embeddings.") or (m and int(m.group(1)) < keep):
+                p.requires_grad_(False)
+        log("sqljev finetune: training the top %d of %d encoder layers and the head" % (min(int(train_layers), n_layers), n_layers))
+
+    use_cuda = device.type == "cuda"
+    amp_dtype = torch.bfloat16 if use_cuda and torch.cuda.is_bf16_supported() else torch.float16
+    scaler = torch.amp.GradScaler("cuda", enabled=use_cuda and amp_dtype == torch.float16)
+    enc = [p for n, p in model.named_parameters() if n.startswith("encoder.") and p.requires_grad]
+    head = [p for n, p in model.named_parameters() if not n.startswith("encoder.") and p.requires_grad]
+    opt = torch.optim.AdamW([{"params": enc, "lr": lr_encoder}, {"params": head, "lr": lr_head}], weight_decay=0.01)
+    updates = max(1, -(-len(items) // (micro_batch * grad_accum)) * epochs)
+    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=updates, eta_min=1e-6)
+    t0 = time.time()
+    for epoch in range(epochs):
+        rng.shuffle(items)
+        sigma = 0.4 + (0.1 - 0.4) * (epoch / max(1, epochs - 1))
+        total, steps = 0.0, 0
+        opt.zero_grad(set_to_none=True)
+        for b in range(0, len(items), micro_batch):
+            ids, att, mpos, mmask, target, qtype = (x.to(device) for x in _collate(items[b:b + micro_batch], tok.pad_token_id))
+            with torch.autocast(device.type, dtype=amp_dtype, enabled=use_cuda):
+                logits, act = model(ids, att, mpos, mmask, qtype)
+            logits = logits.float()
+            k = mmask.sum(-1, keepdim=True).float()
+            eps = torch.randn((group_size,) + logits.shape, device=device) * sigma * mmask
+            eps = (eps - eps.sum(-1, keepdim=True) / k) * mmask
+            z = logits.detach().unsqueeze(0) + eps
+            q = torch.softmax(z.masked_fill(~mmask, -1e4), -1)
+            with torch.no_grad():
+                r = proper_reward(q, target.unsqueeze(0), qtype, mmask, w_sph=0.75, w_rps=1.0)
+                adv = (r - r.mean(0, keepdim=True))
+                adv = adv / (adv.std() + 1e-6)
+            logp = -(((z - logits.unsqueeze(0)) ** 2) * mmask).sum(-1) / (2 * sigma ** 2)
+            loss_ce = -(target * torch.log_softmax(logits.masked_fill(~mmask, -1e4), -1)).sum(-1).mean()
+            loss = (-(adv * logp).mean() + loss_ce) / grad_accum + 0.0 * act.sum()
+            scaler.scale(loss).backward()
+            steps += 1
+            total += loss.item() * grad_accum
+            if steps % grad_accum == 0 or b + micro_batch >= len(items):
+                scaler.unscale_(opt)
+                torch.nn.utils.clip_grad_norm_(enc + head, 1.0)
+                scaler.step(opt)
+                scaler.update()
+                sched.step()
+                opt.zero_grad(set_to_none=True)
+            if steps % 100 == 0:
+                log("  epoch %d/%d  step %d/%d  loss %.4f  %.0fs"
+                    % (epoch + 1, epochs, steps, -(-len(items) // micro_batch), total / steps, time.time() - t0))
+        log("  epoch %d/%d done, mean loss %.4f, %.0fs" % (epoch + 1, epochs, total / max(1, steps), time.time() - t0))
+
+    # Calibrate one temperature per question type on items the run never trained on.
+    model.eval()
+    preds = []
+    with torch.no_grad():
+        for b in range(0, len(calib), 16):
+            chunk = calib[b:b + 16]
+            ids, att, mpos, mmask, target, qtype = (x.to(device) for x in _collate(chunk, tok.pad_token_id))
+            with torch.autocast(device.type, dtype=amp_dtype, enabled=use_cuda):
+                lg, _ = model(ids, att, mpos, mmask, qtype)
+            lg = lg.float().cpu().numpy()
+            preds += [(it["qtype"], lg[i, :len(it["markers"])].tolist(), it["target"]) for i, it in enumerate(chunk)]
+    temps = list(cfg.get("temperature") or [1.2, 1.2, 1.2])
+    for qt in range(3):
+        sel = [(z, t) for q, z, t in preds if q == qt]
+        if len(sel) >= 10:
+            temps[qt] = _fit_temperature(sel)
+
+    os.makedirs(out_dir, exist_ok=True)
+    save_file({k: v.half().contiguous().cpu() for k, v in model.state_dict().items()},
+              os.path.join(out_dir, "model.safetensors"))
+    model.encoder.config.save_pretrained(os.path.join(out_dir, "encoder"))
+    tok.save_pretrained(os.path.join(out_dir, "tokenizer"))
+    cfg.update(fine_tuned=True, model_name="laya-sqljev", temperature=temps)
+    cfg.pop("temperature_by_options", None)
+    with open(os.path.join(out_dir, "rl_agent_config.json"), "w") as f:
+        json.dump(cfg, f, indent=2)
+    questions = sorted({q["instructions"] for r in records for q in r["questions"].values()})
+    with open(os.path.join(out_dir, "sqljev_finetune.json"), "w") as f:
+        json.dump({"base": base, "records": len(records), "items": len(items), "epochs": epochs,
+                   "questions": questions, "seconds": round(time.time() - t0, 1)}, f, indent=2)
+    log("sqljev finetune: saved to %s (%.0fs)" % (out_dir, time.time() - t0))
+    return out_dir
+
+
+# ---------------------------------------------------------------- publishing
+
+def publish(out_dir, repo_id, token=None, private=True, metrics=None):
+    """Upload a fine-tuned checkpoint to the Hugging Face Hub with a model card. Every database then uses it
+    with SQLJEV_MODEL=<repo_id>."""
+    from huggingface_hub import HfApi
+    meta_path = os.path.join(out_dir, "sqljev_finetune.json")
+    meta = json.load(open(meta_path)) if os.path.exists(meta_path) else {}
+    if metrics:
+        meta["metrics"] = metrics
+        with open(meta_path, "w") as f:
+            json.dump(meta, f, indent=2)
+    card = """---
+license: apache-2.0
+base_model: %s
+tags: [laya, sqljev, sql, decision-model]
+---
+
+# %s
+
+A [Laya](https://github.com/NandhaKishorM/laya) checkpoint fine-tuned with [sqljev](https://github.com/singhpratech/sqljev)
+on labelled SQL rows.
+
+Questions it was trained on:
+%s
+
+%s
+Use it from any database sqljev supports:
+
+```bash
+SQLJEV_MODEL=%s sqljev gateway --host 0.0.0.0
+```
+""" % (meta.get("base", LAYA_BASE), repo_id.split("/")[-1],
+       "\n".join("- %s" % q for q in meta.get("questions", [])),
+       ("Held-out accuracy: %s\n" % json.dumps(metrics)) if metrics else "", repo_id)
+    with open(os.path.join(out_dir, "README.md"), "w") as f:
+        f.write(card)
+    api = HfApi(token=token)
+    api.create_repo(repo_id, private=private, exist_ok=True)
+    api.upload_folder(repo_id=repo_id, folder_path=out_dir, commit_message="sqljev finetune")
+    return "https://huggingface.co/" + repo_id

sqljev/gateway.py

@@ -0,0 +1,215 @@
+"""sqljev gateway: one HTTP service that answers the batch protocols databases use to call out.
+
+    sqljev gateway --port 8765            # Laya in-process (backend 'local') by default
+
+Routes (all POST, JSON):
+
+  /v1/eval               sqljev native: {"question", "kind", "options", "rows": [...]} -> {"answers": [...]}
+  /sqlserver             same body; called by jev.judge through sp_invoke_external_rest_endpoint
+  /bigquery              BigQuery remote functions: {"calls": [[args]], "userDefinedContext": {"fn": ...}}
+                         -> {"replies": [...]}
+  /snowflake/<fn>        Snowflake external / service functions: {"data": [[rownum, args...]]}
+                         -> {"data": [[rownum, result]]}
+  /redshift              Redshift Lambda UDF payload over HTTP (the same handler as sqljev.aws_lambda)
+  /v1/systemone          TypeSafe Jev's wire protocol, for Jev clients such as pg-jev: batch requests
+                         {"state": {"condition", "rows": [...]}, "questions": {"r0": ...}} are answered by the
+                         engine, so `SET jev.api_url = 'http://gateway:8765/v1/systemone'` runs pg-jev on Laya
+  GET /health, GET /stats
+
+Each database already sends rows in batches (BigQuery up to max_batching_rows, Snowflake/Redshift in their
+own batch sizes), so every request becomes one engine.call(): rows grouped by question, de-duplicated,
+looked up in the cache and the misses judged in shared forward passes.
+
+Auth: set SQLJEV_GATEWAY_TOKEN and send it as "Authorization: Bearer <token>" or "X-Jev-Token: <token>".
+BigQuery cannot send headers: run the gateway on Cloud Run with IAM auth instead.
+"""
+import hmac
+import json
+import os
+import re
+import sys
+import traceback
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+from .core import FUNCTIONS, JevError, Jev, __version__, function_name
+
+MAX_BODY_BYTES = 64 * 1024 * 1024
+
+
+def handle_eval(engine, body):
+    kind = body.get("kind") or "noul"
+    rows = body.get("rows")
+    if not isinstance(rows, list) or not body.get("question"):
+        raise JevError("sqljev: body needs 'question' and a 'rows' array")
+    return {"answers": engine.evaluate(rows, body["question"], kind, body.get("options"))}
+
+
+def handle_bigquery(engine, body):
+    ctx = body.get("userDefinedContext") or {}
+    fn = ctx.get("fn") or ctx.get("function")
+    if not fn:
+        raise JevError("sqljev: set user_defined_context = [('fn', 'jev_prob')] on the remote function")
+    return {"replies": _json_safe(engine.call(fn, body.get("calls") or []))}
+
+
+def handle_snowflake(engine, body, fn):
+    data = body.get("data") or []
+    results = engine.call(fn, [r[1:] for r in data])    # VARIANT takes jev_eval's object as is
+    return {"data": [[r[0], v] for r, v in zip(data, results)]}
+
+
+def handle_redshift(engine, event):
+    """Redshift Lambda UDF protocol. Errors are reported in-band, as Redshift expects."""
+    try:
+        fn = event.get("external_function") or ""
+        args = event.get("arguments") or []
+        results = engine.call(fn, args)
+        return {"success": True, "num_records": len(results), "results": _json_safe(results)}
+    except Exception as e:    # noqa: BLE001 -- Redshift shows error_msg to the user
+        return {"success": False, "error_msg": str(e)}
+
+
+_ROW_REF = re.compile(r"`rows\[(\d+)\]`")
+_JEV_PREFIX = {"score": "Rate the record `rows[%d]`: ", "choice": "For the record `rows[%d]`: "}
+
+
+def handle_systemone(engine, body):
+    """Answer a Jev batch request (one shared state with `rows`, one question per row, the format pg-jev and
+    sqljev's own jev backend send) with the engine: questions are grouped back into (question, rows) pairs."""
+    state, questions = body.get("state"), body.get("questions")
+    if not isinstance(state, dict) or not isinstance(state.get("rows"), list) or not isinstance(questions, dict):
+        raise JevError("sqljev: /v1/systemone accepts batch requests: state {'rows': [...]} with one question per "
+                       "row that names `rows[i]` (as pg-jev sends); use /v1/eval for everything else")
+    rows, groups = state["rows"], {}
+    for qid, q in questions.items():
+        m = _ROW_REF.search(q.get("instructions") or "")
+        if not m or int(m.group(1)) >= len(rows):
+            raise JevError("sqljev: question %r does not name a row as `rows[i]`" % qid)
+        i, kind = int(m.group(1)), q.get("type")
+        if kind == "noul":
+            query, opts = state.get("condition") or "", None
+        elif kind in _JEV_PREFIX:
+            prefix = _JEV_PREFIX[kind] % i
+            instr = q["instructions"]
+            query = instr[len(prefix):] if instr.startswith(prefix) else instr
+            crit = q.get("criteria")
+            opts = list(crit.keys()) if isinstance(crit, dict) else list(crit or [])
+        else:
+            raise JevError("sqljev: unknown question type %r" % kind)
+        groups.setdefault((kind, query, tuple(opts or ())), []).append((qid, i))
+    answers = {}
+    for (kind, query, opts), items in groups.items():
+        got = engine.evaluate([rows[i] for _, i in items], query, kind, list(opts) or None)
+        answers.update({qid: a for (qid, _), a in zip(items, got)})
+    return {"model": "sqljev-" + engine.cfg["backend"], "answers": answers,
+            "usage": {"input_tokens": 0, "output_tokens": 0}}
+
+
+def _json_safe(values):
+    # jev_eval answers go back to warehouses as JSON text columns (BigQuery JSON/STRING, Snowflake VARIANT
+    # accepts objects directly but a string is safe everywhere).
+    return [json.dumps(v) if isinstance(v, dict) else v for v in values]
+
+
+class Handler(BaseHTTPRequestHandler):
+    protocol_version = "HTTP/1.1"
+    engine = None
+    token = None
+    server_version = "sqljev/" + __version__
+
+    def log_message(self, fmt, *args):
+        if os.environ.get("SQLJEV_GATEWAY_LOG"):
+            sys.stderr.write("sqljev gateway: " + fmt % args + "\n")
+
+    def _send(self, code, obj):
+        data = json.dumps(obj).encode()
+        self.send_response(code)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", str(len(data)))
+        self.end_headers()
+        self.wfile.write(data)
+
+    def _authorized(self):
+        if not self.token:
+            return True
+        got = self.headers.get("X-Jev-Token") or ""
+        auth = self.headers.get("Authorization") or ""
+        if auth.startswith("Bearer "):
+            got = got or auth[7:]
+        return hmac.compare_digest(got.encode(), self.token.encode())
+
+    def do_GET(self):
+        if self.path == "/health":
+            return self._send(200, {"status": "ok", "version": __version__, "backend": self.engine.cfg["backend"],
+                                    "functions": sorted(FUNCTIONS)})
+        if self.path == "/stats":
+            if not self._authorized():
+                return self._send(401, {"error": "unauthorized"})
+            return self._send(200, self.engine.stats())
+        self._send(404, {"error": "not found"})
+
+    def do_POST(self):
+        n = int(self.headers.get("Content-Length") or 0)
+        if n > MAX_BODY_BYTES:
+            return self._send(413, {"error": "body too large"})
+        raw = self.rfile.read(n)
+        if not self._authorized():
+            return self._send(401, {"error": "unauthorized"})
+        path = self.path.split("?", 1)[0].rstrip("/")
+        try:
+            body = json.loads(raw or b"{}")
+            if path in ("/v1/eval", "/sqlserver"):
+                return self._send(200, handle_eval(self.engine, body))
+            if path == "/bigquery":
+                return self._send(200, handle_bigquery(self.engine, body))
+            if path.startswith("/snowflake/"):
+                return self._send(200, handle_snowflake(self.engine, body, function_name(path.rsplit("/", 1)[1])))
+            if path == "/redshift":
+                return self._send(200, handle_redshift(self.engine, body))
+            if path == "/v1/systemone":
+                return self._send(200, handle_systemone(self.engine, body))
+            return self._send(404, {"error": "not found: " + path})
+        except (JevError, ValueError, KeyError, TypeError, IndexError) as e:
+            # 400 is not retried by BigQuery/Snowflake: validation errors surface as query errors.
+            return self._send(400, {"error": str(e), "errorMessage": str(e)})
+        except Exception as e:     # noqa: BLE001
+            traceback.print_exc()
+            return self._send(500, {"error": "internal error: %s" % type(e).__name__,
+                                    "errorMessage": "sqljev gateway internal error"})
+
+
+def make_server(host="127.0.0.1", port=8765, engine=None, token=None):
+    handler = type("SqlJevHandler", (Handler,), {
+        "engine": engine or Jev(),
+        "token": token if token is not None else os.environ.get("SQLJEV_GATEWAY_TOKEN"),
+    })
+    ThreadingHTTPServer.daemon_threads = True
+    return ThreadingHTTPServer((host, port), handler)
+
+
+def serve(host="127.0.0.1", port=8765, engine=None, token=None, certfile=None, keyfile=None):
+    """certfile/keyfile serve HTTPS directly (sp_invoke_external_rest_endpoint only calls https:// URLs);
+    behind a TLS-terminating proxy or Cloud Run, leave them unset."""
+    srv = make_server(host, port, engine, token)
+    scheme = "http"
+    if certfile:
+        import ssl
+        ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+        ctx.load_cert_chain(certfile, keyfile)
+        accept = srv.get_request
+
+        def get_request():
+            # Handshake lazily, in the worker thread: a slow or failing client never blocks accept().
+            sock, addr = accept()
+            return ctx.wrap_socket(sock, server_side=True, do_handshake_on_connect=False), addr
+        srv.get_request = get_request
+        scheme = "https"
+    eng = srv.RequestHandlerClass.engine
+    print("sqljev gateway %s on %s://%s:%d (backend %s)" % (__version__, scheme, host, port, eng.cfg["backend"]),
+          file=sys.stderr)
+    try:
+        srv.serve_forever()
+    except KeyboardInterrupt:
+        pass
+    finally:
+        srv.server_close()

sqljev/spark.py

@@ -0,0 +1,64 @@
+"""Spark / Databricks: the jev functions as pandas (Arrow) UDFs usable from SQL.
+
+    import sqljev.spark
+    sqljev.spark.register(spark)                                  # Laya in-process on each executor
+    # or: sqljev.spark.register(spark, backend="gateway", api_url="https://.../v1/eval",
+    #                           api_key=dbutils.secrets.get("sqljev", "token"))
+
+    SELECT * FROM tickets WHERE jev(to_json(struct(*)), 'the customer is angry')
+
+Spark passes Arrow batches (spark.sql.execution.arrow.maxRecordsPerBatch, 10,000 rows by default) to each
+UDF; each batch becomes one engine.call(). One engine per Python worker keeps the model loaded and the
+answer cache warm across batches. On a GPU cluster set device="cuda".
+"""
+_ENGINES = {}
+
+_RETURNS = {"jev": "boolean", "jev_prob": "double", "jev_score": "double", "jev_score_norm": "double",
+            "jev_choice": "string", "jev_confidence": "double", "jev_eval": "string"}
+_ARITY = {"jev": 2, "jev_prob": 2, "jev_score": 3, "jev_score_norm": 3, "jev_choice": 3,
+          "jev_confidence": 4, "jev_eval": 4}
+
+
+def _engine(settings):
+    from .core import Jev
+    key = tuple(sorted(settings.items()))
+    eng = _ENGINES.get(key)
+    if eng is None:
+        eng = _ENGINES[key] = Jev(**settings)
+    return eng
+
+
+def _make(name, settings):
+    import json
+    import pandas as pd
+    from pyspark.sql.functions import pandas_udf
+
+    n = _ARITY[name]
+
+    def run(*cols):
+        calls = list(zip(*(c.tolist() for c in cols)))
+        out = _engine(settings).call(name, calls)
+        if name == "jev_eval":
+            out = [None if v is None else json.dumps(v) for v in out]
+        return pd.Series(out, dtype="object" if name in ("jev_choice", "jev_eval", "jev") else "float64")
+
+    # pandas_udf needs a fixed-arity signature with type hints.
+    if n == 2:
+        def f(a: pd.Series, b: pd.Series) -> pd.Series:
+            return run(a, b)
+    elif n == 3:
+        def f(a: pd.Series, b: pd.Series, c: pd.Series) -> pd.Series:
+            return run(a, b, c)
+    else:
+        def f(a: pd.Series, b: pd.Series, c: pd.Series, d: pd.Series) -> pd.Series:
+            return run(a, b, c, d)
+    return pandas_udf(f, _RETURNS[name])
+
+
+def register(spark, prefix="", **settings):
+    """Register jev, jev_prob, jev_score, jev_score_norm, jev_choice, jev_confidence, jev_eval for SQL.
+    Settings (backend, api_url, api_key, model, device, ...) are shipped to the executors."""
+    fns = {}
+    for name in _RETURNS:
+        fns[name] = spark.udf.register(prefix + name, _make(name, dict(settings)))
+    return fns