Buckets:
| """E3 step 1: short reasoning notes for 10,000 training decisions (run in the vLLM venv). | |
| The generator sees the gold option (privileged information) and is told not to reveal it. Notes are never | |
| cut: a note that reaches the token cap is flagged (`capped`) and rejected later. Leak filter: answer | |
| phrases, a bullet that is just the gold label, or the gold label's name when the input never mentions it. | |
| """ | |
| import argparse | |
| import json | |
| import random | |
| import re | |
| import sys | |
| from collections import Counter, defaultdict | |
| sys.path.insert(0, "/workspace/ext/jev-dllm") | |
| PROMPT = """You are given a situation, a question and the options. | |
| The correct option is: {gold}. Do NOT state or hint which option is correct. | |
| Write 3-5 short, independent bullet points (max 80 tokens total) listing only the facts and rules from the situation that matter for deciding. | |
| Facts and rules only: no conclusion, no final verdict, and never name or paraphrase any option. | |
| Situation: | |
| {state} | |
| Question: | |
| {question} | |
| Options: | |
| {options} | |
| Write only the bullet points, each starting with "- ".""" | |
| PROMPT_COT = """You are given a situation, a question and the options. | |
| The correct option is: {gold}. | |
| Work out the decision in 3-6 short bullet points (under 90 tokens in total): the intermediate facts and computations that lead to it. The last bullet states the conclusion. | |
| Situation: | |
| {state} | |
| Question: | |
| {question} | |
| Options: | |
| {options} | |
| Write only the bullet points, each starting with "- ".""" | |
| LEAK = re.compile(r"\b(the answer|correct (option|answer|choice|label)|the right (option|answer)|answer is|" | |
| r"should be (labeled|classified|chosen)|is the best (option|choice|answer)|therefore,? (the|it|this))\b", | |
| re.I) | |
| def serialize(state): | |
| return state if isinstance(state, str) else json.dumps(state, ensure_ascii=False, indent=1) | |
| def option_texts(q): | |
| t, crit = q["type"], q.get("criteria") | |
| if t == "noul": | |
| c = crit or {} | |
| return ["No" + (" - " + c["false"] if isinstance(c, dict) and c.get("false") else ""), | |
| "Yes" + (" - " + c["true"] if isinstance(c, dict) and c.get("true") else "")], ["No", "Yes"] | |
| if t == "score": | |
| names = [str(x) for x in crit] | |
| return ["%d: %s" % (i, x) for i, x in enumerate(names)], names | |
| names = list(crit) | |
| return [k + (" - " + v if isinstance(v, str) and v else "") for k, v in crit.items()], names | |
| def leak(note, gold_name, context): | |
| if LEAK.search(note): | |
| return "phrase" | |
| g = gold_name.strip().lower() | |
| for line in note.splitlines(): | |
| if line.strip("-*• .:").strip().lower() == g: | |
| return "standalone_label" | |
| short = g.split(":")[0].strip() # score/choice levels: "Released: collection ..." -> "released" | |
| for name in {g, short}: | |
| if len(name) < 3 or name in ("yes", "no"): | |
| continue | |
| # single words: stem = first max(5, len-4) chars ("entailment"->"entail", "released"->"relea") | |
| stem = name if " " in name else name[:max(5, len(name) - 4)] | |
| pat = r"\b%s" % re.escape(stem) | |
| if re.search(pat, note.lower()) and not re.search(pat, context.lower()): | |
| return "gold_label_named" | |
| return None | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--train", default="/workspace/data/training-legacy/s1/canonical/train.jsonl") | |
| ap.add_argument("--model", default="/root/hf/qwen3.5-9b") | |
| ap.add_argument("--out", default="/workspace/think-decide/runs/e3/notes.jsonl") | |
| ap.add_argument("--n", type=int, default=10000) | |
| ap.add_argument("--max-tokens", type=int, default=256) | |
| ap.add_argument("--gpu-mem", type=float, default=0.45) | |
| ap.add_argument("--limit", type=int, default=0) | |
| ap.add_argument("--first", type=int, default=0, help="use the first N rows of --train (no stratified pick)") | |
| ap.add_argument("--style", choices=["facts", "cot"], default="facts") | |
| a = ap.parse_args() | |
| rows = [json.loads(l) for l in open(a.train)] | |
| by = defaultdict(list) | |
| for i, r in enumerate(rows): | |
| by[r["source"]["name"]].append(i) | |
| rng = random.Random(42) | |
| pick = [] | |
| for s in sorted(by): # stratified: each source in proportion to its share of the 40k | |
| k = round(a.n * len(by[s]) / len(rows)) | |
| pick += rng.sample(by[s], k) | |
| pick = sorted(random.Random(7).sample(pick, a.limit)) if a.limit else sorted(pick) | |
| if a.first: | |
| pick = list(range(min(a.first, len(rows)))) | |
| print(json.dumps(dict(selected=len(pick), by_source=Counter(rows[i]["source"]["name"] for i in pick))), flush=True) | |
| from vllm import LLM, SamplingParams | |
| llm = LLM(a.model, dtype="bfloat16", gpu_memory_utilization=a.gpu_mem, max_model_len=8192, | |
| enable_prefix_caching=True, max_num_seqs=256) | |
| tok = llm.get_tokenizer() | |
| prompts, meta = [], [] | |
| for i in pick: | |
| r = rows[i] | |
| (qid, q), = r["questions"].items() | |
| texts, names = option_texts(q) | |
| probs = r["gold"][qid]["probabilities"] | |
| keys = list(q["criteria"]) if q["type"] == "choice" else ( | |
| [str(j) for j in range(len(q["criteria"]))] if q["type"] == "score" else ["false", "true"]) | |
| gold = max(range(len(keys)), key=lambda j: probs[keys[j]]) | |
| state = serialize(r["state"]) | |
| msg = (PROMPT_COT if a.style == "cot" else PROMPT).format(gold=texts[gold], state=state, question=q["instructions"], options="\n".join(texts)) | |
| prompts.append(tok.apply_chat_template([{"role": "user", "content": msg}], tokenize=False, | |
| add_generation_prompt=True, enable_thinking=False)) | |
| meta.append(dict(case_id=r["case_id"], view_id=r["view_id"], source=r["source"]["name"], | |
| primitive=q["type"], gold_index=gold, gold_name=names[gold], | |
| context=state + "\n" + q["instructions"])) | |
| sp = SamplingParams(temperature=0.0, max_tokens=a.max_tokens) | |
| outs = llm.generate(prompts, sp) | |
| stats = Counter() | |
| with open(a.out, "w") as f: | |
| for m, o in zip(meta, outs): | |
| c = o.outputs[0] | |
| note = c.text.strip() | |
| capped = c.finish_reason == "length" | |
| ctx = m.pop("context") | |
| why = "capped" if capped else (None if a.style == "cot" else leak(note, m["gold_name"], ctx)) | |
| stats[why or "ok"] += 1 | |
| f.write(json.dumps(dict(m, note=note, n_tokens=len(c.token_ids), finish=c.finish_reason, | |
| reject=why)) + "\n") | |
| print(json.dumps(dict(event="done", out=a.out, stats=stats)), flush=True) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 6.64 kB
- Xet hash:
- 2b5923f66ac4f1e317488d45896477bf916491afb5884c163fd9aaae72aa13f5
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.