239 lines
7.6 KiB
Python
Executable File
239 lines
7.6 KiB
Python
Executable File
#!/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())
|