Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """Build the BrowseComp-Plus structures bundle from an AML response JSONL. | |
| Reads the per-question structure-generation responses (from the | |
| ``information-scaffolds`` derive-structure pipeline) and writes one | |
| combined JSONL into ``backend/data/structures.jsonl``. The Flask backend | |
| loads this file at startup and serves it via /api/browsecomp/structures/. | |
| Single-file output (not per-qid shards) keeps the BrowseComp Space: | |
| * inside the LFS rule already in .gitattributes (``backend/data/*.jsonl``), | |
| * cheap to ship inside the Docker image (one COPY, one file open). | |
| Usage: | |
| python scripts/build_structures.py /path/to/response.jsonl | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import re | |
| import sys | |
| DATASET = "browsecomp_plus" | |
| DEFAULT_OUT = os.path.join( | |
| os.path.dirname(os.path.dirname(os.path.abspath(__file__))), | |
| "backend", | |
| "data", | |
| "structures.jsonl", | |
| ) | |
| def main() -> int: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("response", help="Path to the AML response JSONL file") | |
| ap.add_argument( | |
| "--out", | |
| default=DEFAULT_OUT, | |
| help="Output JSONL path (default: backend/data/structures.jsonl)", | |
| ) | |
| ap.add_argument( | |
| "--include-empty", | |
| action="store_true", | |
| help="Keep rows whose answer is empty/whitespace (default: skip — these are " | |
| "model timeouts or content-filter rejections with no useful structure).", | |
| ) | |
| args = ap.parse_args() | |
| records = [] | |
| skipped_empty = 0 | |
| with open(args.response) as f: | |
| for line in f: | |
| line = line.strip() | |
| if not line: | |
| continue | |
| obj = json.loads(line) | |
| if obj.get("dataset") != DATASET: | |
| continue | |
| answer = obj.get("answer", "") or "" | |
| if not args.include_empty and not answer.strip(): | |
| skipped_empty += 1 | |
| continue | |
| prompt_doc_ids = re.findall(r"\[Doc \d+\] id=(\S+)", obj.get("user_prompt", "")) | |
| records.append({ | |
| "qid": str(obj["qid"]), | |
| "question": obj["question"], | |
| "model": obj.get("model"), | |
| "reasoning_effort": obj.get("reasoning_effort"), | |
| "max_completion_tokens": obj.get("max_completion_tokens"), | |
| "available_num_docs": obj.get("available_num_docs"), | |
| "prompt_num_docs": obj.get("prompt_num_docs"), | |
| "system_prompt_file": obj.get("system_prompt_file"), | |
| "prompt_doc_ids": prompt_doc_ids, | |
| "usage": obj.get("usage"), | |
| "latency_ms": obj.get("latency_ms"), | |
| "finish_reason": obj.get("finish_reason"), | |
| "answer": answer, | |
| }) | |
| records.sort(key=lambda r: r["qid"]) | |
| os.makedirs(os.path.dirname(args.out), exist_ok=True) | |
| with open(args.out, "w") as f: | |
| for r in records: | |
| f.write(json.dumps(r, ensure_ascii=False) + "\n") | |
| print( | |
| f"wrote {len(records)} {DATASET} structure record(s) to {args.out}" | |
| + (f" (skipped {skipped_empty} empty-answer row(s))" if skipped_empty else ""), | |
| file=sys.stderr, | |
| ) | |
| return 0 | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |