| """Select a reproducible, domain-balanced LongBench v2 diagnostic subset.""" |
| import hashlib |
| import json |
| import random |
| from pathlib import Path |
| from huggingface_hub import hf_hub_download |
| from transformers import AutoTokenizer |
|
|
| ROOT=Path(__file__).resolve().parent |
| REV='2b48e494f2c7a2f0af81aae178e05c7e1dde0fe9' |
| SOURCE_REV='87420286149d9cce9bd46cd335ef9bda33c37c1b' |
|
|
| def prompt_for(row): |
| return ('Read the context and answer the question.\n\nContext:\n'+row['context']+ |
| '\n\nQuestion: '+row['question']+'\n'+ |
| '\n'.join(letter+'. '+row['choice_'+letter] for letter in 'ABCD')+ |
| '\nPut your final answer letter inside \\boxed{}.') |
|
|
| def main(): |
| raw=json.loads(Path(hf_hub_download('zai-org/LongBench-v2','data.json',repo_type='dataset',revision=REV)).read_text()) |
| tokenizer=AutoTokenizer.from_pretrained('nex-agi/Nex-N2.5-mini',revision=SOURCE_REV) |
| rng=random.Random(20260909) |
| chosen=[] |
| for domain in sorted({r['domain'] for r in raw}): |
| candidates=[r for r in raw if r['domain']==domain and len(r['context'])<900000] |
| rng.shuffle(candidates) |
| bins=[[],[]] |
| for row in candidates: |
| prompt=prompt_for(row) |
| text=tokenizer.apply_chat_template([dict(role='user',content=prompt)],tokenize=False, |
| add_generation_prompt=True,reasoning_effort='high') |
| count=len(tokenizer(text,add_special_tokens=False).input_ids) |
| if 8192 <= count <= 245000: |
| bins[int(count>64000)].append(dict(id=row['_id'],domain=domain,prompt_tokens=count, |
| prompt_sha256=hashlib.sha256(prompt.encode()).hexdigest())) |
| if all(len(b)>=2 for b in bins): |
| break |
| selected=bins[0][:2]+bins[1][:2] |
| selected_ids={r['id'] for r in selected} |
| selected += [r for r in bins[0]+bins[1] if r['id'] not in selected_ids][:4-len(selected)] |
| assert len(selected)==4,(domain,len(selected)) |
| chosen.extend(selected) |
| print(domain,[r['prompt_tokens'] for r in selected],flush=True) |
| result=dict(dataset='zai-org/LongBench-v2',revision=REV,source_revision=SOURCE_REV,seed=20260909, |
| method='Four cases per domain; two <=64K and two >64K tokens where available; 8192..245000 token limit, context <900000 characters, no truncation.',records=chosen) |
| (ROOT/'longbench_selection.json').write_text(json.dumps(result,indent=2)) |
|
|
| if __name__=='__main__':main() |
|
|
