day26_code/alphacodium_flow.py

"""alphacodium_flow.py — harness track: AlphaCodium-style flow (Ridnik 2024, arXiv:2401.08500) measured
on the SAME sealed held-out as greedy@1, so flow@1 vs greedy@1 is apples-to-apples.

The literature's most RELIABLE held-out gain for code without moving weights (AlphaCodium: 19%->44% pass@5).
It compounds capability via an agentic flow, sidestepping the weight-transfer wall (SFT/GRPO/DPO/ReST-EM
all failed to lift held-out here). NO ground-truth test leak: the flow uses the model's OWN generated
tests + execution self-debugging to pick/repair a solution; the final answer is graded ONCE on the real
sealed suite (same verdict as greedy@1).

Flow per problem (one final solution returned -> flow@1):
  1. PLAN     : model analyzes the problem, lists approach + edge cases (reflection).
  2. GEN-TESTS: model writes a few of its OWN (input -> expected stdout) cases from the spec (no oracle).
  3. CODE     : initial solution conditioned on plan.
  4. REPAIR   : run candidate on the model's own tests; on mismatch, feed the failing case + actual output
                back and ask for a fix; up to R rounds (execution-feedback self-debug).
  5. FINAL    : the surviving solution -> graded on the real sealed suite (the only ground-truth use).
Baseline greedy@1 recomputed here for a matched comparison. Deterministic eval. Qwen2.5-72B + best-of-N
optional (FLOW_K samples of the whole flow, keep first that passes the model's own tests)."""
import os, sys, json, subprocess, tempfile, re, hashlib, signal
sys.path.insert(0,"/workspace/RSI")
from concurrent.futures import ThreadPoolExecutor
from src.utils.external_benchmarks import _try_load_from_datasets, _extract_code
from src.utils.vllm_backend import VLLMModelLoader

MODEL=os.environ.get("MODEL","/workspace/RSI/expanded_models/qwen72")
BNB={"load_in_4bit":True,"bnb_4bit_compute_dtype":"bfloat16"}
PREREG="/workspace/RSI/outputs/prereg.json"; SET=os.environ.get("FLOW_SET","hard_holdout")
OUT=os.environ.get("FLOW_OUT","/workspace/RSI/outputs/alphacodium_flow.jsonl")
R=int(os.environ.get("FLOW_REPAIR","3")); MAXTOK=int(os.environ.get("MAXTOK","900"))
EVAL_TC=int(os.environ.get("EVAL_TC","40")); EVAL_TO=int(os.environ.get("EVAL_TO","8"))
NGEN_TESTS=int(os.environ.get("NGEN_TESTS","3")); _GW=int(os.environ.get("GRADE_WORKERS","48"))
_POOL=ThreadPoolExecutor(max_workers=_GW)

def _norm(s): return "\n".join(l.rstrip() for l in str(s).strip().splitlines())
def _thash(it): return hashlib.md5((it.prompt+repr(it.meta.get("inputs"))+repr(it.meta.get("outputs"))).encode()).hexdigest()
def run_stdin(code,inp,t):
    d=tempfile.mkdtemp(); pth=os.path.join(d,"s.py"); open(pth,"w").write(code)
    try:
        p=subprocess.Popen([sys.executable,pth],stdin=subprocess.PIPE,stdout=subprocess.PIPE,
                           stderr=subprocess.PIPE,text=True,start_new_session=True)
        try: out,err=p.communicate(input=str(inp),timeout=t); return out,err
        except subprocess.TimeoutExpired:
            try: os.killpg(os.getpgid(p.pid),signal.SIGKILL)
            except Exception: pass
            try: p.communicate(timeout=2)
            except Exception: pass
            return "","TIMEOUT"
        except Exception: return "",""
    except Exception: return "",""
    finally:
        try: os.remove(pth); os.rmdir(d)
        except Exception: pass
def grade(code,item,maxtc=None,to=EVAL_TO):
    ins=item.meta.get("inputs") or []; outs=item.meta.get("outputs") or []
    if not code or not ins: return 0.0
    tc=list(zip(ins,outs))
    if maxtc: step=max(1,len(tc)//maxtc); tc=tc[::step][:maxtc]
    res=list(_POOL.map(lambda io:_norm(run_stdin(code,io[0],to)[0])==_norm(io[1]), tc))
    return sum(res)/max(1,len(res))

def main():
    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=True,enforce_eager=True)
    loader.load()
    def gen(ps,temp=0.0): return list(loader.generate_batch(ps,max_new_tokens=MAXTOK,temperature=temp,top_p=(1.0 if temp==0 else 0.95)))
    PR=json.load(open(PREREG)); IDS=set(PR[SET])
    apps=[it for it in (_try_load_from_datasets("apps") or []) if not (it.meta.get("fn_name") or "").strip() and it.meta.get("inputs") and it.meta.get("outputs")]
    items=[it for it in apps if _thash(it) in IDS]
    print(f"[flow] {SET}: {len(items)} problems | repair rounds={R} | self-tests={NGEN_TESTS}",flush=True)

    # matched greedy@1 baseline
    g=gen([it.prompt for it in items],0.0)
    greedy=[grade(_extract_code(o) or o,it,EVAL_TC)>=0.999 for it,o in zip(items,g)]
    print(f"[flow] greedy@1 = {sum(greedy)}/{len(items)} ({sum(greedy)/len(items):.4f})",flush=True)

    def parse_tests(txt):
        # expect blocks: ###INPUT\n...\n###OUTPUT\n...  (model-generated, NO oracle)
        tests=[]
        for m in re.finditer(r"###INPUT\s*\n(.*?)\n###OUTPUT\s*\n(.*?)(?=\n###INPUT|\Z)", txt, re.S):
            tests.append((m.group(1).strip("\n"), m.group(2).strip("\n")))
        return tests[:NGEN_TESTS]

    flow_pass=[]
    for idx,it in enumerate(items):
        # 1. PLAN
        plan=gen([it.prompt+"\n\nBriefly: outline the algorithm and list tricky edge cases. <=120 words."],0.0)[0]
        # 2. GEN model's own tests (no oracle)
        tt=gen([it.prompt+"\n\nWrite "+str(NGEN_TESTS)+" small test cases you are confident about, each EXACTLY as:\n###INPUT\n<stdin>\n###OUTPUT\n<stdout>"],0.0)[0]
        mytests=parse_tests(tt)
        # 3. CODE (conditioned on plan)
        code=_extract_code(gen([it.prompt+"\n\nApproach:\n"+plan+"\n\nWrite the complete Python solution (reads stdin, writes stdout). Output ONLY a ```python block."],0.0)[0]) or ""
        # 4. REPAIR against the model's OWN tests using execution feedback
        for r in range(R):
            if not code.strip() or not mytests: break
            fail=None
            for tin,tout in mytests:
                got,err=run_stdin(code,tin,EVAL_TO)
                if _norm(got)!=_norm(tout):
                    fail=(tin,tout,_norm(got) if got else ("ERROR: "+err[:200])); break
            if fail is None: break  # passes own tests
            fixp=(it.prompt+"\n\nYour solution:\n```python\n"+code+"\n```\nOn input:\n"+fail[0]+"\nexpected:\n"+fail[1]+"\nbut got:\n"+fail[2]+"\nFix the bug. Output ONLY a corrected ```python block.")
            code=_extract_code(gen([fixp],0.0)[0]) or code
        # 5. FINAL grade on the real sealed suite (only ground-truth use)
        ok=grade(code,it,EVAL_TC)>=0.999; flow_pass.append(ok)
        if (idx+1)%10==0: print(f"[flow] ...{idx+1}/{len(items)} done, flow solved so far {sum(flow_pass)}",flush=True)
    gp=sum(greedy); fp=sum(flow_pass)
    rec={"set":SET,"n":len(items),"greedy@1":gp,"flow@1":fp,"greedy_rate":round(gp/len(items),4),"flow_rate":round(fp/len(items),4),"repair":R}
    open(OUT,"a").write(json.dumps(rec)+"\n")
    print(f"[flow] RESULT greedy@1={gp}/{len(items)} -> flow@1={fp}/{len(items)} (+{fp-gp} = +{(fp-gp)/len(items):.1%} absolute, no weights)",flush=True)
    print("[flow] DONE",flush=True); sys.exit(0)

if __name__=="__main__":
    main()