File size: 4,458 Bytes
e46c127 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """Client for an openjev checkpoint served by SGLang (serve_sglang.sh). Same interface as
`OpenJevCrossEncoder.predict/rerank`, plus `--check` to compare the server against transformers.
python sglang_client.py --premise "A man is playing a guitar." --hypothesis "Someone is making music."
python sglang_client.py --check ckpt/qwen3.5-0.8b-nli-v2s-long # max |p_sglang - p_hf| on a few pairs
"""
from __future__ import annotations
import argparse
import base64
from concurrent.futures import ThreadPoolExecutor
import numpy as np
import requests
CON, ENT, NEU = 0, 1, 2
LABELS = ["contradiction", "entailment", "neutral"]
DEFAULT_TEMPLATE = "Premise: {premise}\nHypothesis: {hypothesis}"
IMG_MARK = "<<IMG>>"
IMG_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>" # SGLang expands the single pad to the patch count
class OpenJevSGLang:
def __init__(self, url: str = "http://127.0.0.1:30000", template: str = DEFAULT_TEMPLATE, bs: int = 32,
workers: int = 16):
# `workers` requests of `bs` pairs are in flight at once; the server batches across them
self.url, self.template, self.bs, self.workers = url.rstrip("/"), template, bs, workers
def _post(self, payload):
r = requests.post(f"{self.url}/classify", json=payload, timeout=600)
r.raise_for_status()
out = r.json()
return [o["embedding"] for o in (out if isinstance(out, list) else [out])]
def logits(self, pairs, images=None) -> np.ndarray:
"""Raw [contradiction, entailment, neutral] logits. `images[i]` is a path / URL / bytes or None and is
spliced into premise i at the `<<IMG>>` marker (appended if the marker is absent)."""
payloads = []
for i in range(0, len(pairs), self.bs):
texts, imgs = [], []
for j, (p, h) in enumerate(pairs[i:i + self.bs]):
img = images[i + j] if images else None
if img is not None:
p = p.replace(IMG_MARK, IMG_BLOCK) if IMG_MARK in p else f"{p.rstrip()} {IMG_BLOCK}"
img = base64.b64encode(img).decode() if isinstance(img, bytes) else img
texts.append(self.template.format(premise=p.strip(), hypothesis=h.strip()))
imgs.append(img)
payload = {"text": texts}
if any(x is not None for x in imgs):
payload["image_data"] = imgs
payloads.append(payload)
with ThreadPoolExecutor(self.workers) as ex:
out = [z for part in ex.map(self._post, payloads) for z in part]
return np.asarray(out, np.float32)
def predict(self, pairs, images=None) -> np.ndarray:
"""Softmax probabilities [contradiction, entailment, neutral] per pair."""
z = self.logits(pairs, images)
z = np.exp(z - z.max(-1, keepdims=True))
return z / z.sum(-1, keepdims=True)
def rerank(self, question: str, options, hyp_fmt: str = "The correct answer is: {}") -> int:
return int(self.predict([(question, hyp_fmt.format(o)) for o in options])[:, ENT].argmax())
CHECK_PAIRS = [
("A man is playing a guitar.", "Someone is making music."),
("A man is playing a guitar.", "The man is asleep."),
("A man is playing a guitar.", "The man is a professional musician."),
("The Eiffel Tower is in Paris. " * 40, "The Eiffel Tower is in Berlin."),
("What is the capital of France?", "The correct answer is: Paris"),
]
if __name__ == "__main__":
ap = argparse.ArgumentParser()
ap.add_argument("--url", default="http://127.0.0.1:30000")
ap.add_argument("--premise")
ap.add_argument("--hypothesis")
ap.add_argument("--image")
ap.add_argument("--check", metavar="CKPT", help="compare against transformers on this checkpoint")
a = ap.parse_args()
client = OpenJevSGLang(a.url)
if a.check:
from modeling_openjev import OpenJevCrossEncoder
ours, ref = client.predict(CHECK_PAIRS), OpenJevCrossEncoder(a.check).predict(CHECK_PAIRS)
for (p, h), x, y in zip(CHECK_PAIRS, ours, ref):
print(f"{h[:50]:52s} sglang {np.round(x, 4)} hf {np.round(y, 4)}")
print("max abs diff:", float(np.abs(ours - ref).max()), "| argmax agree:", bool((ours.argmax(1) == ref.argmax(1)).all()))
else:
p = client.predict([(a.premise, a.hypothesis)], [a.image] if a.image else None)[0]
print({l: round(float(v), 4) for l, v in zip(LABELS, p)})
|