#!/usr/bin/env python3 """Convert the PlantGeneBench 100-question Markdown table to JSONL tasks/gold.""" from __future__ import annotations import argparse import json import re from pathlib import Path from typing import Any LAYER_CATEGORY = { "TU": "task_understanding", "EP": "execution_planning", "EC": "evidence_constraint", "BI": "biological_interpretation", "BC": "boundary_control", } ANALYSIS_TYPE_BY_STAGE = { "rnaseq_qc": "rnaseq", "rnaseq_align": "rnaseq", "rnaseq_quant": "rnaseq", "rnaseq_sample_qc": "rnaseq", "rnaseq_de": "rnaseq", "rnaseq_enrich": "rnaseq", "rnaseq_wgcna": "rnaseq", "rnaseq_fgsea": "rnaseq", "rnaseq_candidate_family": "rnaseq", } def split_md_row(line: str) -> list[str]: return [cell.strip() for cell in line.strip().strip("|").split("|")] def parse_list_expr(text: str, key: str) -> list[str]: match = re.search(rf"{re.escape(key)}=\[([^\]]*)\]", text) if not match: return [] value = match.group(1) return [ item.strip().strip("'\"` ") for item in re.split(r",|，", value) if item.strip().strip("'\"` ") ] def parse_target_stages(text: str) -> list[str]: stages = parse_list_expr(text, "target_stages") if stages: return stages stages = parse_list_expr(text, "target") if stages: return stages # A few rows describe executable stages without literal target_stages. if "map to " in text: mapped: list[str] = [] for stage in [ "identify", "motif", "domain", "phylogeny", "synteny", "interspecies_synteny", "promoter", "pep", "kaks", ]: if re.search(rf"\b{stage}\b", text): mapped.append(stage) return mapped return [] def parse_execution_plan(text: str, target_stages: list[str]) -> list[str]: stages = parse_list_expr(text, "execution_plan") if stages: return stages arrow = re.search(r"([a-zA-Z0-9_]+(?:\s*->\s*[a-zA-Z0-9_]+)+)", text) if arrow: return [part.strip() for part in arrow.group(1).split("->")] return target_stages def parse_params(text: str) -> dict[str, Any]: params: dict[str, Any] = {} patterns = { "promoter_length": r"promoter_length=([0-9]+)", "max_motifs": r"max_motifs=([0-9]+)", "bootstrap": r"bootstrap=([0-9]+)", "geneFamily": r"family=([A-Za-z0-9\-]+)", "species": r"species=([A-Za-z_]+)", } for key, pattern in patterns.items(): match = re.search(pattern, text) if match: value: Any = match.group(1) if key in {"promoter_length", "max_motifs", "bootstrap"}: value = int(value) params[key] = value if "tree_method=ML" in text: params["tree_method"] = "ML" if "analysis_input=DEG_result" in text: params["analysis_input"] = "DEG_result" return params def infer_analysis_type(stages: list[str], category: str) -> str: if category in {"evidence_constraint", "biological_interpretation", "boundary_control"}: return category if any(stage in ANALYSIS_TYPE_BY_STAGE for stage in stages): return "rnaseq" return "gene_family" def forbidden_stages(must_not: str) -> list[str]: found: list[str] = [] stage_terms = { "identify": ["identify"], "phylogeny": ["phylogeny"], "motif": ["motif"], "domain": ["domain"], "structure": ["structure"], "promoter": ["promoter"], "synteny": ["synteny"], "interspecies_synteny": ["interspecies_synteny"], "karyotype": ["karyotype"], "combo": ["combo"], "ppi": ["PPI", "ppi"], "expression": ["expression"], "pep": ["pep"], "rnaseq_qc": ["rnaseq_qc"], "rnaseq_align": ["rnaseq_align"], "rnaseq_quant": ["rnaseq_quant"], "rnaseq_de": ["rnaseq_de"], "rnaseq_enrich": ["rnaseq_enrich"], } action_patterns = [ "不要安排", "不要运行", "不要重新跑", "不要加入", "不要调用", "不要转到", "不要启动", "不要强制", "不要直接", "不要只运行", "不要创建", ] for stage, terms in stage_terms.items(): for term in terms: if any(f"{action}{term}" in must_not or f"{action} {term}" in must_not for action in action_patterns): found.append(stage) if "不要安排 expression、motif 或 RNA-seq" in must_not or "不要运行 raw RNA-seq" in must_not: found.extend(["rnaseq_qc", "rnaseq_align", "rnaseq_quant", "rnaseq_de", "rnaseq_enrich"]) return sorted(set(found)) def convert(md_path: Path) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: tasks: list[dict[str, Any]] = [] gold: list[dict[str, Any]] = [] for line in md_path.read_text(encoding="utf-8").splitlines(): if not line.startswith("| PGB-"): continue cells = split_md_row(line) if len(cells) < 5: continue task_id, difficulty, query, expected_text, must_not = cells[:5] layer_code = task_id.split("-")[1] category = LAYER_CATEGORY.get(layer_code, "unknown") target_stages = parse_target_stages(expected_text) execution_plan = parse_execution_plan(expected_text, target_stages) params = parse_params(expected_text) analysis_type = infer_analysis_type(target_stages or execution_plan, category) task = { "task_id": task_id, "category": category, "difficulty": difficulty, "source_type": "synthetic_expert_curated", "tags": [category], "user_query": query, "input_context": { "job_id": "", "prefix": "", "intent": { "analysis_type": analysis_type, "stages": target_stages, "params": params, }, }, "gold_summary": expected_text, "must_not_do": must_not, } tasks.append(task) gold.append( { "task_id": task_id, "gold": { "analysis_type": analysis_type, "expected_stages": target_stages, "expected_execution_plan": execution_plan, "required_params": params, "required_dependencies": {}, "forbidden_stages": forbidden_stages(must_not), "can_execute": "depends_on_inputs", "rationale": expected_text, "must_not_do": must_not, }, } ) return tasks, gold def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", encoding="utf-8") as handle: for row in rows: handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n") def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--md", type=Path, required=True) parser.add_argument("--tasks-out", type=Path, required=True) parser.add_argument("--gold-out", type=Path, required=True) args = parser.parse_args() tasks, gold = convert(args.md) write_jsonl(args.tasks_out, tasks) write_jsonl(args.gold_out, gold) print(json.dumps({"tasks": len(tasks), "gold": len(gold)}, ensure_ascii=False)) return 0 if __name__ == "__main__": raise SystemExit(main())