| |
| from __future__ import annotations |
|
|
| import argparse |
| import contextlib |
| import fcntl |
| import io |
| import json |
| import os |
| import random |
| import re |
| import subprocess |
| import sys |
| import tempfile |
| import textwrap |
| import types |
| from pathlib import Path |
| from typing import Any |
|
|
|
|
| PROJECT_ROOT = Path("/225040511/project/React+code_labbench") |
| LAB_BENCH_ROOT = Path("/225040511/project/LAB-Bench") |
| DEFAULT_OUTPUT_ROOT = PROJECT_ROOT / "results" |
| DEFAULT_DEV_SIZE = 45 |
| DEFAULT_TEST_SIZE = 315 |
| DEFAULT_SEED = 20260514 |
|
|
| ANSWER_RE = re.compile(r"\[ANSWER\]\s*([A-Z])\s*\[/ANSWER\]", re.IGNORECASE) |
| SOLUTION_RE = re.compile(r"<solution>\s*(.*?)\s*</solution>", re.IGNORECASE | re.DOTALL) |
| PY_BLOCK_RE = re.compile(r"```(?:python)?\s*(.*?)```", re.IGNORECASE | re.DOTALL) |
| EXEC_RE = re.compile(r"<execute(?:\s+type=[\"']python[\"'])?\s*>(.*?)</execute>", re.IGNORECASE | re.DOTALL) |
| LETTER_RE = re.compile(r"\b([A-Z])\b", re.IGNORECASE) |
|
|
| sys.path.insert(0, str(LAB_BENCH_ROOT)) |
|
|
|
|
| def install_labbench_import_stubs() -> None: |
| if "vertexai" not in sys.modules: |
| vertexai = types.ModuleType("vertexai") |
| vertexai.init = lambda *_args, **_kwargs: None |
| sys.modules["vertexai"] = vertexai |
| if "google.auth" not in sys.modules: |
| google = sys.modules.setdefault("google", types.ModuleType("google")) |
| auth = types.ModuleType("google.auth") |
| auth.default = lambda *_args, **_kwargs: (types.SimpleNamespace(refresh=lambda *_a, **_k: None, token=""), None) |
| transport = types.ModuleType("google.auth.transport") |
| requests = types.ModuleType("google.auth.transport.requests") |
| requests.Request = lambda *_args, **_kwargs: None |
| transport.requests = requests |
| auth.transport = transport |
| google.auth = auth |
| sys.modules["google.auth"] = auth |
| sys.modules["google.auth.transport"] = transport |
| sys.modules["google.auth.transport.requests"] = requests |
| if "chembench" not in sys.modules: |
| chembench = types.ModuleType("chembench") |
| sys.modules["chembench"] = chembench |
| constant = types.ModuleType("chembench.constant") |
| constant.COT_PROMPT = "Think step by step." |
| constant.MCQ_REGEX_TEMPLATE_1 = r"\[ANSWER\]\s*([A-Z])\s*\[/ANSWER\]" |
| sys.modules["chembench.constant"] = constant |
| prompter = types.ModuleType("chembench.prompter") |
| prompter.prepare_mcq_answer = lambda text, *_args, **_kwargs: text |
| sys.modules["chembench.prompter"] = prompter |
| utils = types.ModuleType("chembench.utils") |
| utils.create_multiple_choice_regex = lambda letters: r"\b(" + "|".join(letters) + r")\b" |
| utils.post_process_prompts = lambda text: text |
| utils.run_regex = lambda _regex, text, return_first=True: None |
| sys.modules["chembench.utils"] = utils |
|
|
|
|
| install_labbench_import_stubs() |
| import labbench |
|
|
|
|
| def load_dotenv_files(paths: list[Path]) -> None: |
| for path in paths: |
| if not path.exists(): |
| continue |
| for raw_line in path.read_text(encoding="utf-8", errors="replace").splitlines(): |
| line = raw_line.strip() |
| if not line or line.startswith("#") or "=" not in line: |
| continue |
| key, value = line.split("=", 1) |
| os.environ.setdefault(key.strip(), value.strip().strip('"').strip("'")) |
|
|
|
|
| def load_eval(eval_name: str) -> labbench.Evaluator: |
| return labbench.Evaluator(labbench.Eval(eval_name), debug=False, open_answer=False, use_hf=False) |
|
|
|
|
| def select_instances(args: argparse.Namespace) -> list[tuple[str, Any]]: |
| evaluator = load_eval(args.eval) |
| instances = list(evaluator.eval_set.instances) |
| rng = random.Random(f"{args.seed}:{args.eval}:question-set") |
| rng.shuffle(instances) |
| if args.debug: |
| selected = instances[: min(3, len(instances))] |
| elif args.split == "dev": |
| selected = instances[: min(args.dev_size, len(instances))] |
| elif args.split == "test": |
| start = min(args.dev_size, len(instances)) |
| selected = instances[start : min(start + args.test_size, len(instances))] |
| else: |
| selected = instances |
| if args.shard_count > 1: |
| total = len(selected) |
| chunk_size = (total + args.shard_count - 1) // args.shard_count |
| selected = selected[ |
| min(total, args.shard_index * chunk_size) : min(total, (args.shard_index + 1) * chunk_size) |
| ] |
| return selected |
|
|
|
|
| def load_completed_results(path: Path) -> tuple[set[str], set[str]]: |
| if not path.exists(): |
| return set(), set() |
| completed_ids: set[str] = set() |
| completed_questions: set[str] = set() |
| for raw_line in path.read_text(encoding="utf-8", errors="replace").splitlines(): |
| if not raw_line.strip(): |
| continue |
| try: |
| record = json.loads(raw_line) |
| except json.JSONDecodeError: |
| continue |
| task_id = str(record.get("id") or record.get("task_id") or "").strip() |
| if task_id: |
| completed_ids.add(task_id) |
| question = str(record.get("question") or "").strip() |
| if question: |
| completed_questions.add(question) |
| return completed_ids, completed_questions |
|
|
|
|
| def append_text_locked(path: Path, text: str) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| with path.open("a", encoding="utf-8") as handle: |
| fcntl.flock(handle.fileno(), fcntl.LOCK_EX) |
| handle.write(text) |
| handle.flush() |
| os.fsync(handle.fileno()) |
| fcntl.flock(handle.fileno(), fcntl.LOCK_UN) |
|
|
|
|
| def append_jsonl_locked(path: Path, payload: dict[str, Any]) -> None: |
| append_text_locked(path, json.dumps(payload, ensure_ascii=False, default=str) + "\n") |
|
|
|
|
| def parse_answer(text: str, n_choices: int) -> str: |
| valid = set("ABCDEFGHIJKLMNOPQRSTUVWXYZ"[:n_choices]) |
| if match := ANSWER_RE.search(text or ""): |
| letter = match.group(1).upper() |
| if letter in valid: |
| return letter |
| if match := SOLUTION_RE.search(text or ""): |
| return parse_answer(match.group(1), n_choices) |
| for match in LETTER_RE.finditer(text or ""): |
| letter = match.group(1).upper() |
| if letter in valid: |
| return letter |
| return "" |
|
|
|
|
| def build_prompt(input_obj: Any, eval_name: str) -> str: |
| choices = "\n".join(input_obj.choices) |
| return f""" |
| You are a biology reasoning agent with a Python scratchpad. |
| Answer this multiple-choice LAB-Bench question from {eval_name}. |
| |
| You may reason and write Python code for sequence analysis. If you need code, |
| return it in either <execute>...</execute> or a Python fenced block. |
| When finished, return exactly: <solution>[ANSWER]X[/ANSWER]</solution> |
| |
| Question: |
| {input_obj.question} |
| |
| Choices: |
| {choices} |
| """.strip() |
|
|
|
|
| def execute_python(code: str, timeout: int) -> str: |
| with tempfile.TemporaryDirectory(prefix="react_code_labbench_") as tmp: |
| script = Path(tmp) / "scratch.py" |
| script.write_text(code, encoding="utf-8") |
| try: |
| completed = subprocess.run( |
| [sys.executable, str(script)], |
| cwd=tmp, |
| text=True, |
| capture_output=True, |
| timeout=timeout, |
| ) |
| except subprocess.TimeoutExpired: |
| return "Execution timed out." |
| output = completed.stdout |
| if completed.stderr: |
| output += "\n[stderr]\n" + completed.stderr |
| return output[-12000:] |
|
|
|
|
| def resolve_client(args: argparse.Namespace): |
| try: |
| from openai import OpenAI |
| except ImportError as exc: |
| raise SystemExit("openai package is required for React+code runner.") from exc |
| api_key = ( |
| args.api_key |
| or os.getenv("REACT_CODE_API_KEY") |
| or os.getenv("DEEPSEEK_API_KEY") |
| or os.getenv("BIOMNI_CUSTOM_API_KEY") |
| or os.getenv("OPENAI_API_KEY") |
| ) |
| if not api_key: |
| raise SystemExit("Missing API key. Set DEEPSEEK_API_KEY, REACT_CODE_API_KEY, BIOMNI_CUSTOM_API_KEY, or OPENAI_API_KEY.") |
| return OpenAI(api_key=api_key, base_url=args.base_url), args.model |
|
|
|
|
| def react_code_answer(client: Any, model: str, prompt: str, max_iters: int, timeout: int) -> tuple[str, str, dict[str, int]]: |
| messages = [ |
| { |
| "role": "system", |
| "content": "Use concise reasoning. Use Python when helpful. Return a final <solution>[ANSWER]X[/ANSWER]</solution>.", |
| }, |
| {"role": "user", "content": prompt}, |
| ] |
| transcript = [] |
| usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "llm_call_count": 0} |
| for _ in range(max_iters): |
| response = client.chat.completions.create(model=model, temperature=0, messages=messages) |
| usage["llm_call_count"] += 1 |
| if getattr(response, "usage", None): |
| usage["prompt_tokens"] += int(getattr(response.usage, "prompt_tokens", 0) or 0) |
| usage["completion_tokens"] += int(getattr(response.usage, "completion_tokens", 0) or 0) |
| usage["total_tokens"] += int(getattr(response.usage, "total_tokens", 0) or 0) |
| content = response.choices[0].message.content or "" |
| transcript.append(content) |
| if ANSWER_RE.search(content) or SOLUTION_RE.search(content): |
| return content, "\n\n".join(transcript), usage |
| code = None |
| if match := EXEC_RE.search(content): |
| code = match.group(1).strip() |
| elif match := PY_BLOCK_RE.search(content): |
| code = match.group(1).strip() |
| if not code: |
| messages.append({"role": "assistant", "content": content}) |
| messages.append({"role": "user", "content": "Finish with <solution>[ANSWER]X[/ANSWER]</solution>."}) |
| continue |
| observation = execute_python(code, timeout=timeout) |
| transcript.append("[observation]\n" + observation) |
| messages.append({"role": "assistant", "content": content}) |
| messages.append({"role": "user", "content": "Python observation:\n" + observation}) |
| return "", "\n\n".join(transcript), usage |
|
|
|
|
| def run_one(args: argparse.Namespace, client: Any, model: str, subset: str, instance: Any) -> dict[str, Any]: |
| input_obj, target_output, _unsure = instance.get_input_output() |
| final, transcript, usage = react_code_answer( |
| client, |
| model, |
| build_prompt(input_obj, args.eval), |
| max_iters=args.max_iterations, |
| timeout=args.command_timeout, |
| ) |
| answer = parse_answer(final or transcript, len(input_obj.choices)) |
| record = { |
| "id": str(instance.id), |
| "task_id": str(instance.id), |
| "subset": subset, |
| "question": str(input_obj.question), |
| "answer": str(target_output), |
| "agent_answer": answer, |
| "method": "React+code", |
| "model": model, |
| **usage, |
| } |
| append_jsonl_locked(args.result_file, record) |
| append_text_locked( |
| args.reasoning_log, |
| "\n".join( |
| [ |
| "=" * 80, |
| f"eval: {args.eval}", |
| f"split: {args.split}", |
| f"id: {instance.id}", |
| f"answer: {target_output}", |
| f"agent_answer: {answer}", |
| "", |
| "[transcript]", |
| transcript, |
| "", |
| ] |
| ), |
| ) |
| return record |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description="Run LAB-Bench with a lightweight React+code baseline.") |
| parser.add_argument("--eval", choices=[member.value for member in labbench.Eval], default="CloningScenarios") |
| parser.add_argument("--split", choices=["dev", "test", "all"], default="test") |
| parser.add_argument("--dev-size", type=int, default=DEFAULT_DEV_SIZE) |
| parser.add_argument("--test-size", type=int, default=DEFAULT_TEST_SIZE) |
| parser.add_argument("--seed", type=int, default=DEFAULT_SEED) |
| parser.add_argument("--shard-index", type=int, default=0) |
| parser.add_argument("--shard-count", type=int, default=1) |
| parser.add_argument("--output-root", type=Path, default=DEFAULT_OUTPUT_ROOT) |
| parser.add_argument("--result-file", type=Path, default=None) |
| parser.add_argument("--reasoning-log", type=Path, default=None) |
| parser.add_argument("--skip-existing-results", action="store_true") |
| parser.add_argument("--debug", action="store_true") |
| parser.add_argument("--model", default=os.getenv("REACT_CODE_MODEL", os.getenv("DEEPSEEK_MODEL_NAME", "deepseek-chat"))) |
| parser.add_argument("--base-url", default=os.getenv("REACT_CODE_BASE_URL", os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1"))) |
| parser.add_argument("--api-key", default=None) |
| parser.add_argument("--max-iterations", type=int, default=int(os.getenv("REACT_CODE_MAX_ITERATIONS", "6"))) |
| parser.add_argument("--command-timeout", type=int, default=int(os.getenv("REACT_CODE_TIMEOUT", "60"))) |
| parser.add_argument("--env-file", action="append", type=Path, default=[]) |
| return parser.parse_args() |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| if args.shard_count < 1: |
| raise SystemExit("--shard-count must be at least 1.") |
| if args.shard_index < 0 or args.shard_index >= args.shard_count: |
| raise SystemExit("--shard-index must satisfy 0 <= shard-index < shard-count.") |
| load_dotenv_files([PROJECT_ROOT / ".env", LAB_BENCH_ROOT / ".env", Path("/225040511/project/.env"), *args.env_file]) |
| args.output_root.mkdir(parents=True, exist_ok=True) |
| eval_lower = args.eval.lower() |
| args.result_file = args.result_file or args.output_root / f"{eval_lower}_results.jsonl" |
| args.reasoning_log = args.reasoning_log or args.output_root / f"{eval_lower}_reasoning.log" |
| selected = select_instances(args) |
| if args.skip_existing_results: |
| completed_ids, completed_questions = load_completed_results(args.result_file) |
| selected = [ |
| (subset, instance) |
| for subset, instance in selected |
| if str(instance.id) not in completed_ids |
| and str(instance.get_input_output()[0].question).strip() not in completed_questions |
| ] |
| client, model = resolve_client(args) |
| for subset, instance in selected: |
| print(json.dumps(run_one(args, client, model, subset, instance), ensure_ascii=False), flush=True) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
