File size: 3,705 Bytes
867d0f3 | 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 | """
Second pass over segments whose greedy translation failed validation.
Two recovery strategies, tried in order, keeping the first candidate that validates:
1. beam search (beam_width=4) — what the Seed-X authors recommend
2. sampled best-of-8 — a different part of the distribution when beam search repeats the error
Recovered translations are appended to the cache, overriding the greedy result.
"""
import json
import os
import sys
from pathlib import Path
sys.path.insert(0, ".")
from build_dataset import check
from gsm_common import SRC_PROMPT
OUT = Path("out_gsm")
CACHE = OUT / "translations.jsonl"
MODEL = "./models/Seed-X-PPO-7B"
def load():
trans = {}
with open(CACHE, encoding="utf-8") as fh:
for line in fh:
try:
r = json.loads(line)
except json.JSONDecodeError:
continue
trans[r["src"]] = r["tgt"]
rows = [json.loads(l) for l in open(OUT / "selected_rows.jsonl", encoding="utf-8")]
kind = {}
for r in rows:
kind.setdefault(r["question"], "question")
kind.setdefault(r["thinking"], "thinking")
return trans, kind
def main():
trans, kind = load()
failed = [
s for s, t in trans.items()
if check(s, t, strict=(kind.get(s) == "thinking")) is not None
]
print(f"[*] {len(failed)}/{len(trans)} segments failed validation ({len(failed)/len(trans):.2%})")
if not failed:
return
from vllm import LLM, SamplingParams
from vllm.sampling_params import BeamSearchParams
llm = LLM(model=MODEL, max_num_seqs=256, gpu_memory_utilization=0.92, max_model_len=1024)
prompts = [SRC_PROMPT.format(text=s) for s in failed]
recovered = {}
# Beam search is disabled by default: vLLM 0.8.5 runs it as a Python-level loop and it took
# >50 min on ~4.8k segments without finishing, versus ~2 min for the batched sampling path
# below. Set RETRY_BEAM=1 to use it anyway.
if os.environ.get("RETRY_BEAM") == "1":
print("[*] pass 1: beam search")
outs = llm.beam_search([{"prompt": p} for p in prompts], BeamSearchParams(beam_width=4, max_tokens=256))
still = []
for src, o in zip(failed, outs):
strict = kind.get(src) == "thinking"
for seq in o.sequences:
cand = seq.text.strip()
if check(src, cand, strict=strict) is None:
recovered[src] = cand
break
else:
still.append(src)
print(f" recovered {len(recovered)}, still failing {len(still)}")
else:
print("[*] pass 1: skipped (beam search disabled)")
still = list(failed)
if still:
print("[*] pass 2: sampled best-of-8")
params = SamplingParams(n=8, temperature=0.8, top_p=0.95, max_tokens=256, skip_special_tokens=True)
outs = llm.generate([SRC_PROMPT.format(text=s) for s in still], params)
final = []
for src, o in zip(still, outs):
strict = kind.get(src) == "thinking"
for cand in o.outputs:
text = cand.text.strip()
if check(src, text, strict=strict) is None:
recovered[src] = text
break
else:
final.append(src)
print(f" recovered {len(recovered)} total, unrecoverable {len(final)}")
with open(CACHE, "a", encoding="utf-8") as fh:
for src, tgt in recovered.items():
fh.write(json.dumps({"src": src, "tgt": tgt}, ensure_ascii=False) + "\n")
print(f"[+] appended {len(recovered)} recovered translations to {CACHE}")
if __name__ == "__main__":
main()
|