File size: 6,368 Bytes
9f5fbb5 285bb31 9f5fbb5 285bb31 9f5fbb5 285bb31 9f5fbb5 285bb31 9f5fbb5 285bb31 9f5fbb5 285bb31 9f5fbb5 | 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 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 | #!/usr/bin/env python3
"""Token-shape shim between lm-eval and llama-server. Stdlib only.
llama-server and lm-eval disagree on three shapes; this translates and
reverse-proxies everything else untouched:
- GET /tokenizer_info (missing on llama-server): served from the server's
own /props (bos/eos strings, canonical chat template).
- POST /tokenize {"prompt"} -> upstream {"content"} -> {"tokens"}.
- POST /detokenize {"tokens"} -> upstream -> {"prompt"}.
- /v1/completions logprobs: server returns chat-style
content[{token, logprob, top_logprobs}]; harness expects legacy
{token_logprobs[], top_logprobs[]}. Translated field-for-field.
- /v1/chat/completions: system messages are folded into the following user
turn, but only for families without a system role (Gemma). See families.py.
- Sidecar: every completions request + raw response (incl. reasoning_content)
is appended to thinking-log.jsonl for later inspection. Scoring untouched.
Usage: python3 shim.py [upstream] [port]
"""
import json
import os
import sys
import time
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
UPSTREAM = sys.argv[1] if len(sys.argv) > 1 else "http://host.docker.internal:8080"
PORT = int(sys.argv[2]) if len(sys.argv) > 2 else 8081
THINK_LOG = os.path.join(os.path.dirname(os.path.abspath(__file__)), "thinking-log.jsonl")
def _up(method, path, body=None):
data = json.dumps(body).encode() if body is not None else None
req = urllib.request.Request(
UPSTREAM + path, data=data, method=method,
headers={"Content-Type": "application/json"},
)
with urllib.request.urlopen(req, timeout=300) as r:
return r.status, r.read()
def _log(path, req_body, resp_raw):
try:
resp = json.loads(resp_raw)
label = (req_body or {}).get("model", "model")
for ch in resp.get("choices", []):
msg = ch.get("message", {})
# server echoes its local file path as the model id: personal,
# meaningless to others. Log our run label instead.
if isinstance(msg, dict):
msg.pop("model", None)
if "model" in resp:
resp["model"] = label
with open(THINK_LOG, "a") as f:
f.write(json.dumps({
"ts": time.time(),
"path": path,
"request": req_body,
"response": resp,
}) + "\n")
except Exception:
pass
def _info():
try:
_, raw = _up("GET", "/props")
p = json.loads(raw)
eos = p.get("eos_token", "<eos>")
bos = p.get("bos_token", "<bos>")
tmpl = p.get("chat_template") or ""
except Exception:
eos, bos, tmpl = "<eos>", "<bos>", ""
return {"eos_token": eos, "bos_token": bos, "pad_token": eos,
"chat_template": tmpl}
def _family():
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import families
try:
_, raw = _up("GET", "/props")
tmpl = json.loads(raw).get("chat_template", "")
except Exception:
tmpl = ""
name = families.detect(tmpl, sys.argv[3] if len(sys.argv) > 3 else "")
sys.stderr.write(f"shim: family={name}\n")
return families.profile(name)
def _fix_logprobs(raw):
try:
obj = json.loads(raw)
except Exception:
return raw
for ch in obj.get("choices", []):
lp = (ch.get("logprobs") or {}).get("content")
if isinstance(lp, list):
ch["logprobs"] = {
"token_logprobs": [e.get("logprob", 0.0) for e in lp],
"top_logprobs": [
{t.get("token", ""): t.get("logprob", 0.0)
for t in (e.get("top_logprobs") or [])}
or {e.get("token", ""): e.get("logprob", 0.0)}
for e in lp
],
}
return json.dumps(obj).encode()
def _fold_system(body):
try:
msgs = body.get("messages")
if isinstance(msgs, list) and msgs and msgs[0].get("role") == "system":
rest = msgs[1:]
if rest and rest[0].get("role") == "user":
rest[0] = dict(rest[0], content=msgs[0].get("content", "")
+ "\n" + rest[0].get("content", ""))
return dict(body, messages=rest)
except Exception:
pass
return body
INFO = _info()
FAMILY = _family()
class H(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def log_message(self, *a):
pass
def _send(self, code, obj=None, raw=None):
body = raw if raw is not None else json.dumps(obj).encode()
self.send_response(code)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def _read_json(self):
try:
n = int(self.headers.get("Content-Length") or 0)
except ValueError:
n = 0
if not n:
return {}
return json.loads(self.rfile.read(n) or b"{}")
def do_GET(self):
if self.path == "/tokenizer_info":
return self._send(200, INFO)
code, raw = _up("GET", self.path)
self._send(code, None, raw)
def do_POST(self):
body = self._read_json()
if self.path == "/tokenize":
code, raw = _up("POST", "/tokenize",
{"content": body.get("prompt", "")})
toks = json.loads(raw).get("tokens", [])
return self._send(code, {"tokens": toks})
if self.path == "/detokenize":
code, raw = _up("POST", "/detokenize",
{"tokens": body.get("tokens", [])})
text = json.loads(raw).get("content", "")
return self._send(code, {"prompt": text})
if self.path == "/v1/chat/completions" and FAMILY["fold_system"]:
body = _fold_system(body)
try:
code, raw = _up("POST", self.path, body)
except Exception as e:
return self._send(502, {"error": f"upstream: {e}"})
if "completions" in self.path:
_log(self.path, body, raw)
self._send(code, None, _fix_logprobs(raw))
if __name__ == "__main__":
ThreadingHTTPServer(("127.0.0.1", PORT), H).serve_forever()
|