File size: 4,534 Bytes
77f690c | 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 | """serve_bon.py — the shippable "smarter version": chat with base + best-of-N + real-exec verifier.
Day-25 proved best-of-N+verifier lifts sealed-holdout solve 11x (greedy 2/60 -> bo32 22/60) with NO weight
changes. This exposes that as a usable thing: an interactive REPL (and a one-shot --prompt mode) that, for a
coding task with stdin/stdout test cases, samples K candidates, runs each against the provided tests, and
returns the FIRST that passes (or the best-partial if none fully pass). For free-form chat (no tests) it
returns the greedy answer. This is base-model + inference search — the "smarter version" used by default.
Usage (pod):
PYTHONPATH=/workspace/RSI VLLM_USE_FLASHINFER_SAMPLER=0 /venv/main/bin/python3 serve_bon.py # REPL
... serve_bon.py --prompt "problem text" --tests tests.json # one-shot
tests.json = {"inputs": ["..."], "outputs": ["..."]} (same shape as apps meta). Omit for plain chat.
Env: MODEL, BON_K (default 16), BON_TEMP (0.9), MAXTOK (640), ADAPTER (optional LoRA dir)."""
import os, sys, json, argparse, subprocess, tempfile
sys.path.insert(0,"/workspace/RSI")
from concurrent.futures import ThreadPoolExecutor
from src.utils.external_benchmarks import _extract_code
from src.utils.vllm_backend import VLLMModelLoader
MODEL=os.environ.get("MODEL","/workspace/RSI/expanded_models/qwen3_32b")
BNB={"load_in_4bit":True,"bnb_4bit_compute_dtype":"bfloat16"}
K=int(os.environ.get("BON_K","16")); TEMP=float(os.environ.get("BON_TEMP","0.9"))
MAXTOK=int(os.environ.get("MAXTOK","640")); ADAPTER=os.environ.get("ADAPTER","")
_POOL=ThreadPoolExecutor(max_workers=int(os.environ.get("GRADE_WORKERS","32")))
def _norm(s): return "\n".join(l.rstrip() for l in str(s).strip().splitlines())
def run_stdin(code,inp,t=8):
d=tempfile.mkdtemp(); p=os.path.join(d,"s.py"); open(p,"w").write(code)
try:
try: return subprocess.run([sys.executable,p],input=str(inp),capture_output=True,text=True,timeout=t).stdout
except Exception: return ""
finally:
try: os.remove(p); os.rmdir(d)
except: pass
def score(code,tests):
ins=tests.get("inputs") or []; outs=tests.get("outputs") or []
if not code or not ins: return 0.0
tc=list(zip(ins,outs)); res=list(_POOL.map(lambda io:_norm(run_stdin(code,io[0]))==_norm(io[1]),tc))
return sum(res)/max(1,len(res))
def main():
ap=argparse.ArgumentParser(); ap.add_argument("--prompt"); ap.add_argument("--tests"); ap.add_argument("--k",type=int,default=K)
a=ap.parse_args()
loader=VLLMModelLoader(model_path=MODEL,dtype="bfloat16",max_model_len=4096,gpu_memory_utilization=0.85,
allow_remote_code=True,quantization_config=BNB,max_lora_rank=128,enable_chunked_prefill=False,
enable_lora=bool(ADAPTER),enforce_eager=True)
loader.load()
if ADAPTER and os.path.exists(ADAPTER): loader.set_lora_adapter(ADAPTER); print(f"[serve] adapter {ADAPTER}",flush=True)
def solve(prompt,tests,k):
if not tests: # plain chat -> greedy
return loader.generate_batch([prompt],max_new_tokens=MAXTOK,temperature=0.0,top_p=1.0)[0], None
# best-of-K: sample k, return first that fully passes; else best-partial
cands=list(loader.generate_batch([prompt]*k,max_new_tokens=MAXTOK,temperature=TEMP,top_p=0.95))
best=None; best_s=-1.0
for c in cands:
code=_extract_code(c) or c; s=score(code,tests)
if s>=0.999: return c, 1.0
if s>best_s: best,best_s=c,s
return best, best_s
if a.prompt is not None:
tests=json.load(open(a.tests)) if a.tests else None
ans,s=solve(a.prompt,tests,a.k); print(f"\n[serve] {'SOLVED' if s==1.0 else ('best-partial %.2f'%s if s is not None else 'chat')}\n{ans}")
return
print(f"[serve] REPL ready — base+best-of-{K}+verifier. Paste problem; for verified solve add a line '#TESTS <path.json>'. 'exit' quits.")
while True:
try: u=input("\nyou> ").strip()
except (EOFError,KeyboardInterrupt): break
if u=="exit": break
if not u: continue
tests=None
if "#TESTS" in u:
parts=u.split("#TESTS"); u=parts[0].strip()
tp=parts[1].strip()
if os.path.exists(tp): tests=json.load(open(tp))
ans,s=solve(u,tests,a.k)
tag=("SOLVED via best-of-%d"%a.k) if s==1.0 else ("best-partial %.2f"%s if s is not None else "chat")
print(f"\nmodel[{tag}]> {ans}")
if __name__=="__main__":
main()
|