SkillOpt v0.1.0: initial release
- Skill optimization framework with training loop analogy - 11 benchmarks, 4 model backends (Azure OpenAI, Claude, Codex, Qwen) - WebUI for browser-based training control - Pluggable architecture for extending benchmarks and backends
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""ReflACT environment adapters."""
|
||||
@@ -0,0 +1,19 @@
|
||||
# Benchmark Template
|
||||
|
||||
This directory provides scaffold files for adding a new benchmark to SkillOpt.
|
||||
|
||||
## Files
|
||||
|
||||
- `env_template.py` — Environment adapter template
|
||||
- `loader_template.py` — Data loader template
|
||||
- `config_template.yaml` — Config file template
|
||||
|
||||
## Usage
|
||||
|
||||
1. Copy this directory: `cp -r skillopt/envs/_template skillopt/envs/your_benchmark`
|
||||
2. Rename files: remove `_template` suffix
|
||||
3. Implement the `TODO` sections
|
||||
4. Register in `skillopt/envs/__init__.py`
|
||||
5. Create config at `configs/your_benchmark/default.yaml`
|
||||
|
||||
See the [documentation](../../docs/guide/new-benchmark.md) for the full guide.
|
||||
@@ -0,0 +1,45 @@
|
||||
# ──────────────────────────────────────────────────
|
||||
# SkillOpt Config Template — <Your Benchmark Name>
|
||||
# ──────────────────────────────────────────────────
|
||||
# Copy this file to configs/<your_benchmark>/default.yaml
|
||||
# and customize the values below.
|
||||
|
||||
# Inherit global defaults
|
||||
_base_: ['../_base_/default.yaml']
|
||||
|
||||
# ── Environment ──────────────────────────────────
|
||||
env:
|
||||
name: your_benchmark # Must match registry key
|
||||
data_path: data/your_benchmark # Path to your data
|
||||
split_mode: ratio # "ratio" or "split_dir"
|
||||
split_ratio: "2:1:7" # train:val:test
|
||||
exec_timeout: 120 # Per-task timeout (seconds)
|
||||
|
||||
# ── Training ─────────────────────────────────────
|
||||
train:
|
||||
num_epochs: 4 # Number of epochs
|
||||
batch_size: 40 # Tasks per step (batch size)
|
||||
seed: 42
|
||||
|
||||
# ── Gradient (Reflection) ───────────────────────
|
||||
gradient:
|
||||
analyst_workers: 16 # Parallel reflection workers
|
||||
minibatch_size: 8
|
||||
|
||||
# ── Optimizer ────────────────────────────────────
|
||||
optimizer:
|
||||
learning_rate: 4 # Max edits per step (edit budget)
|
||||
lr_scheduler: cosine # cosine | linear | constant | autonomous
|
||||
use_slow_update: true # Epoch-boundary momentum
|
||||
use_meta_skill: true # Cross-epoch teacher memory
|
||||
|
||||
# ── Evaluation ───────────────────────────────────
|
||||
evaluation:
|
||||
use_gate: true # Validation gating
|
||||
eval_test: true # Run test eval after training
|
||||
|
||||
# ── Model ────────────────────────────────────────
|
||||
model:
|
||||
backend: azure_openai # azure_openai | openai_chat | claude_code_exec | qwen
|
||||
teacher: gpt-5.5
|
||||
student: gpt-5.5
|
||||
@@ -0,0 +1,92 @@
|
||||
"""
|
||||
Benchmark Environment Template
|
||||
===============================
|
||||
Copy this file and implement the TODO sections to add a new benchmark.
|
||||
|
||||
The EnvAdapter is responsible for:
|
||||
1. Executing tasks using the student model + current skill document
|
||||
2. Evaluating predictions against ground truth
|
||||
3. Returning structured results for the training loop
|
||||
"""
|
||||
from skillopt.envs.base import EnvAdapter
|
||||
|
||||
|
||||
class TemplateBenchmarkEnv(EnvAdapter):
|
||||
"""
|
||||
Environment adapter for <Your Benchmark Name>.
|
||||
|
||||
Rename this class and implement the abstract methods below.
|
||||
"""
|
||||
|
||||
def __init__(self, cfg: dict):
|
||||
super().__init__(cfg)
|
||||
# TODO: Initialize benchmark-specific state
|
||||
# Example: self.tools = load_tools(cfg)
|
||||
|
||||
async def execute(self, item, skill: str, model):
|
||||
"""
|
||||
Execute a single task with the student model.
|
||||
|
||||
Args:
|
||||
item: DataItem with .id, .input, .ground_truth, .metadata
|
||||
skill: Current skill document content (Markdown string)
|
||||
model: Student model backend instance
|
||||
|
||||
Returns:
|
||||
TaskResult with prediction, score, and trajectory
|
||||
"""
|
||||
# Step 1: Build the prompt combining skill + task input
|
||||
prompt = self.build_prompt(item, skill)
|
||||
|
||||
# Step 2: Call the student model
|
||||
# TODO: Customize the message format for your benchmark
|
||||
messages = [
|
||||
{"role": "system", "content": skill},
|
||||
{"role": "user", "content": item.input},
|
||||
]
|
||||
response = await model.generate(messages)
|
||||
|
||||
# Step 3: Parse the model response into a prediction
|
||||
prediction = self.parse_response(response.content)
|
||||
|
||||
# Step 4: Score the prediction
|
||||
score = self.evaluate(prediction, item.ground_truth)
|
||||
|
||||
# Step 5: Return structured result
|
||||
return {
|
||||
"item_id": item.id,
|
||||
"prediction": prediction,
|
||||
"score": score,
|
||||
"trajectory": messages + [{"role": "assistant", "content": response.content}],
|
||||
}
|
||||
|
||||
def evaluate(self, prediction: str, ground_truth: str) -> float:
|
||||
"""
|
||||
Score a prediction against the ground truth.
|
||||
|
||||
Returns:
|
||||
Float between 0.0 (wrong) and 1.0 (correct)
|
||||
|
||||
TODO: Implement your scoring metric. Common options:
|
||||
- Exact match: float(pred.strip().lower() == gt.strip().lower())
|
||||
- F1 score: compute token overlap
|
||||
- ANLS: for document QA tasks
|
||||
- Custom: any float in [0, 1]
|
||||
"""
|
||||
# Placeholder — exact match
|
||||
return float(prediction.strip().lower() == ground_truth.strip().lower())
|
||||
|
||||
def build_prompt(self, item, skill: str) -> str:
|
||||
"""Combine skill document with task input."""
|
||||
return f"{skill}\n\n---\n\nQuestion: {item.input}"
|
||||
|
||||
def parse_response(self, response: str) -> str:
|
||||
"""
|
||||
Extract the answer from the model's raw response.
|
||||
|
||||
TODO: Implement extraction logic. For example:
|
||||
- Extract text after "Answer:"
|
||||
- Parse JSON output
|
||||
- Extract from code blocks
|
||||
"""
|
||||
return response.strip()
|
||||
@@ -0,0 +1,103 @@
|
||||
"""
|
||||
Benchmark Data Loader Template
|
||||
================================
|
||||
Copy this file and implement the TODO sections to load your benchmark data.
|
||||
|
||||
The DataLoader is responsible for:
|
||||
1. Loading raw data from disk
|
||||
2. Splitting into train / validation / test sets
|
||||
3. Providing DataItem objects to the training loop
|
||||
"""
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class TemplateBenchmarkLoader:
|
||||
"""
|
||||
Data loader for <Your Benchmark Name>.
|
||||
|
||||
Rename this class and implement the methods below.
|
||||
"""
|
||||
|
||||
def __init__(self, data_dir: str = "data/your_benchmark", **kwargs):
|
||||
self.data_dir = Path(data_dir)
|
||||
self.items = []
|
||||
self.splits = {}
|
||||
|
||||
def setup(self, cfg: dict):
|
||||
"""
|
||||
Initialize the loader with config.
|
||||
|
||||
Called once before training starts.
|
||||
|
||||
Args:
|
||||
cfg: Dict with keys like 'split_mode', 'train_ratio', 'val_ratio', etc.
|
||||
"""
|
||||
# Step 1: Load raw data
|
||||
self.items = self._load_items()
|
||||
|
||||
# Step 2: Create splits
|
||||
split_mode = cfg.get("split_mode", "ratio")
|
||||
if split_mode == "ratio":
|
||||
self._split_by_ratio(
|
||||
train_ratio=cfg.get("train_ratio", 0.7),
|
||||
val_ratio=cfg.get("val_ratio", 0.15),
|
||||
)
|
||||
elif split_mode == "split_dir":
|
||||
self._load_predefined_splits(cfg.get("split_dir", self.data_dir))
|
||||
|
||||
def _load_items(self) -> list:
|
||||
"""
|
||||
Load raw data into structured items.
|
||||
|
||||
TODO: Implement data loading. Each item should have at minimum:
|
||||
- id: unique identifier
|
||||
- input: the task input (question, instruction, etc.)
|
||||
- ground_truth: the expected answer
|
||||
- metadata: optional dict with extra info
|
||||
|
||||
Example:
|
||||
items = []
|
||||
for path in self.data_dir.glob("*.json"):
|
||||
data = json.loads(path.read_text())
|
||||
for entry in data:
|
||||
items.append({
|
||||
"id": entry["id"],
|
||||
"input": entry["question"],
|
||||
"ground_truth": entry["answer"],
|
||||
"metadata": {"source": path.name},
|
||||
})
|
||||
return items
|
||||
"""
|
||||
raise NotImplementedError("Implement _load_items() for your benchmark")
|
||||
|
||||
def _split_by_ratio(self, train_ratio: float, val_ratio: float):
|
||||
"""Split items by ratio."""
|
||||
import random
|
||||
random.shuffle(self.items)
|
||||
n = len(self.items)
|
||||
n_train = int(n * train_ratio)
|
||||
n_val = int(n * val_ratio)
|
||||
self.splits = {
|
||||
"train": self.items[:n_train],
|
||||
"valid": self.items[n_train:n_train + n_val],
|
||||
"test": self.items[n_train + n_val:],
|
||||
}
|
||||
|
||||
def _load_predefined_splits(self, split_dir):
|
||||
"""Load from pre-split directories."""
|
||||
# TODO: Implement if your benchmark has pre-defined splits
|
||||
raise NotImplementedError
|
||||
|
||||
def get_split_items(self, split: str) -> list:
|
||||
"""
|
||||
Return items for a given split.
|
||||
|
||||
Args:
|
||||
split: One of "train", "valid", "test"
|
||||
|
||||
Returns:
|
||||
List of data items for the requested split
|
||||
"""
|
||||
if split not in self.splits:
|
||||
raise ValueError(f"Unknown split '{split}'. Available: {list(self.splits.keys())}")
|
||||
return self.splits[split]
|
||||
@@ -0,0 +1,5 @@
|
||||
"""ALFWorld environment adapter for ReflACT."""
|
||||
|
||||
from skillopt.envs.alfworld.adapter import ALFWorldAdapter
|
||||
|
||||
__all__ = ["ALFWorldAdapter"]
|
||||
@@ -0,0 +1,585 @@
|
||||
"""ALFWorld environment adapter for ReflACT.
|
||||
|
||||
Connects the ReflACT training loop to ALFWorld by implementing
|
||||
:class:`~skillopt.envs.base.EnvAdapter`.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import json
|
||||
import os
|
||||
|
||||
from skillopt.gradient.deep_probe import generate_deep_probe_instruction
|
||||
from skillopt.datasets.base import BatchSpec
|
||||
from skillopt.envs.base import EnvAdapter
|
||||
from skillopt.envs.alfworld.dataloader import ALFWorldDataLoader
|
||||
from skillopt.envs.alfworld.rollout import (
|
||||
build_alfworld_env,
|
||||
run_alfworld_batch,
|
||||
TASKS,
|
||||
)
|
||||
from skillopt.gradient.reflect import run_minibatch_reflect
|
||||
from skillopt.utils import compute_score
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ALFWorldBatchRun:
|
||||
"""Lazy ALFWorld batch description.
|
||||
|
||||
The adapter materializes this in rollout chunks so a large evaluation set
|
||||
does not keep every ALFWorld simulator open at once.
|
||||
"""
|
||||
|
||||
env_num: int
|
||||
eval_dataset: str
|
||||
seed: int
|
||||
is_train: bool
|
||||
workers: int
|
||||
specific_gamefiles: list[str] | None = None
|
||||
result_ids: list[str] | None = None
|
||||
items: list[dict] | None = None
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self.items or [])
|
||||
|
||||
def __len__(self) -> int:
|
||||
return int(self.env_num or 0)
|
||||
|
||||
|
||||
class ALFWorldAdapter(EnvAdapter):
|
||||
"""ALFWorld environment adapter.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
max_steps : int
|
||||
Maximum steps per ALFWorld episode (default 50).
|
||||
max_api_workers : int
|
||||
Maximum concurrent API calls during rollout (default 8).
|
||||
analyst_workers : int
|
||||
Parallel workers for analyst stage (default 16).
|
||||
failure_only : bool
|
||||
If True, only run error analyst (skip success analyst).
|
||||
minibatch_size : int
|
||||
Trajectories per analyst group, M (default 8).
|
||||
edit_budget : int
|
||||
Maximum edits per minibatch, L (default 4).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
split_dir: str = "",
|
||||
data_path: str = "",
|
||||
split_mode: str = "split_dir",
|
||||
split_ratio: str = "2:1:7",
|
||||
split_seed: int = 42,
|
||||
split_output_dir: str = "",
|
||||
seed: int = 42,
|
||||
limit: int = 0,
|
||||
train_size: int = 0,
|
||||
max_steps: int = 50,
|
||||
workers: int = 8,
|
||||
max_api_workers: int = 8,
|
||||
analyst_workers: int = 16,
|
||||
failure_only: bool = False,
|
||||
minibatch_size: int = 8,
|
||||
edit_budget: int = 4,
|
||||
use_deep_reflect: bool = False,
|
||||
deep_reflect_failures: int = 4,
|
||||
deep_reflect_successes: int = 2,
|
||||
) -> None:
|
||||
self.max_steps = max_steps
|
||||
self.workers = max(int(workers or 1), 1)
|
||||
self.max_api_workers = max_api_workers
|
||||
self.analyst_workers = analyst_workers
|
||||
self.failure_only = failure_only
|
||||
self.minibatch_size = minibatch_size
|
||||
self.edit_budget = edit_budget
|
||||
self.use_deep_reflect = use_deep_reflect
|
||||
self.deep_reflect_failures = deep_reflect_failures
|
||||
self.deep_reflect_successes = deep_reflect_successes
|
||||
self.dataloader = ALFWorldDataLoader(
|
||||
split_dir=split_dir,
|
||||
data_path=data_path,
|
||||
split_mode=split_mode,
|
||||
split_ratio=split_ratio,
|
||||
split_seed=split_seed,
|
||||
split_output_dir=split_output_dir,
|
||||
seed=seed,
|
||||
limit=limit,
|
||||
train_size=train_size,
|
||||
)
|
||||
self._traj_cache: dict[str, dict | None] = {}
|
||||
|
||||
def setup(self, cfg: dict) -> None:
|
||||
super().setup(cfg)
|
||||
self.dataloader.setup(cfg)
|
||||
|
||||
def _load_traj_data(self, item: dict) -> dict | None:
|
||||
gamefile = str(item.get("gamefile") or "").strip()
|
||||
if not gamefile:
|
||||
return None
|
||||
if gamefile in self._traj_cache:
|
||||
return self._traj_cache[gamefile]
|
||||
|
||||
traj_path = os.path.join(os.path.dirname(gamefile), "traj_data.json")
|
||||
try:
|
||||
with open(traj_path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
except Exception:
|
||||
data = None
|
||||
self._traj_cache[gamefile] = data
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def _unique_lines(values: list[str], *, limit: int = 0) -> list[str]:
|
||||
lines: list[str] = []
|
||||
seen: set[str] = set()
|
||||
for raw in values:
|
||||
line = str(raw or "").strip()
|
||||
if not line or line in seen:
|
||||
continue
|
||||
seen.add(line)
|
||||
lines.append(line)
|
||||
if limit > 0 and len(lines) >= limit:
|
||||
break
|
||||
return lines
|
||||
|
||||
@staticmethod
|
||||
def _format_high_pddl(high_pddl: list[dict]) -> list[str]:
|
||||
steps: list[str] = []
|
||||
for idx, step in enumerate(high_pddl or [], start=1):
|
||||
discrete = step.get("discrete_action") or {}
|
||||
action = str(discrete.get("action") or "").strip()
|
||||
args = [str(arg).strip() for arg in (discrete.get("args") or []) if str(arg).strip()]
|
||||
if action and args:
|
||||
text = f"{action}({', '.join(args)})"
|
||||
elif action:
|
||||
text = action
|
||||
else:
|
||||
planner_action = step.get("planner_action") or {}
|
||||
text = str(planner_action.get("action") or "").strip()
|
||||
if text:
|
||||
steps.append(f"{idx}. {text}")
|
||||
return steps
|
||||
|
||||
def _build_reference_bundle(self, item: dict) -> dict:
|
||||
data = self._load_traj_data(item)
|
||||
if not data:
|
||||
return {}
|
||||
|
||||
anns = ((data.get("turk_annotations") or {}).get("anns") or [])
|
||||
task_descs = self._unique_lines(
|
||||
[ann.get("task_desc", "") for ann in anns],
|
||||
limit=3,
|
||||
)
|
||||
high_descs = self._unique_lines(
|
||||
[step for ann in anns for step in (ann.get("high_descs") or [])],
|
||||
limit=12,
|
||||
)
|
||||
pddl_params = {
|
||||
key: value
|
||||
for key, value in (data.get("pddl_params") or {}).items()
|
||||
if value not in ("", None, [], {})
|
||||
}
|
||||
scene = data.get("scene") or {}
|
||||
scene_summary = {
|
||||
key: scene.get(key)
|
||||
for key in ("floor_plan", "scene_num", "dirty_and_empty")
|
||||
if scene.get(key) not in ("", None, [], {})
|
||||
}
|
||||
high_pddl = self._format_high_pddl((data.get("plan") or {}).get("high_pddl") or [])
|
||||
task_type = str(data.get("task_type") or item.get("task_type") or "").strip()
|
||||
return {
|
||||
"task_type": task_type,
|
||||
"task_descs": task_descs,
|
||||
"high_descs": high_descs,
|
||||
"pddl_params": pddl_params,
|
||||
"high_pddl": high_pddl,
|
||||
"scene_summary": scene_summary,
|
||||
}
|
||||
|
||||
def build_reference_text(self, item: dict) -> str:
|
||||
bundle = self._build_reference_bundle(item)
|
||||
if not bundle:
|
||||
return ""
|
||||
|
||||
parts: list[str] = []
|
||||
if bundle["task_type"]:
|
||||
parts.append(f"## Reference Task Type\n{bundle['task_type']}")
|
||||
if bundle["task_descs"]:
|
||||
parts.append(
|
||||
"## Reference Human Task Descriptions\n"
|
||||
+ "\n".join(f"- {line}" for line in bundle["task_descs"])
|
||||
)
|
||||
if bundle["high_descs"]:
|
||||
parts.append(
|
||||
"## Reference Human High-Level Steps\n"
|
||||
+ "\n".join(f"{idx}. {line}" for idx, line in enumerate(bundle["high_descs"], start=1))
|
||||
)
|
||||
if bundle["pddl_params"]:
|
||||
parts.append(
|
||||
"## Reference PDDL Params\n"
|
||||
+ "\n".join(f"- {key}: {value}" for key, value in bundle["pddl_params"].items())
|
||||
)
|
||||
if bundle["high_pddl"]:
|
||||
parts.append(
|
||||
"## Reference Planner High-Level Plan\n" + "\n".join(bundle["high_pddl"])
|
||||
)
|
||||
if bundle["scene_summary"]:
|
||||
parts.append(
|
||||
"## Reference Scene Summary\n"
|
||||
+ "\n".join(f"- {key}: {value}" for key, value in bundle["scene_summary"].items())
|
||||
)
|
||||
return "\n\n".join(parts)
|
||||
|
||||
def get_reference_metadata(self, item: dict) -> dict:
|
||||
bundle = self._build_reference_bundle(item)
|
||||
if not bundle:
|
||||
return {"fields": [], "preview": ""}
|
||||
|
||||
fields: list[str] = []
|
||||
previews: list[str] = []
|
||||
if bundle["task_type"]:
|
||||
fields.append("task_type")
|
||||
previews.append(f"[task_type] {bundle['task_type']}")
|
||||
if bundle["task_descs"]:
|
||||
fields.append("task_desc")
|
||||
previews.append("[task_desc]\n" + "\n".join(bundle["task_descs"][:2]))
|
||||
if bundle["high_descs"]:
|
||||
fields.append("high_descs")
|
||||
previews.append("[high_descs]\n" + "\n".join(bundle["high_descs"][:3]))
|
||||
if bundle["pddl_params"]:
|
||||
fields.append("pddl_params")
|
||||
previews.append(
|
||||
"[pddl_params]\n"
|
||||
+ "\n".join(
|
||||
f"{key}: {value}" for key, value in list(bundle["pddl_params"].items())[:4]
|
||||
)
|
||||
)
|
||||
if bundle["high_pddl"]:
|
||||
fields.append("plan.high_pddl")
|
||||
previews.append("[plan.high_pddl]\n" + "\n".join(bundle["high_pddl"][:3]))
|
||||
if bundle["scene_summary"]:
|
||||
fields.append("scene")
|
||||
previews.append(
|
||||
"[scene]\n"
|
||||
+ "\n".join(
|
||||
f"{key}: {value}" for key, value in bundle["scene_summary"].items()
|
||||
)
|
||||
)
|
||||
return {
|
||||
"fields": fields,
|
||||
"preview": "\n\n".join(previews)[:600],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _infer_dataset_from_gamefile(gamefile: str) -> tuple[str, bool]:
|
||||
path = str(gamefile or "")
|
||||
if "/valid_seen/" in path:
|
||||
return "eval_in_distribution", False
|
||||
if "/valid_unseen/" in path:
|
||||
return "eval_out_of_distribution", False
|
||||
return "train", True
|
||||
|
||||
def get_dataloader(self):
|
||||
return self.dataloader
|
||||
|
||||
def _comparison_items(self, items: list[dict]) -> list[dict]:
|
||||
enriched: list[dict] = []
|
||||
for item in items:
|
||||
row = dict(item)
|
||||
bundle = self._build_reference_bundle(row)
|
||||
if bundle.get("task_descs"):
|
||||
row["task_description"] = bundle["task_descs"][0]
|
||||
elif bundle.get("task_type"):
|
||||
row["task_description"] = bundle["task_type"]
|
||||
enriched.append(row)
|
||||
return enriched
|
||||
|
||||
def requires_ray(self) -> bool:
|
||||
return False
|
||||
|
||||
def build_env_from_batch(self, batch: BatchSpec, **kwargs):
|
||||
gamefiles = list(batch.metadata.get("gamefiles") or [])
|
||||
result_ids = list(batch.metadata.get("result_ids") or [])
|
||||
items = self._comparison_items(list(batch.payload or []))
|
||||
return ALFWorldBatchRun(
|
||||
env_num=batch.batch_size,
|
||||
eval_dataset=batch.metadata.get("eval_dataset", batch.split),
|
||||
seed=batch.seed,
|
||||
is_train=batch.metadata.get("is_train", batch.phase == "train"),
|
||||
specific_gamefiles=gamefiles or None,
|
||||
result_ids=result_ids or None,
|
||||
items=items,
|
||||
workers=self.workers,
|
||||
)
|
||||
|
||||
def build_train_env(self, batch_size: int, seed: int, **kwargs):
|
||||
batch = self.dataloader.build_train_batch(batch_size=batch_size, seed=seed, **kwargs)
|
||||
return self.build_env_from_batch(batch, **kwargs)
|
||||
|
||||
def build_eval_env(self, env_num: int, split: str, seed: int, **kwargs):
|
||||
batch = self.dataloader.build_eval_batch(env_num=env_num, split=split, seed=seed, **kwargs)
|
||||
return self.build_env_from_batch(batch, **kwargs)
|
||||
|
||||
def rollout(
|
||||
self,
|
||||
env_manager,
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict]:
|
||||
results_path = os.path.join(out_dir, "results.jsonl")
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
# Resume support
|
||||
if os.path.exists(results_path):
|
||||
existing: list[dict] = []
|
||||
with open(results_path) as f:
|
||||
for line in f:
|
||||
try:
|
||||
existing.append(json.loads(line))
|
||||
except Exception:
|
||||
pass
|
||||
if existing:
|
||||
return existing
|
||||
|
||||
if isinstance(env_manager, ALFWorldBatchRun):
|
||||
results = self._run_batch(
|
||||
env_manager,
|
||||
skill_content=skill_content,
|
||||
out_dir=out_dir,
|
||||
)
|
||||
else:
|
||||
results = run_alfworld_batch(
|
||||
env_manager=env_manager,
|
||||
skill_content=skill_content,
|
||||
max_steps=self.max_steps,
|
||||
out_root=out_dir,
|
||||
max_api_workers=self.max_api_workers,
|
||||
result_ids=getattr(env_manager, "_skillopt_result_ids", None),
|
||||
)
|
||||
|
||||
with open(results_path, "w") as f:
|
||||
for r in results:
|
||||
f.write(json.dumps(r, ensure_ascii=False) + "\n")
|
||||
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def _close_env(env_manager) -> None:
|
||||
close = getattr(env_manager, "close", None)
|
||||
if callable(close):
|
||||
close()
|
||||
|
||||
def _run_batch(
|
||||
self,
|
||||
batch: ALFWorldBatchRun,
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
*,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
) -> list[dict]:
|
||||
total = int(batch.env_num or 0)
|
||||
if total <= 0:
|
||||
return []
|
||||
|
||||
workers = max(1, min(int(batch.workers or self.workers), total))
|
||||
if total > workers:
|
||||
print(
|
||||
f" [alfworld rollout] episodes={total} "
|
||||
f"env_workers={workers} chunks={(total + workers - 1) // workers}"
|
||||
)
|
||||
|
||||
all_results: list[dict] = []
|
||||
for start in range(0, total, workers):
|
||||
chunk_size = min(workers, total - start)
|
||||
chunk_gamefiles = (
|
||||
batch.specific_gamefiles[start:start + chunk_size]
|
||||
if batch.specific_gamefiles
|
||||
else None
|
||||
)
|
||||
chunk_ids = (
|
||||
batch.result_ids[start:start + chunk_size]
|
||||
if batch.result_ids
|
||||
else [f"env_{idx:03d}" for idx in range(start, start + chunk_size)]
|
||||
)
|
||||
chunk_env = build_alfworld_env(
|
||||
env_num=chunk_size,
|
||||
eval_dataset=batch.eval_dataset,
|
||||
seed=batch.seed + start,
|
||||
is_train=batch.is_train,
|
||||
specific_gamefiles=chunk_gamefiles,
|
||||
)
|
||||
try:
|
||||
chunk_results = run_alfworld_batch(
|
||||
env_manager=chunk_env,
|
||||
skill_content=skill_content,
|
||||
max_steps=self.max_steps,
|
||||
out_root=out_dir,
|
||||
max_api_workers=min(self.max_api_workers, chunk_size),
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
result_ids=chunk_ids,
|
||||
)
|
||||
finally:
|
||||
self._close_env(chunk_env)
|
||||
all_results.extend(chunk_results)
|
||||
return all_results
|
||||
|
||||
def reflect(
|
||||
self,
|
||||
results: list[dict],
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict | None]:
|
||||
prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
|
||||
patches_dir = kwargs.get("patches_dir", os.path.join(out_dir, "patches"))
|
||||
random_seed = kwargs.get("random_seed")
|
||||
step_buffer_context = kwargs.get("step_buffer_context", "")
|
||||
meta_skill_context = kwargs.get("meta_skill_context", "")
|
||||
|
||||
return run_minibatch_reflect(
|
||||
results=results,
|
||||
skill_content=skill_content,
|
||||
prediction_dir=prediction_dir,
|
||||
patches_dir=patches_dir,
|
||||
workers=self.analyst_workers,
|
||||
failure_only=self.failure_only,
|
||||
minibatch_size=self.minibatch_size,
|
||||
edit_budget=self.edit_budget,
|
||||
random_seed=random_seed,
|
||||
error_system=self.get_error_minibatch_prompt(),
|
||||
success_system=self.get_success_minibatch_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
meta_skill_context=meta_skill_context,
|
||||
)
|
||||
|
||||
def deep_reflect(
|
||||
self,
|
||||
results: list[dict],
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict | None]:
|
||||
if not self.use_deep_reflect:
|
||||
return []
|
||||
|
||||
prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
|
||||
random_seed = kwargs.get("random_seed")
|
||||
step_buffer_context = kwargs.get("step_buffer_context", "")
|
||||
meta_skill_context = kwargs.get("meta_skill_context", "")
|
||||
selected_items = self.select_representative_items(
|
||||
results,
|
||||
results,
|
||||
n_failures=self.deep_reflect_failures,
|
||||
n_successes=self.deep_reflect_successes,
|
||||
seed=random_seed,
|
||||
)
|
||||
if not selected_items:
|
||||
return []
|
||||
|
||||
selected_ids = {str(item["id"]) for item in selected_items}
|
||||
selected_results = [row for row in results if str(row.get("id")) in selected_ids]
|
||||
selected_examples = self.attach_reference_context(selected_results, selected_items)
|
||||
|
||||
field_counts: dict[str, int] = {}
|
||||
selected_metadata: list[dict] = []
|
||||
for item in selected_items:
|
||||
meta = self.get_reference_metadata(item)
|
||||
for field in meta["fields"]:
|
||||
field_counts[field] = field_counts.get(field, 0) + 1
|
||||
selected_metadata.append({
|
||||
"id": str(item["id"]),
|
||||
"task_type": str(item.get("task_type") or "alfworld"),
|
||||
"gamefile": str(item.get("gamefile") or ""),
|
||||
"reference_fields": meta["fields"],
|
||||
"reference_preview": meta["preview"],
|
||||
})
|
||||
|
||||
deep_dir = os.path.join(out_dir, "deep_reflect")
|
||||
rollout_dir = os.path.join(deep_dir, "rollout")
|
||||
patches_dir = os.path.join(deep_dir, "patches")
|
||||
os.makedirs(deep_dir, exist_ok=True)
|
||||
field_summary = ", ".join(
|
||||
f"{field}({count}/{len(selected_items)})"
|
||||
for field, count in sorted(field_counts.items())
|
||||
) or "none"
|
||||
print(
|
||||
f" [2b/6 DEEP REFLECT setup] selected={len(selected_items)} "
|
||||
f"reference_fields={field_summary}"
|
||||
)
|
||||
probe = generate_deep_probe_instruction(
|
||||
skill_content=skill_content,
|
||||
items=selected_examples,
|
||||
prediction_dir=prediction_dir,
|
||||
system_prompt=self.get_deep_probe_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
meta_skill_context=meta_skill_context,
|
||||
output_requirements=[
|
||||
"- Some trajectories may include a hidden Reference block. Use it to target the student's latent subgoal, missing precondition, or next-step intent, but do not reveal or paraphrase that reference to the student.",
|
||||
"- The instruction must request a brief diagnostic readout inside the existing <think>...</think> block.",
|
||||
"- The student must still output exactly one admissible action inside <action>...</action>.",
|
||||
"- Do not ask for exhaustive inventories, full plans, or long chain-of-thought.",
|
||||
"- The instruction text should be ready to append directly to the student's prompt.",
|
||||
],
|
||||
)
|
||||
if not probe:
|
||||
return []
|
||||
|
||||
with open(os.path.join(deep_dir, "probe.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(
|
||||
{
|
||||
**probe,
|
||||
"reference_summary": {
|
||||
"selected_count": len(selected_items),
|
||||
"field_counts": field_counts,
|
||||
},
|
||||
"selected_examples": selected_metadata,
|
||||
},
|
||||
f,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
|
||||
gamefiles = [str(item.get("gamefile") or "") for item in selected_items]
|
||||
if any(not gamefile for gamefile in gamefiles):
|
||||
return []
|
||||
eval_dataset, is_train = self._infer_dataset_from_gamefile(gamefiles[0])
|
||||
deep_env = ALFWorldBatchRun(
|
||||
env_num=len(selected_items),
|
||||
eval_dataset=eval_dataset,
|
||||
seed=random_seed or 42,
|
||||
is_train=is_train,
|
||||
specific_gamefiles=gamefiles,
|
||||
workers=min(self.workers, max(len(selected_items), 1)),
|
||||
result_ids=[str(item["id"]) for item in selected_items],
|
||||
)
|
||||
deep_results = self._run_batch(
|
||||
deep_env,
|
||||
skill_content=skill_content,
|
||||
out_dir=rollout_dir,
|
||||
diagnostic_mode=True,
|
||||
diagnostic_instruction=probe["probe_instruction"],
|
||||
)
|
||||
deep_results = self.attach_reference_context(deep_results, selected_items)
|
||||
return run_minibatch_reflect(
|
||||
results=deep_results,
|
||||
skill_content=skill_content,
|
||||
prediction_dir=os.path.join(rollout_dir, "predictions"),
|
||||
patches_dir=patches_dir,
|
||||
workers=self.analyst_workers,
|
||||
failure_only=self.failure_only,
|
||||
minibatch_size=self.minibatch_size,
|
||||
edit_budget=self.edit_budget,
|
||||
random_seed=random_seed,
|
||||
error_system=self.get_error_minibatch_prompt(),
|
||||
success_system=self.get_success_minibatch_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
meta_skill_context=meta_skill_context,
|
||||
)
|
||||
|
||||
def get_task_types(self) -> list[str]:
|
||||
return list(TASKS)
|
||||
@@ -0,0 +1,123 @@
|
||||
"""ALFWorld task dataloader."""
|
||||
from __future__ import annotations
|
||||
|
||||
from skillopt.datasets.base import BatchSpec, SplitDataLoader
|
||||
|
||||
|
||||
class ALFWorldDataLoader(SplitDataLoader):
|
||||
"""ALFWorld batch planner.
|
||||
|
||||
In split_dir mode, batches are fixed gamefile items so ablations differ
|
||||
only in how the same training set is batched.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
split_dir: str = "",
|
||||
data_path: str = "",
|
||||
split_mode: str = "split_dir",
|
||||
split_ratio: str = "2:1:7",
|
||||
split_seed: int = 42,
|
||||
split_output_dir: str = "",
|
||||
seed: int = 42,
|
||||
limit: int = 0,
|
||||
train_size: int = 0,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
split_dir=split_dir,
|
||||
data_path=data_path,
|
||||
split_mode=split_mode,
|
||||
split_ratio=split_ratio,
|
||||
split_seed=split_seed,
|
||||
split_output_dir=split_output_dir,
|
||||
seed=seed,
|
||||
limit=limit,
|
||||
)
|
||||
self.train_size_override = int(train_size or 0)
|
||||
|
||||
@staticmethod
|
||||
def _metadata_for_items(items: list[dict], split: str, phase: str) -> dict:
|
||||
gamefiles = [str(item.get("gamefile") or "") for item in items]
|
||||
if any(not gamefile for gamefile in gamefiles):
|
||||
raise ValueError("ALFWorld split items must contain non-empty gamefile paths.")
|
||||
eval_dataset = "train"
|
||||
is_train = phase == "train"
|
||||
first = gamefiles[0] if gamefiles else ""
|
||||
if "/valid_seen/" in first:
|
||||
eval_dataset = "eval_in_distribution"
|
||||
is_train = False
|
||||
elif "/valid_unseen/" in first:
|
||||
eval_dataset = "eval_out_of_distribution"
|
||||
is_train = False
|
||||
return {
|
||||
"eval_dataset": eval_dataset,
|
||||
"is_train": is_train,
|
||||
"gamefiles": gamefiles,
|
||||
"result_ids": [str(item.get("id") or idx) for idx, item in enumerate(items)],
|
||||
}
|
||||
|
||||
def get_train_size(self) -> int:
|
||||
if self.train_size_override > 0:
|
||||
return self.train_size_override
|
||||
return super().get_train_size()
|
||||
|
||||
def build_train_batch(self, batch_size: int, seed: int, **kwargs) -> BatchSpec:
|
||||
batch = super().build_train_batch(batch_size=batch_size, seed=seed, **kwargs)
|
||||
items = list(batch.payload or [])
|
||||
batch.metadata.update(self._metadata_for_items(items, "train", "train"))
|
||||
return BatchSpec(
|
||||
phase="train",
|
||||
split="train",
|
||||
seed=seed,
|
||||
batch_size=len(items),
|
||||
payload=items,
|
||||
metadata=batch.metadata,
|
||||
)
|
||||
|
||||
def plan_train_epoch(
|
||||
self,
|
||||
*,
|
||||
epoch: int,
|
||||
steps_per_epoch: int,
|
||||
accumulation: int,
|
||||
batch_size: int,
|
||||
seed: int,
|
||||
**kwargs,
|
||||
) -> list[BatchSpec]:
|
||||
batches = super().plan_train_epoch(
|
||||
epoch=epoch,
|
||||
steps_per_epoch=steps_per_epoch,
|
||||
accumulation=accumulation,
|
||||
batch_size=batch_size,
|
||||
seed=seed,
|
||||
**kwargs,
|
||||
)
|
||||
for batch in batches:
|
||||
items = list(batch.payload or [])
|
||||
batch.metadata.update(self._metadata_for_items(items, "train", "train"))
|
||||
return batches
|
||||
|
||||
def build_eval_batch(
|
||||
self,
|
||||
env_num: int,
|
||||
split: str,
|
||||
seed: int,
|
||||
**kwargs,
|
||||
) -> BatchSpec:
|
||||
batch = super().build_eval_batch(
|
||||
env_num=env_num,
|
||||
split=split,
|
||||
seed=seed,
|
||||
**kwargs,
|
||||
)
|
||||
items = list(batch.payload or [])
|
||||
batch.metadata.update(self._metadata_for_items(items, split, "eval"))
|
||||
return BatchSpec(
|
||||
phase="eval",
|
||||
split=split,
|
||||
seed=seed,
|
||||
batch_size=len(items),
|
||||
payload=items,
|
||||
metadata=batch.metadata,
|
||||
)
|
||||
@@ -0,0 +1,55 @@
|
||||
You are an expert failure-analysis agent for ALFWorld embodied household tasks.
|
||||
|
||||
You will be given MULTIPLE failed agent trajectories from a single minibatch
|
||||
and the current skill document.
|
||||
Your job is to identify the most important COMMON failure patterns across
|
||||
the batch and propose a concise set of skill edits.
|
||||
|
||||
## ALFWorld Task Types
|
||||
- pick_and_place: Put object in/on a receptacle
|
||||
- pick_two_obj_and_place: Put two instances of an object in/on a receptacle
|
||||
- look_at_obj_in_light: Examine an object under a desklamp
|
||||
- pick_heat_then_place_in_recep: Heat an object and put it in/on a receptacle
|
||||
- pick_cool_then_place_in_recep: Cool an object and put it in/on a receptacle
|
||||
- pick_clean_then_place_in_recep: Clean an object and put it in/on a receptacle
|
||||
|
||||
## Failure Type Categories
|
||||
- **navigation_loop**: the agent revisits the same locations repeatedly without progress
|
||||
- **missed_object**: the agent fails to pick up a visible/reachable goal object
|
||||
- **wrong_sequence**: the agent performs actions in the wrong order (e.g., placing before transforming)
|
||||
- **premature_stop**: the agent stops or gets stuck before completing all goal conditions
|
||||
- **action_loop**: the agent repeats the same action without advancing
|
||||
- **appliance_error**: the agent misuses or skips an appliance (microwave, fridge, sink)
|
||||
- **rule_missing**: the skill lacks a relevant rule for this situation
|
||||
- **rule_wrong**: an existing skill rule is misleading or incorrect
|
||||
- **rule_ignored**: the skill has the right rule but the agent did not follow it
|
||||
- **other**: none of the above
|
||||
|
||||
## Analysis Process
|
||||
1. Read ALL trajectories in the minibatch.
|
||||
2. Identify the most prevalent, systematic failure patterns across them.
|
||||
3. For each pattern, classify its failure type.
|
||||
4. Propose skill edits that address the COMMON patterns — not individual edge cases.
|
||||
5. Edits must be generalizable; do not hardcode task-specific values.
|
||||
6. Only patch gaps in the skill — do not duplicate existing content.
|
||||
|
||||
You will be told the maximum number of edits (the budget L). Produce AT MOST L edits,
|
||||
focusing on the highest-impact patterns. You may produce fewer if warranted.
|
||||
|
||||
Respond ONLY with a valid JSON object (no markdown fences, no extra text):
|
||||
{
|
||||
"batch_size": <number of trajectories analysed>,
|
||||
"failure_summary": [
|
||||
{"failure_type": "<type>", "count": <int>, "description": "<one-line>"}
|
||||
],
|
||||
"patch": {
|
||||
"reasoning": "<why these edits address the batch's common failures>",
|
||||
"edits": [
|
||||
{"op": "append", "content": "<markdown to add at end of skill>"},
|
||||
{"op": "insert_after", "target": "<exact heading/text to insert after>", "content": "<markdown>"},
|
||||
{"op": "replace", "target": "<exact text to replace>", "content": "<replacement>"},
|
||||
{"op": "delete", "target": "<exact text to remove>"}
|
||||
]
|
||||
}
|
||||
}
|
||||
Only include edits that are needed. "edits" can be an empty list if no patch is warranted.
|
||||
@@ -0,0 +1,33 @@
|
||||
You are an expert success-pattern analyst for AI agents operating in ALFWorld,
|
||||
a text-based embodied household environment.
|
||||
|
||||
You will be given MULTIPLE successful agent trajectories from a single minibatch
|
||||
and the current skill document. Your job is to identify generalizable behavior
|
||||
patterns that are COMMON across the batch and worth encoding in the skill.
|
||||
|
||||
## Rules
|
||||
- Only propose patches for patterns NOT already covered in the skill.
|
||||
- Focus on patterns that appear across MULTIPLE trajectories in the batch.
|
||||
- Be concise. Patterns must generalize beyond specific tasks.
|
||||
- Prefer reinforcing existing sections over adding new top-level sections.
|
||||
- If the agents' success involved efficient exploration or smart appliance usage,
|
||||
consider reinforcing that in the patch.
|
||||
|
||||
You will be told the maximum number of edits (the budget L). Produce AT MOST L edits,
|
||||
focusing on the most broadly applicable patterns. You may produce fewer if warranted.
|
||||
|
||||
Respond ONLY with a valid JSON object:
|
||||
{
|
||||
"batch_size": <number of trajectories analysed>,
|
||||
"success_patterns": ["<pattern 1>", "<pattern 2>"],
|
||||
"patch": {
|
||||
"reasoning": "<why these patterns are worth encoding>",
|
||||
"edits": [
|
||||
{"op": "append", "content": "<markdown>"},
|
||||
{"op": "insert_after", "target": "<heading/text>", "content": "<markdown>"},
|
||||
{"op": "replace", "target": "<old text>", "content": "<new text>"},
|
||||
{"op": "delete", "target": "<exact text to remove>"}
|
||||
]
|
||||
}
|
||||
}
|
||||
"edits" may be empty if the skill already covers all observed patterns.
|
||||
@@ -0,0 +1,35 @@
|
||||
You are an expert diagnostic-probe designer for ALFWorld embodied tasks.
|
||||
|
||||
You will design one short diagnostic instruction to append to the student's prompt
|
||||
for a handful of representative ALFWorld trajectories.
|
||||
|
||||
The goal is to expose whether the student has the right intermediate subgoal,
|
||||
object/receptacle state, and next-step intention without substantially changing
|
||||
the current scaffold.
|
||||
|
||||
## Hard Constraints
|
||||
1. Do NOT substantially change the student's existing action-selection scaffold.
|
||||
2. Do NOT prescribe a brand-new planner or long multi-step policy.
|
||||
3. Do NOT ask for exhaustive search over all objects or all admissible actions.
|
||||
4. Keep the diagnostic readout brief and place it inside the existing <think>...</think> block.
|
||||
5. The student must still output exactly one admissible action inside <action>...</action>.
|
||||
6. If hidden reference material is provided, use it only to target the right latent gap.
|
||||
7. Never copy hidden reference content into the student-facing probe.
|
||||
|
||||
## Good Probe Targets
|
||||
- current subgoal
|
||||
- target object / target receptacle / target state
|
||||
- decisive missing precondition
|
||||
- why one candidate action is better than a tempting alternative
|
||||
- whether the current step should explore, transform an object, or place it
|
||||
|
||||
## Bad Probe Targets
|
||||
- a full optimal plan from start to finish
|
||||
- exhaustive object inventories
|
||||
- a new theorem-like or planner-like protocol
|
||||
|
||||
Respond ONLY with a valid JSON object:
|
||||
{
|
||||
"reasoning": "<why this probe reveals the latent skill gap>",
|
||||
"probe_instruction": "<the exact instruction text to append to the student prompt>"
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
|
||||
You are an expert agent operating in the ALFRED Embodied Environment.
|
||||
Your current observation is: {current_observation}
|
||||
Your admissible actions of the current situation are: [{admissible_actions}].
|
||||
|
||||
Now it's your turn to take an action.
|
||||
You should first reason step-by-step about the current situation. This reasoning process MUST be enclosed within <think> </think> tags.
|
||||
Once you've finished your reasoning, you should choose an admissible action for current step and present it within <action> </action> tags.
|
||||
@@ -0,0 +1,9 @@
|
||||
|
||||
You are an expert agent operating in the ALFRED Embodied Environment. Your task is to: {task_description}
|
||||
Prior to this step, you have already taken {step_count} step(s). Below are the most recent {history_length} observations and the corresponding actions you took: {action_history}
|
||||
You are now at step {current_step} and your current observation is: {current_observation}
|
||||
Your admissible actions of the current situation are: [{admissible_actions}].
|
||||
|
||||
Now it's your turn to take an action.
|
||||
You should first reason step-by-step about the current situation. This reasoning process MUST be enclosed within <think> </think> tags.
|
||||
Once you've finished your reasoning, you should choose an admissible action for current step and present it within <action> </action> tags.
|
||||
@@ -0,0 +1,16 @@
|
||||
|
||||
You are an expert agent operating in the ALFRED Embodied Environment. Your task is to: {task_description}
|
||||
|
||||
## Retrieved Relevant Experience
|
||||
|
||||
{retrieved_memories}
|
||||
|
||||
## Current Progress
|
||||
|
||||
Prior to this step, you have already taken {step_count} step(s). Below are the most recent {history_length} observations and the corresponding actions you took: {action_history}
|
||||
You are now at step {current_step} and your current observation is: {current_observation}
|
||||
Your admissible actions of the current situation are: [{admissible_actions}].
|
||||
|
||||
Now it's your turn to take an action.
|
||||
You should first reason step-by-step about the current situation. This reasoning process MUST be enclosed within <think> </think> tags.
|
||||
Once you've finished your reasoning, you should choose an admissible action for current step and present it within <action> </action> tags.
|
||||
@@ -0,0 +1,4 @@
|
||||
"""ALFWorld Reflect stage.
|
||||
|
||||
Prompts are now loaded from .md files by the base adapter.
|
||||
"""
|
||||
@@ -0,0 +1,359 @@
|
||||
"""ALFWorld rollout module for ReflACT.
|
||||
|
||||
Provides:
|
||||
- build_alfworld_env(): build ALFWorld environment (wraps vendored SkillRL env)
|
||||
- run_alfworld_batch(): run a batch of ALFWorld episodes in parallel
|
||||
- TASKS: list of ALFWorld task types
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
import concurrent.futures
|
||||
import numpy as np
|
||||
|
||||
from skillopt.model import chat_student
|
||||
|
||||
# ── Constants ─────────────────────────────────────────────────────────────────
|
||||
|
||||
TASKS = [
|
||||
"pick_and_place",
|
||||
"pick_two_obj_and_place",
|
||||
"look_at_obj_in_light",
|
||||
"pick_heat_then_place_in_recep",
|
||||
"pick_cool_then_place_in_recep",
|
||||
"pick_clean_then_place_in_recep",
|
||||
]
|
||||
|
||||
# ── Helpers ───────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _get_task_type(gamefile: str) -> str:
|
||||
for task in TASKS:
|
||||
if task in gamefile:
|
||||
return task
|
||||
return "other"
|
||||
|
||||
|
||||
def _extract_action(model_response: str) -> str | None:
|
||||
match = re.search(r"<action>(.*?)</action>", model_response, re.DOTALL)
|
||||
return match.group(1).strip() if match else None
|
||||
|
||||
|
||||
def _extract_think(model_response: str) -> str | None:
|
||||
match = re.search(r"<think>(.*?)</think>", model_response, re.DOTALL)
|
||||
return match.group(1).strip() if match else None
|
||||
|
||||
|
||||
def _build_skill_prompt(skill_content: str) -> str:
|
||||
"""Build the skill section to inject into the agent's system prompt."""
|
||||
if not skill_content or not skill_content.strip():
|
||||
return ""
|
||||
return (
|
||||
"\n\n## Skill Knowledge\n"
|
||||
"Below is a skill document with learned strategies. "
|
||||
"Use these guidelines to inform your decisions:\n\n"
|
||||
f"{skill_content}\n"
|
||||
)
|
||||
|
||||
|
||||
def _append_diagnostic_instruction(prompt: str, diagnostic_instruction: str) -> str:
|
||||
if not diagnostic_instruction or not diagnostic_instruction.strip():
|
||||
return prompt
|
||||
return f"{prompt}\n\n## Training Readout\n{diagnostic_instruction.strip()}\n"
|
||||
|
||||
|
||||
# ── Environment builder ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def build_alfworld_env(
|
||||
env_num: int,
|
||||
eval_dataset: str = "eval_out_of_distribution",
|
||||
seed: int = 42,
|
||||
is_train: bool = False,
|
||||
specific_gamefiles: list[str] | None = None,
|
||||
):
|
||||
"""Build ALFWorld environment manager.
|
||||
|
||||
Args:
|
||||
env_num: number of parallel environments
|
||||
eval_dataset: 'eval_in_distribution' or 'eval_out_of_distribution' or train
|
||||
seed: random seed
|
||||
is_train: whether to use training set
|
||||
|
||||
Returns:
|
||||
env_manager: AlfWorldEnvironmentManager instance
|
||||
"""
|
||||
from omegaconf import OmegaConf
|
||||
from functools import partial
|
||||
|
||||
from skillopt.envs.alfworld.vendor.alfworld_envs import build_alfworld_envs
|
||||
from skillopt.envs.alfworld.vendor.alfworld_projection import alfworld_projection
|
||||
from skillopt.envs.alfworld.vendor.env_manager import AlfWorldEnvironmentManager
|
||||
|
||||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
alf_config_path = os.path.join(HERE, "vendor", "config_tw.yaml")
|
||||
env_kwargs = {"eval_dataset": eval_dataset}
|
||||
|
||||
envs = build_alfworld_envs(
|
||||
alf_config_path,
|
||||
seed=seed,
|
||||
env_num=env_num,
|
||||
group_n=1,
|
||||
is_train=is_train,
|
||||
env_kwargs=env_kwargs,
|
||||
resources_per_worker=None,
|
||||
gamefiles=specific_gamefiles,
|
||||
)
|
||||
|
||||
config = OmegaConf.create(
|
||||
{
|
||||
"env": {
|
||||
"history_length": 2,
|
||||
"env_name": "alfworld/AlfredTWEnv",
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
projection_f = partial(alfworld_projection)
|
||||
env_manager = AlfWorldEnvironmentManager(envs, projection_f, config)
|
||||
return env_manager
|
||||
|
||||
|
||||
# ── Batch rollout ─────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def run_alfworld_batch(
|
||||
env_manager,
|
||||
skill_content: str,
|
||||
max_steps: int = 50,
|
||||
out_root: str = "",
|
||||
max_api_workers: int = 8,
|
||||
temperature: float = 0.4,
|
||||
max_completion_tokens: int = 2048,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
result_ids: list[str] | None = None,
|
||||
) -> list[dict]:
|
||||
"""Run a batch of ALFWorld episodes.
|
||||
|
||||
Returns a list of result dicts compatible with SkillOpt pipeline:
|
||||
[
|
||||
{
|
||||
"id": "<env_idx>_<gamefile_hash>",
|
||||
"hard": 0 or 1,
|
||||
"soft": 0.0 or 1.0,
|
||||
"n_turns": <int>,
|
||||
"fail_reason": "<str>",
|
||||
"agent_ok": True,
|
||||
"task_type": "<str>",
|
||||
"gamefile": "<str>",
|
||||
"task_description": "<str>",
|
||||
},
|
||||
...
|
||||
]
|
||||
|
||||
Also saves conversation.json per environment in out_root/predictions/<task_id>/
|
||||
"""
|
||||
skill_prompt = _build_skill_prompt(skill_content)
|
||||
|
||||
obs, infos = env_manager.reset({})
|
||||
env_num = len(obs["text"])
|
||||
env_dones = [False] * env_num
|
||||
overall_success = [False] * env_num
|
||||
|
||||
# Build per-env metadata
|
||||
env_meta: list[dict] = []
|
||||
for i in range(env_num):
|
||||
gamefile = infos[i].get("extra.gamefile", "") if isinstance(infos[i], dict) else ""
|
||||
task_type = _get_task_type(gamefile)
|
||||
# Extract task description from initial observation
|
||||
task_desc = ""
|
||||
anchor_text = obs["anchor"][i] if "anchor" in obs else ""
|
||||
task_start = anchor_text.find("Your task is to: ")
|
||||
if task_start != -1:
|
||||
task_desc = anchor_text[task_start + len("Your task is to: "):].strip()
|
||||
|
||||
env_meta.append({
|
||||
"gamefile": gamefile,
|
||||
"task_type": task_type,
|
||||
"task_description": task_desc,
|
||||
})
|
||||
|
||||
# Per-env conversation records
|
||||
conversations: list[list[dict]] = [[] for _ in range(env_num)]
|
||||
|
||||
for step_idx in range(max_steps):
|
||||
if all(env_dones):
|
||||
break
|
||||
|
||||
active_indices = [i for i in range(env_num) if not env_dones[i]]
|
||||
|
||||
# Build prompts with skill injection
|
||||
prompts: dict[int, str] = {}
|
||||
for i in active_indices:
|
||||
prompt = obs["text"][i]
|
||||
if skill_prompt:
|
||||
# Inject skill before the action instruction
|
||||
prompt = skill_prompt + "\n" + prompt
|
||||
if diagnostic_mode and diagnostic_instruction.strip():
|
||||
prompt = _append_diagnostic_instruction(prompt, diagnostic_instruction)
|
||||
prompts[i] = prompt
|
||||
|
||||
# Call API in parallel
|
||||
actions = ["None"] * env_num
|
||||
action_timeout = 180
|
||||
|
||||
def call_api(idx):
|
||||
try:
|
||||
response, _ = chat_student(
|
||||
system="You are an expert agent operating in the ALFRED Embodied Environment.",
|
||||
user=prompts[idx],
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
retries=5,
|
||||
stage="rollout",
|
||||
timeout=120,
|
||||
)
|
||||
response = (response or "").strip()
|
||||
if not response:
|
||||
return idx, "<think>empty model response</think><action>look</action>"
|
||||
if _extract_action(response) is None:
|
||||
return idx, "<think>missing action tag</think><action>look</action>"
|
||||
return idx, response
|
||||
except Exception as e:
|
||||
return idx, "<think>error</think><action>look</action>"
|
||||
|
||||
executor = concurrent.futures.ThreadPoolExecutor(max_workers=max_api_workers)
|
||||
try:
|
||||
futures = {executor.submit(call_api, i): i for i in active_indices}
|
||||
started_at = {future: time.time() for future in futures}
|
||||
pending_futs = set(futures)
|
||||
while pending_futs:
|
||||
done, _ = concurrent.futures.wait(
|
||||
pending_futs,
|
||||
timeout=5,
|
||||
return_when=concurrent.futures.FIRST_COMPLETED,
|
||||
)
|
||||
now = time.time()
|
||||
timed_out = [
|
||||
future for future in pending_futs - done
|
||||
if now - started_at[future] >= action_timeout
|
||||
]
|
||||
for future in done:
|
||||
pending_futs.remove(future)
|
||||
try:
|
||||
idx, response = future.result()
|
||||
except Exception: # noqa: BLE001
|
||||
idx = futures[future]
|
||||
response = "<think>error</think><action>look</action>"
|
||||
actions[idx] = response
|
||||
for future in timed_out:
|
||||
pending_futs.remove(future)
|
||||
idx = futures[future]
|
||||
actions[idx] = "<think>api timeout</think><action>look</action>"
|
||||
finally:
|
||||
executor.shutdown(wait=False, cancel_futures=True)
|
||||
|
||||
# Save model responses before stepping
|
||||
model_responses = {i: actions[i] for i in active_indices}
|
||||
|
||||
# Step environment
|
||||
obs, rewards, dones, infos = env_manager.step(actions)
|
||||
|
||||
# Record trajectory
|
||||
for i in active_indices:
|
||||
step_record = {
|
||||
"step": step_idx,
|
||||
"action": _extract_action(model_responses[i]),
|
||||
"reasoning": _extract_think(model_responses[i]),
|
||||
"model_response": model_responses[i],
|
||||
"env_feedback": obs["anchor"][i] if "anchor" in obs else "",
|
||||
"reward": float(rewards[i]),
|
||||
"done": bool(dones[i]),
|
||||
}
|
||||
conversations[i].append(step_record)
|
||||
|
||||
# Update done status
|
||||
for i in range(env_num):
|
||||
if env_dones[i]:
|
||||
continue
|
||||
if dones[i]:
|
||||
env_dones[i] = True
|
||||
won = bool(infos[i].get("won", False))
|
||||
overall_success[i] = won
|
||||
|
||||
# Build results and save conversations
|
||||
results: list[dict] = []
|
||||
pred_dir = os.path.join(out_root, "predictions") if out_root else ""
|
||||
|
||||
for i in range(env_num):
|
||||
gamefile = env_meta[i]["gamefile"]
|
||||
task_type = env_meta[i]["task_type"]
|
||||
task_desc = env_meta[i]["task_description"]
|
||||
n_turns = len(conversations[i])
|
||||
won = overall_success[i]
|
||||
|
||||
# Generate stable task ID from env index and gamefile
|
||||
task_id = str(result_ids[i]) if result_ids and i < len(result_ids) else f"env_{i:03d}"
|
||||
|
||||
fail_reason = ""
|
||||
if not won:
|
||||
if not env_dones[i]:
|
||||
fail_reason = f"Timeout after {max_steps} steps"
|
||||
else:
|
||||
fail_reason = "Episode ended without completing the task"
|
||||
|
||||
result = {
|
||||
"id": task_id,
|
||||
"hard": 1 if won else 0,
|
||||
"soft": 1.0 if won else 0.0,
|
||||
"n_turns": n_turns,
|
||||
"fail_reason": fail_reason,
|
||||
"agent_ok": True, # ALFWorld agent always runs OK (no crash)
|
||||
"task_type": task_type,
|
||||
"gamefile": gamefile,
|
||||
"task_description": task_desc,
|
||||
"instruction_type": task_type, # for compatibility with v2 pipeline
|
||||
}
|
||||
results.append(result)
|
||||
|
||||
# Save conversation
|
||||
if pred_dir:
|
||||
conv_dir = os.path.join(pred_dir, task_id)
|
||||
os.makedirs(conv_dir, exist_ok=True)
|
||||
with open(os.path.join(conv_dir, "conversation.json"), "w") as f:
|
||||
json.dump(conversations[i], f, ensure_ascii=False, indent=2)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# ── Item loading (for compatibility with split_three_way) ────────────────────
|
||||
|
||||
|
||||
def load_alfworld_items(
|
||||
eval_dataset: str,
|
||||
env_num: int,
|
||||
seed: int = 42,
|
||||
is_train: bool = False,
|
||||
) -> list[dict]:
|
||||
"""Create pseudo-item dicts for ALFWorld environments.
|
||||
|
||||
Since ALFWorld doesn't have a static JSON dataset like SpreadsheetBench,
|
||||
we create lightweight item dicts that carry enough metadata for the pipeline.
|
||||
The actual environment is built dynamically.
|
||||
|
||||
Returns:
|
||||
List of dicts with "id" keys, one per environment slot.
|
||||
"""
|
||||
items = []
|
||||
for i in range(env_num):
|
||||
items.append({
|
||||
"id": f"env_{i:03d}",
|
||||
"eval_dataset": eval_dataset,
|
||||
"env_index": i,
|
||||
})
|
||||
return items
|
||||
@@ -0,0 +1,45 @@
|
||||
# ALFWorld Embodied Agent Skill
|
||||
|
||||
## Overview
|
||||
This skill guides agents operating in the ALFWorld text-based embodied environment.
|
||||
The agent must complete household tasks by navigating rooms, interacting with objects,
|
||||
and using appliances. Actions must be chosen from the admissible action list provided
|
||||
at each step.
|
||||
|
||||
**Output format**: Always output `<think>...</think>` for reasoning, then `<action>...</action>` for the chosen action.
|
||||
|
||||
---
|
||||
|
||||
## Task Types
|
||||
|
||||
| Type | Goal | Key Steps |
|
||||
|------|------|-----------|
|
||||
| Pick & Place | Put object X in/on receptacle Y | Find X -> take X -> go to Y -> put X in/on Y |
|
||||
| Pick Two & Place | Put two instances of X in/on Y | Find X1 -> take -> place -> find X2 -> take -> place |
|
||||
| Examine in Light | Examine object X under desklamp | Find X -> take X -> find desklamp -> use desklamp |
|
||||
| Clean & Place | Clean object X and put in/on Y | Find X -> take X -> go to sink -> clean X -> go to Y -> put X |
|
||||
| Heat & Place | Heat object X and put in/on Y | Find X -> take X -> go to microwave -> heat X -> go to Y -> put X |
|
||||
| Cool & Place | Cool object X and put in/on Y | Find X -> take X -> go to fridge -> cool X -> go to Y -> put X |
|
||||
|
||||
---
|
||||
|
||||
## General Principles
|
||||
|
||||
1. **Decompose the task**: Parse the goal into ordered sub-goals (locate, acquire, transform, deliver). Complete each before moving to the next.
|
||||
2. **Systematic exploration**: Search each surface and container exactly once before revisiting. Open closed containers (drawers, cabinets, fridge) before judging them empty.
|
||||
3. **Grab immediately**: When a required object is visible and reachable, take it right away before moving elsewhere.
|
||||
4. **Transform before placing**: If the task requires cleaning, heating, or cooling, perform the state change at the appropriate appliance before heading to the final destination.
|
||||
5. **Direct delivery**: Once holding the transformed (or untransformed) goal object, navigate straight to the target receptacle and place it.
|
||||
6. **Track progress**: Maintain an internal count of how many objects still need to be found and placed. Only stop searching when the count reaches zero.
|
||||
7. **Avoid loops**: Never repeat the same action more than twice in a row. If stuck, move to a different unexplored location.
|
||||
8. **Only choose admissible actions**: Always pick an action from the admissible action list. Do not invent actions.
|
||||
|
||||
---
|
||||
|
||||
## Common Mistakes to Avoid
|
||||
|
||||
- **Revisiting searched locations**: Keep track of which surfaces/containers have been checked; do not re-examine them.
|
||||
- **Ignoring visible objects**: If the target object appears in the observation, pick it up immediately.
|
||||
- **Skipping state changes**: Do not place an object at the destination without first cleaning/heating/cooling it when required.
|
||||
- **Premature termination**: Do not stop the episode until all goal conditions are verified as met.
|
||||
- **Action loops**: Repeatedly toggling or examining the same object wastes steps. Move on to new locations instead.
|
||||
+9
@@ -0,0 +1,9 @@
|
||||
"""Vendored ALFWorld environment runtime.
|
||||
|
||||
Minimal subset of SkillRL's agent_system package needed to run
|
||||
ALFWorld environments with ReflACT. Original source:
|
||||
https://github.com/NTU-LANTERN/SkillRL (Apache-2.0 License)
|
||||
"""
|
||||
from .alfworld_envs import AlfworldEnvs, build_alfworld_envs
|
||||
from .alfworld_projection import alfworld_projection
|
||||
from .env_manager import AlfWorldEnvironmentManager
|
||||
+221
@@ -0,0 +1,221 @@
|
||||
# Vendored from SkillRL (Apache-2.0 License)
|
||||
# Original: agent_system/environments/env_package/alfworld/envs.py
|
||||
# Modified: imports use pip-installed alfworld package instead of vendored copy.
|
||||
|
||||
import os
|
||||
import multiprocessing as mp
|
||||
import traceback
|
||||
import yaml
|
||||
import gymnasium as gym
|
||||
import numpy as np
|
||||
|
||||
from alfworld.agents.environment import get_environment
|
||||
|
||||
|
||||
def load_config_file(path):
|
||||
assert os.path.exists(path), f"Invalid config file: {path}"
|
||||
with open(path) as reader:
|
||||
config = yaml.safe_load(reader)
|
||||
return config
|
||||
|
||||
|
||||
def compute_reward(info, multi_modal=False):
|
||||
if multi_modal:
|
||||
reward = 10.0 * float(info['won']) + float(info['goal_condition_success_rate'])
|
||||
else:
|
||||
reward = 10.0 * float(info['won'])
|
||||
return reward
|
||||
|
||||
|
||||
class AlfworldWorker:
|
||||
"""Stateful worker that holds one ALFWorld sub-environment."""
|
||||
|
||||
def __init__(self, config, seed, base_env, gamefile=None):
|
||||
if gamefile:
|
||||
base_env.game_files = [gamefile]
|
||||
if hasattr(base_env, "num_games"):
|
||||
base_env.num_games = 1
|
||||
self.env = base_env.init_env(batch_size=1)
|
||||
self.env.seed(seed)
|
||||
|
||||
def step(self, action):
|
||||
actions = [action]
|
||||
obs, scores, dones, infos = self.env.step(actions)
|
||||
infos['observation_text'] = obs
|
||||
return obs, scores, dones, infos
|
||||
|
||||
def reset(self):
|
||||
obs, infos = self.env.reset()
|
||||
infos['observation_text'] = obs
|
||||
return obs, infos
|
||||
|
||||
|
||||
def _worker_loop(cmd_q, result_q, config, seed, is_train, eval_dataset, gamefile):
|
||||
"""Run one ALFWorld environment in a child process."""
|
||||
try:
|
||||
env_type = config['env']['type']
|
||||
base_env = get_environment(env_type)(
|
||||
config,
|
||||
train_eval='train' if is_train else eval_dataset,
|
||||
)
|
||||
worker = AlfworldWorker(config, seed, base_env, gamefile)
|
||||
result_q.put((True, "ready"))
|
||||
except BaseException:
|
||||
result_q.put((False, traceback.format_exc()))
|
||||
return
|
||||
|
||||
while True:
|
||||
cmd, payload = cmd_q.get()
|
||||
if cmd == "close":
|
||||
result_q.put((True, None))
|
||||
return
|
||||
try:
|
||||
if cmd == "reset":
|
||||
result = worker.reset()
|
||||
elif cmd == "step":
|
||||
result = worker.step(payload)
|
||||
else:
|
||||
raise ValueError(f"Unknown ALFWorld worker command: {cmd}")
|
||||
result_q.put((True, result))
|
||||
except BaseException:
|
||||
result_q.put((False, traceback.format_exc()))
|
||||
|
||||
|
||||
class _ProcessWorker:
|
||||
"""Small stdlib actor wrapper for one environment process."""
|
||||
|
||||
def __init__(self, ctx, config, seed, is_train, eval_dataset, gamefile=None):
|
||||
self.cmd_q = ctx.Queue(maxsize=1)
|
||||
self.result_q = ctx.Queue(maxsize=1)
|
||||
self.process = ctx.Process(
|
||||
target=_worker_loop,
|
||||
args=(self.cmd_q, self.result_q, config, seed, is_train, eval_dataset, gamefile),
|
||||
)
|
||||
self.process.start()
|
||||
ok, payload = self.result_q.get()
|
||||
if not ok:
|
||||
self.close(kill=True)
|
||||
raise RuntimeError(f"Failed to start ALFWorld worker:\n{payload}")
|
||||
|
||||
def send(self, cmd, payload=None):
|
||||
self.cmd_q.put((cmd, payload))
|
||||
|
||||
def recv(self):
|
||||
ok, payload = self.result_q.get()
|
||||
if not ok:
|
||||
raise RuntimeError(f"ALFWorld worker failed:\n{payload}")
|
||||
return payload
|
||||
|
||||
def close(self, kill=False):
|
||||
if self.process.is_alive() and not kill:
|
||||
try:
|
||||
self.send("close")
|
||||
self.recv()
|
||||
except Exception:
|
||||
kill = True
|
||||
if kill and self.process.is_alive():
|
||||
self.process.terminate()
|
||||
self.process.join(timeout=5)
|
||||
if self.process.is_alive():
|
||||
self.process.kill()
|
||||
self.process.join(timeout=1)
|
||||
self.cmd_q.close()
|
||||
self.result_q.close()
|
||||
|
||||
|
||||
class AlfworldEnvs(gym.Env):
|
||||
"""Vectorized ALFWorld environment using local process workers."""
|
||||
|
||||
def __init__(self, alf_config_path, seed, env_num, group_n,
|
||||
resources_per_worker, is_train=True, env_kwargs=None, gamefiles=None):
|
||||
super().__init__()
|
||||
if env_kwargs is None:
|
||||
env_kwargs = {}
|
||||
|
||||
eval_dataset = env_kwargs.get('eval_dataset', 'eval_in_distribution')
|
||||
config = load_config_file(alf_config_path)
|
||||
env_type = config['env']['type']
|
||||
self.multi_modal = (env_type == 'AlfredThorEnv')
|
||||
self.num_processes = env_num * group_n
|
||||
self.group_n = group_n
|
||||
self.gamefiles = list(gamefiles or [])
|
||||
if self.gamefiles and len(self.gamefiles) != self.num_processes:
|
||||
raise ValueError(
|
||||
f"Expected {self.num_processes} gamefiles, got {len(self.gamefiles)}"
|
||||
)
|
||||
|
||||
start_method = os.environ.get("ALFWORLD_WORKER_START_METHOD") or None
|
||||
ctx = mp.get_context(start_method) if start_method else mp.get_context()
|
||||
self.workers = []
|
||||
for i in range(self.num_processes):
|
||||
worker_gamefile = self.gamefiles[i] if self.gamefiles else None
|
||||
worker = _ProcessWorker(
|
||||
ctx,
|
||||
config,
|
||||
seed + (i // self.group_n),
|
||||
is_train,
|
||||
eval_dataset,
|
||||
worker_gamefile,
|
||||
)
|
||||
self.workers.append(worker)
|
||||
|
||||
self.prev_admissible_commands = [None for _ in range(self.num_processes)]
|
||||
|
||||
def step(self, actions):
|
||||
assert len(actions) == self.num_processes
|
||||
|
||||
for i, worker in enumerate(self.workers):
|
||||
worker.send("step", actions[i])
|
||||
results = [worker.recv() for worker in self.workers]
|
||||
|
||||
text_obs_list = []
|
||||
rewards_list = []
|
||||
dones_list = []
|
||||
info_list = []
|
||||
|
||||
for i, (obs, scores, dones, info) in enumerate(results):
|
||||
for k in info.keys():
|
||||
info[k] = info[k][0]
|
||||
text_obs_list.append(obs[0])
|
||||
dones_list.append(dones[0])
|
||||
info_list.append(info)
|
||||
self.prev_admissible_commands[i] = info['admissible_commands']
|
||||
rewards_list.append(compute_reward(info, self.multi_modal))
|
||||
|
||||
image_obs_list = None
|
||||
return text_obs_list, image_obs_list, rewards_list, dones_list, info_list
|
||||
|
||||
def reset(self):
|
||||
for worker in self.workers:
|
||||
worker.send("reset")
|
||||
results = [worker.recv() for worker in self.workers]
|
||||
|
||||
text_obs_list = []
|
||||
info_list = []
|
||||
|
||||
for i, (obs, info) in enumerate(results):
|
||||
for k in info.keys():
|
||||
info[k] = info[k][0]
|
||||
text_obs_list.append(obs[0])
|
||||
self.prev_admissible_commands[i] = info['admissible_commands']
|
||||
info_list.append(info)
|
||||
|
||||
image_obs_list = None
|
||||
return text_obs_list, image_obs_list, info_list
|
||||
|
||||
@property
|
||||
def get_admissible_commands(self):
|
||||
return self.prev_admissible_commands
|
||||
|
||||
def close(self):
|
||||
for worker in self.workers:
|
||||
worker.close()
|
||||
|
||||
|
||||
def build_alfworld_envs(alf_config_path, seed, env_num, group_n,
|
||||
resources_per_worker, is_train=True, env_kwargs=None, gamefiles=None):
|
||||
"""Build vectorized ALFWorld environments."""
|
||||
return AlfworldEnvs(
|
||||
alf_config_path, seed, env_num, group_n,
|
||||
resources_per_worker, is_train, env_kwargs, gamefiles,
|
||||
)
|
||||
@@ -0,0 +1,60 @@
|
||||
# Vendored from SkillRL (Apache-2.0 License)
|
||||
# Original: agent_system/environments/env_package/alfworld/projection.py
|
||||
|
||||
from typing import List
|
||||
import re
|
||||
|
||||
|
||||
def alfworld_projection(actions: List[str], action_pools: List[List[str]]):
|
||||
"""Process raw model outputs into valid ALFWorld actions.
|
||||
|
||||
Extracts text from ``<action>...</action>`` tags and validates that
|
||||
the response also contains ``<think>...</think>`` tags.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
actions : list[str]
|
||||
Raw model outputs, one per environment.
|
||||
action_pools : list[list[str]]
|
||||
Admissible action lists per environment (unused but kept for API compat).
|
||||
|
||||
Returns
|
||||
-------
|
||||
actions : list[str]
|
||||
Cleaned action strings.
|
||||
valids : list[int]
|
||||
1 if the action was successfully parsed, 0 otherwise.
|
||||
"""
|
||||
valids = [0] * len(actions)
|
||||
|
||||
for i in range(len(actions)):
|
||||
original_str = actions[i]
|
||||
actions[i] = actions[i].lower()
|
||||
|
||||
start_tag = "<action>"
|
||||
end_tag = "</action>"
|
||||
start_idx = actions[i].find(start_tag)
|
||||
end_idx = actions[i].find(end_tag)
|
||||
try:
|
||||
if start_idx == -1 or end_idx == -1:
|
||||
actions[i] = actions[i][-30:]
|
||||
continue
|
||||
|
||||
extracted_action = actions[i][start_idx + len(start_tag):end_idx].strip().lower()
|
||||
actions[i] = extracted_action
|
||||
valids[i] = 1
|
||||
|
||||
except Exception:
|
||||
actions[i] = actions[i][-30:]
|
||||
|
||||
# Require <think>...</think>
|
||||
think_start_idx = original_str.find("<think>")
|
||||
think_end_idx = original_str.find("</think>")
|
||||
if think_start_idx == -1 or think_end_idx == -1:
|
||||
valids[i] = 0
|
||||
|
||||
# Reject responses containing Chinese characters
|
||||
if re.search(r'[\u4e00-\u9fff]', original_str):
|
||||
valids[i] = 0
|
||||
|
||||
return actions, valids
|
||||
@@ -0,0 +1,8 @@
|
||||
# Vendored from SkillRL (Apache-2.0 License)
|
||||
# Original: agent_system/environments/prompts/alfworld.py
|
||||
|
||||
from skillopt.prompts import load_prompt
|
||||
|
||||
ALFWORLD_TEMPLATE_NO_HIS = load_prompt("rollout_no_history", env="alfworld")
|
||||
ALFWORLD_TEMPLATE = load_prompt("rollout_with_history", env="alfworld")
|
||||
ALFWORLD_TEMPLATE_WITH_MEMORY = load_prompt("rollout_with_memory", env="alfworld")
|
||||
+145
@@ -0,0 +1,145 @@
|
||||
dataset:
|
||||
data_path: '$ALFWORLD_DATA/json_2.1.1/train'
|
||||
eval_id_data_path: '$ALFWORLD_DATA/json_2.1.1/valid_seen' # null/None to disable
|
||||
eval_ood_data_path: '$ALFWORLD_DATA/json_2.1.1/valid_unseen' # null/None to disable
|
||||
num_train_games: -1 # max training games (<=0 indicates full dataset)
|
||||
num_eval_games: -1 # max evaluation games (<=0 indicates full dataset)
|
||||
|
||||
logic:
|
||||
domain: '$ALFWORLD_DATA/logic/alfred.pddl' # PDDL domain file that defines the world dynamics
|
||||
grammar: '$ALFWORLD_DATA/logic/alfred.twl2' # Grammar file that defines the text feedbacks
|
||||
|
||||
env:
|
||||
type: 'AlfredTWEnv' # 'AlfredTWEnv' or 'AlfredThorEnv' or 'AlfredHybrid'
|
||||
# regen_game_files: False # check if game is solvable by expert and save to game.tw-pddl file
|
||||
domain_randomization: False # shuffle Textworld print order and object id nums
|
||||
task_types: [1, 2, 3, 4, 5, 6] # task-type ids: 1 - Pick & Place, 2 - Examine in Light, 3 - Clean & Place, 4 - Heat & Place, 5 - Cool & Place, 6 - Pick Two & Place
|
||||
expert_timeout_steps: 150 # max steps before timeout for expert to solve the task
|
||||
expert_type: "handcoded" # 'handcoded' or 'planner'. Note: the planner is very slow for real-time use
|
||||
goal_desc_human_anns_prob: 0.0 # prob of using human-annotated goal language instead of templated goals (1.0 indicates all human annotations from ALFRED)
|
||||
|
||||
hybrid:
|
||||
start_eps: 100000 # starting episode of hybrid training, tw-only training upto this point
|
||||
thor_prob: 0.5 # prob of AlfredThorEnv during hybrid training
|
||||
eval_mode: "tw" # 'tw' or 'thor' - env used for evaluation during hybrid training
|
||||
|
||||
thor:
|
||||
screen_width: 300 # width of THOR window
|
||||
screen_height: 300 # height of THOR window
|
||||
smooth_nav: False # smooth rotations, looks, and translations during navigation (very slow)
|
||||
save_frames_to_disk: False # save frame PNGs to disk (useful for making videos)
|
||||
save_frames_path: './videos/' # path to save frame PNGs
|
||||
|
||||
controller:
|
||||
type: 'oracle' # 'oracle' or 'oracle_astar' or 'mrcnn' or 'mrcnn_astar' (aka BUTLER)
|
||||
debug: False
|
||||
load_receps: True # load receptacle locations from precomputed dict (if available)
|
||||
|
||||
mask_rcnn:
|
||||
pretrained_model_path: '$ALFWORLD_DATA/detectors/mrcnn.pth'
|
||||
|
||||
general:
|
||||
random_seed: 42
|
||||
use_cuda: True # disable this when running on machine without cuda
|
||||
visdom: False # plot training/eval curves, run with visdom server
|
||||
task: 'alfred'
|
||||
training_method: 'dagger' # 'dqn' or 'dagger'
|
||||
save_path: './training/' # path to save pytorch models
|
||||
observation_pool_capacity: 3 # k-size queue, 0 indicates no observation
|
||||
hide_init_receptacles: False # remove initial observation containing navigable receptacles
|
||||
|
||||
training:
|
||||
batch_size: 10
|
||||
max_episode: 50000
|
||||
smoothing_eps: 0.1
|
||||
optimizer:
|
||||
learning_rate: 0.001
|
||||
clip_grad_norm: 5
|
||||
|
||||
evaluate:
|
||||
run_eval: True
|
||||
batch_size: 10
|
||||
env:
|
||||
type: "AlfredTWEnv"
|
||||
|
||||
checkpoint:
|
||||
report_frequency: 1000 # report every N episode
|
||||
experiment_tag: 'test' # name of experiment
|
||||
load_pretrained: False # during test, enable this so that the agent load your pretrained model
|
||||
load_from_tag: 'not loading anything' # name of pre-trained model to load in save_path
|
||||
|
||||
model:
|
||||
encoder_layers: 1
|
||||
decoder_layers: 1
|
||||
encoder_conv_num: 5
|
||||
block_hidden_dim: 64
|
||||
n_heads: 1
|
||||
dropout: 0.1
|
||||
block_dropout: 0.1
|
||||
recurrent: True
|
||||
|
||||
rl:
|
||||
action_space: "admissible" # 'admissible' (candidates from text engine) or 'generation' (seq2seq-style generation) or 'beam_search_choice' or 'exhaustive' (not working)
|
||||
max_target_length: 20 # max token length for seq2seq generation
|
||||
beam_width: 10 # 1 means greedy
|
||||
generate_top_k: 3
|
||||
|
||||
training:
|
||||
max_nb_steps_per_episode: 50 # terminate after this many steps
|
||||
learn_start_from_this_episode: 0 # delay updates until this epsiode
|
||||
target_net_update_frequency: 500 # sync target net with online net per this many epochs
|
||||
|
||||
replay:
|
||||
accumulate_reward_from_final: True
|
||||
count_reward_lambda: 0.0 # 0 to disable
|
||||
novel_object_reward_lambda: 0.0 # 0 to disable
|
||||
discount_gamma_game_reward: 0.9
|
||||
discount_gamma_count_reward: 0.5
|
||||
discount_gamma_novel_object_reward: 0.5
|
||||
replay_memory_capacity: 500000 # adjust this depending on your RAM size
|
||||
replay_memory_priority_fraction: 0.5
|
||||
update_per_k_game_steps: 5
|
||||
replay_batch_size: 64
|
||||
multi_step: 3
|
||||
replay_sample_history_length: 4
|
||||
replay_sample_update_from: 2
|
||||
|
||||
epsilon_greedy:
|
||||
noisy_net: False # if this is true, then epsilon greedy is disabled
|
||||
epsilon_anneal_episodes: 1000 # -1 if not annealing
|
||||
epsilon_anneal_from: 0.3
|
||||
epsilon_anneal_to: 0.1
|
||||
|
||||
dagger:
|
||||
action_space: "generation" # 'admissible' (candidates from text engine) or 'generation' (seq2seq-style generation) or 'exhaustive' (not working)
|
||||
max_target_length: 20 # max token length for seq2seq generation
|
||||
beam_width: 10 # 1 means greedy
|
||||
generate_top_k: 5
|
||||
unstick_by_beam_search: False # use beam-search for failed actions, set True during evaluation
|
||||
|
||||
training:
|
||||
max_nb_steps_per_episode: 50 # terminate after this many steps
|
||||
|
||||
fraction_assist:
|
||||
fraction_assist_anneal_episodes: 50000
|
||||
fraction_assist_anneal_from: 1.0
|
||||
fraction_assist_anneal_to: 0.01
|
||||
|
||||
fraction_random:
|
||||
fraction_random_anneal_episodes: 0
|
||||
fraction_random_anneal_from: 0.0
|
||||
fraction_random_anneal_to: 0.0
|
||||
|
||||
replay:
|
||||
replay_memory_capacity: 500000
|
||||
update_per_k_game_steps: 5
|
||||
replay_batch_size: 64
|
||||
replay_sample_history_length: 4
|
||||
replay_sample_update_from: 2
|
||||
|
||||
vision_dagger:
|
||||
model_type: "resnet" # 'resnet' (whole image features) or 'maskrcnn_whole' (whole image MaskRCNN feats) or 'maskrcnn' (top k MaskRCNN detection feats) or 'no_vision' (zero vision input)
|
||||
resnet_fc_dim: 64
|
||||
maskrcnn_top_k_boxes: 10 # top k box features
|
||||
use_exploration_frame_feats: False # append feats from initial exploration (memory intensive!)
|
||||
sequence_aggregation_method: "average" # 'sum' or 'average' or 'rnn'
|
||||
+84
@@ -0,0 +1,84 @@
|
||||
# Vendored from SkillRL (Apache-2.0 License)
|
||||
# Original: agent_system/environments/base.py
|
||||
# Trimmed to only include what ALFWorld needs.
|
||||
|
||||
from typing import List, Tuple, Dict, Any
|
||||
import numpy as np
|
||||
from collections import defaultdict
|
||||
|
||||
|
||||
def to_numpy(data):
|
||||
"""Convert data to numpy array."""
|
||||
# Lazy-check for torch.Tensor to avoid hard dependency on torch
|
||||
_torch_tensor = None
|
||||
try:
|
||||
import torch
|
||||
_torch_tensor = torch.Tensor
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
if _torch_tensor is not None and isinstance(data, _torch_tensor):
|
||||
data = data.detach().cpu().numpy()
|
||||
elif isinstance(data, np.ndarray):
|
||||
pass
|
||||
elif isinstance(data, (int, float, bool, Tuple, List)):
|
||||
data = np.array(data)
|
||||
else:
|
||||
raise ValueError(f"Unsupported type: {type(data)})")
|
||||
return data
|
||||
|
||||
|
||||
class EnvironmentManagerBase:
|
||||
"""Base class for vectorized environment managers.
|
||||
|
||||
Manages a set of parallel environments, handles action projection,
|
||||
observation post-processing, and history tracking.
|
||||
"""
|
||||
|
||||
def __init__(self, envs, projection_f, config):
|
||||
self.envs = envs
|
||||
self.projection_f = projection_f
|
||||
self.config = config
|
||||
|
||||
def reset(self, kwargs) -> Dict[str, Any]:
|
||||
obs, infos = self.envs.reset()
|
||||
return {'text': None, 'image': obs, 'anchor': None}, infos
|
||||
|
||||
def step(self, text_actions: List[str]):
|
||||
actions, valids = self.projection_f(text_actions)
|
||||
next_obs, rewards, dones, infos = self.envs.step(actions)
|
||||
|
||||
next_observations = {
|
||||
'text': None,
|
||||
'image': next_obs,
|
||||
'anchor': None,
|
||||
}
|
||||
for i, info in enumerate(infos):
|
||||
info['is_action_valid'] = to_numpy(valids[i])
|
||||
|
||||
rewards = to_numpy(rewards)
|
||||
dones = to_numpy(dones)
|
||||
return next_observations, rewards, dones, infos
|
||||
|
||||
def close(self) -> None:
|
||||
self.envs.close()
|
||||
|
||||
def success_evaluator(self, *args, **kwargs) -> Dict[str, np.ndarray]:
|
||||
total_infos = kwargs['total_infos']
|
||||
total_batch_list = kwargs['total_batch_list']
|
||||
batch_size = len(total_batch_list)
|
||||
|
||||
success = defaultdict(list)
|
||||
for bs in range(batch_size):
|
||||
self._process_batch(bs, total_batch_list, total_infos, success)
|
||||
assert len(success['success_rate']) == batch_size
|
||||
return {key: np.array(value) for key, value in success.items()}
|
||||
|
||||
def _process_batch(self, batch_idx, total_batch_list, total_infos, success):
|
||||
for i in reversed(range(len(total_batch_list[batch_idx]))):
|
||||
batch_item = total_batch_list[batch_idx][i]
|
||||
if batch_item['active_masks']:
|
||||
info = total_infos[batch_idx][i]
|
||||
won_value = float(info['won'])
|
||||
success['success_rate'].append(won_value)
|
||||
return
|
||||
+139
@@ -0,0 +1,139 @@
|
||||
# Vendored from SkillRL (Apache-2.0 License)
|
||||
# Original: agent_system/environments/env_manager.py
|
||||
# Trimmed to only include AlfWorldEnvironmentManager and its helpers.
|
||||
|
||||
from typing import List, Dict, Any
|
||||
from collections import defaultdict
|
||||
import numpy as np
|
||||
|
||||
from skillopt.envs.alfworld.vendor.env_base import EnvironmentManagerBase, to_numpy
|
||||
from skillopt.envs.alfworld.vendor.alfworld_prompts import (
|
||||
ALFWORLD_TEMPLATE,
|
||||
ALFWORLD_TEMPLATE_NO_HIS,
|
||||
ALFWORLD_TEMPLATE_WITH_MEMORY,
|
||||
)
|
||||
from skillopt.envs.alfworld.vendor.memory import SimpleMemory
|
||||
|
||||
|
||||
def parse_gamefile(infos):
|
||||
gamefile = []
|
||||
for info in infos:
|
||||
if 'extra.gamefile' in info:
|
||||
gamefile.append(info['extra.gamefile'])
|
||||
else:
|
||||
gamefile.append(None)
|
||||
return gamefile
|
||||
|
||||
|
||||
def set_gamefile(infos, gamefile):
|
||||
for i in range(len(infos)):
|
||||
if 'extra.gamefile' in infos[i]:
|
||||
infos[i]['extra.gamefile'] = gamefile[i]
|
||||
else:
|
||||
infos[i]['extra.gamefile'] = None
|
||||
return infos
|
||||
|
||||
|
||||
class AlfWorldEnvironmentManager(EnvironmentManagerBase):
|
||||
"""Manages parallel ALFWorld environments with observation templating."""
|
||||
|
||||
def __init__(self, envs, projection_f, config):
|
||||
self.memory = SimpleMemory()
|
||||
self.retrieval_memory = None
|
||||
super().__init__(envs, projection_f, config)
|
||||
|
||||
def reset(self, kwargs):
|
||||
text_obs, image_obs, infos = self.envs.reset()
|
||||
self.gamefile = parse_gamefile(infos)
|
||||
self.memory.reset(batch_size=len(text_obs))
|
||||
self.tasks = []
|
||||
self.pre_text_obs = text_obs
|
||||
self.extract_task(text_obs)
|
||||
|
||||
full_text_obs = self.build_text_obs(text_obs, self.envs.get_admissible_commands, init=True)
|
||||
return {'text': full_text_obs, 'image': image_obs, 'anchor': text_obs}, infos
|
||||
|
||||
def step(self, text_actions: List[str]):
|
||||
actions, valids = self.projection_f(text_actions, self.envs.get_admissible_commands)
|
||||
text_obs, image_obs, rewards, dones, infos = self.envs.step(actions)
|
||||
self.memory.store({'text_obs': self.pre_text_obs, 'action': actions})
|
||||
self.pre_text_obs = text_obs
|
||||
|
||||
full_text_obs = self.build_text_obs(text_obs, self.envs.get_admissible_commands)
|
||||
if infos[0].get("extra.gamefile") is None:
|
||||
infos = set_gamefile(infos, self.gamefile)
|
||||
|
||||
for i, info in enumerate(infos):
|
||||
info['is_action_valid'] = to_numpy(valids[i])
|
||||
|
||||
next_observations = {'text': full_text_obs, 'image': image_obs, 'anchor': text_obs}
|
||||
rewards = to_numpy(rewards)
|
||||
dones = to_numpy(dones)
|
||||
return next_observations, rewards, dones, infos
|
||||
|
||||
def extract_task(self, text_obs: List[str]):
|
||||
for obs in text_obs:
|
||||
task_start = obs.find('Your task is to: ')
|
||||
if task_start != -1:
|
||||
self.tasks.append(obs[task_start + len('Your task is to: '):].strip())
|
||||
else:
|
||||
raise ValueError("Task description not found in text observation.")
|
||||
|
||||
def build_text_obs(self, text_obs: List[str], admissible_actions: List[List[str]], init: bool = False) -> List[str]:
|
||||
postprocess_text_obs = []
|
||||
if not init and self.config.env.history_length > 0:
|
||||
memory_contexts, valid_lens = self.memory.fetch(
|
||||
self.config.env.history_length,
|
||||
obs_key="text_obs",
|
||||
action_key="action",
|
||||
)
|
||||
|
||||
for i in range(len(text_obs)):
|
||||
reformatted_admissible_actions = "\n ".join(
|
||||
f"'{s}'" for s in admissible_actions[i] if s != 'help'
|
||||
)
|
||||
|
||||
if init or self.config.env.history_length <= 0:
|
||||
obs = ALFWORLD_TEMPLATE_NO_HIS.format(
|
||||
current_observation=text_obs[i],
|
||||
admissible_actions=reformatted_admissible_actions,
|
||||
)
|
||||
else:
|
||||
obs = ALFWORLD_TEMPLATE.format(
|
||||
task_description=self.tasks[i],
|
||||
step_count=len(self.memory[i]),
|
||||
history_length=valid_lens[i],
|
||||
action_history=memory_contexts[i],
|
||||
current_step=len(self.memory[i]) + 1,
|
||||
current_observation=text_obs[i],
|
||||
admissible_actions=reformatted_admissible_actions,
|
||||
)
|
||||
postprocess_text_obs.append(obs)
|
||||
return postprocess_text_obs
|
||||
|
||||
def _process_batch(self, batch_idx, total_batch_list, total_infos, success):
|
||||
for i in reversed(range(len(total_batch_list[batch_idx]))):
|
||||
batch_item = total_batch_list[batch_idx][i]
|
||||
if batch_item['active_masks']:
|
||||
info = total_infos[batch_idx][i]
|
||||
won_value = float(info['won'])
|
||||
success['success_rate'].append(won_value)
|
||||
|
||||
gamefile = info.get("extra.gamefile")
|
||||
if gamefile:
|
||||
self._process_gamefile(gamefile, won_value, success)
|
||||
return
|
||||
|
||||
def _process_gamefile(self, gamefile, won_value, success):
|
||||
tasks = [
|
||||
"pick_and_place",
|
||||
"pick_two_obj_and_place",
|
||||
"look_at_obj_in_light",
|
||||
"pick_heat_then_place_in_recep",
|
||||
"pick_cool_then_place_in_recep",
|
||||
"pick_clean_then_place_in_recep",
|
||||
]
|
||||
for task in tasks:
|
||||
if task in gamefile:
|
||||
success[f"{task}_success_rate"].append(won_value)
|
||||
break
|
||||
+87
@@ -0,0 +1,87 @@
|
||||
# Vendored from SkillRL (Apache-2.0 License)
|
||||
# Original: agent_system/memory/base.py + agent_system/memory/memory.py
|
||||
# Merged into a single file for simplicity.
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Dict, Any, Tuple
|
||||
|
||||
|
||||
class BaseMemory(ABC):
|
||||
"""Base class for memory management."""
|
||||
|
||||
@abstractmethod
|
||||
def __len__(self):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def __getitem__(self, idx: int):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def reset(self, batch_size: int):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def store(self, record: Dict[str, List[Any]]):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def fetch(self, step: int):
|
||||
pass
|
||||
|
||||
|
||||
class SimpleMemory(BaseMemory):
|
||||
"""Per-environment history buffer for storing observations and actions."""
|
||||
|
||||
def __init__(self):
|
||||
self._data = None
|
||||
self.keys = None
|
||||
self.batch_size = 0
|
||||
|
||||
def __len__(self):
|
||||
return len(self._data)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return self._data[idx]
|
||||
|
||||
def reset(self, batch_size: int):
|
||||
if self._data is not None:
|
||||
self._data.clear()
|
||||
self._data = [[] for _ in range(batch_size)]
|
||||
self.batch_size = batch_size
|
||||
self.keys = None
|
||||
|
||||
def store(self, record: Dict[str, List[Any]]):
|
||||
if self.keys is None:
|
||||
self.keys = list(record.keys())
|
||||
assert self.keys == list(record.keys())
|
||||
|
||||
for env_idx in range(self.batch_size):
|
||||
self._data[env_idx].append({k: record[k][env_idx] for k in self.keys})
|
||||
|
||||
def fetch(
|
||||
self,
|
||||
history_length: int,
|
||||
obs_key: str = "text_obs",
|
||||
action_key: str = "action",
|
||||
) -> Tuple[List[str], List[int]]:
|
||||
memory_contexts, valid_lengths = [], []
|
||||
|
||||
for env_idx in range(self.batch_size):
|
||||
recent = self._data[env_idx][-history_length:]
|
||||
valid_len = len(recent)
|
||||
start_idx = len(self._data[env_idx]) - valid_len
|
||||
|
||||
lines = []
|
||||
for j, rec in enumerate(recent):
|
||||
step_num = start_idx + j + 1
|
||||
act = rec[action_key]
|
||||
obs = rec[obs_key]
|
||||
lines.append(
|
||||
f"[Observation {step_num}: '{obs}', Action {step_num}: '{act}']"
|
||||
)
|
||||
|
||||
memory_contexts.append("\n".join(lines))
|
||||
valid_lengths.append(valid_len)
|
||||
|
||||
return memory_contexts, valid_lengths
|
||||
@@ -0,0 +1 @@
|
||||
"""BabyVision environment package for ReflACT."""
|
||||
@@ -0,0 +1,267 @@
|
||||
"""BabyVision environment adapter for ReflACT."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
from skillopt.gradient.deep_probe import generate_deep_probe_instruction
|
||||
from skillopt.datasets.base import BatchSpec
|
||||
from skillopt.gradient.reflect import run_minibatch_reflect
|
||||
from skillopt.envs.base import EnvAdapter
|
||||
from skillopt.envs.babyvision.dataloader import BabyVisionDataLoader
|
||||
from skillopt.envs.babyvision.rollout import run_batch
|
||||
from skillopt.model import get_student_backend
|
||||
|
||||
|
||||
class BabyVisionAdapter(EnvAdapter):
|
||||
"""BabyVision adapter."""
|
||||
|
||||
def build_reference_text(self, item: dict) -> str:
|
||||
cot = str(item.get("cot") or "").strip()
|
||||
if not cot:
|
||||
return ""
|
||||
return f"## Reference CoT\n{cot}"
|
||||
|
||||
def get_reference_metadata(self, item: dict) -> dict:
|
||||
cot = str(item.get("cot") or "").strip()
|
||||
if not cot:
|
||||
return {"fields": [], "preview": ""}
|
||||
return {
|
||||
"fields": ["cot"],
|
||||
"preview": cot[:400],
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
split_dir: str = "",
|
||||
data_path: str = "",
|
||||
split_mode: str = "ratio",
|
||||
split_ratio: str = "2:1:7",
|
||||
split_seed: int = 42,
|
||||
split_output_dir: str = "",
|
||||
max_turns: int = 1,
|
||||
workers: int = 32,
|
||||
analyst_workers: int = 16,
|
||||
failure_only: bool = False,
|
||||
minibatch_size: int = 8,
|
||||
edit_budget: int = 4,
|
||||
seed: int = 42,
|
||||
limit: int = 0,
|
||||
image_detail: str = "auto",
|
||||
judge_model: str = "gpt-5.4",
|
||||
judge_max_completion_tokens: int = 256,
|
||||
judge_retries: int = 5,
|
||||
use_deep_reflect: bool = False,
|
||||
deep_reflect_failures: int = 4,
|
||||
deep_reflect_successes: int = 2,
|
||||
) -> None:
|
||||
self.max_turns = max_turns
|
||||
self.workers = workers
|
||||
self.analyst_workers = analyst_workers
|
||||
self.failure_only = failure_only
|
||||
self.minibatch_size = minibatch_size
|
||||
self.edit_budget = edit_budget
|
||||
self.image_detail = image_detail
|
||||
self.judge_model = judge_model
|
||||
self.judge_max_completion_tokens = judge_max_completion_tokens
|
||||
self.judge_retries = judge_retries
|
||||
self.use_deep_reflect = use_deep_reflect
|
||||
self.deep_reflect_failures = deep_reflect_failures
|
||||
self.deep_reflect_successes = deep_reflect_successes
|
||||
self.dataloader = BabyVisionDataLoader(
|
||||
split_dir=split_dir,
|
||||
data_path=data_path,
|
||||
split_mode=split_mode,
|
||||
split_ratio=split_ratio,
|
||||
split_seed=split_seed,
|
||||
split_output_dir=split_output_dir,
|
||||
seed=seed,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
def setup(self, cfg: dict) -> None:
|
||||
super().setup(cfg)
|
||||
self.dataloader.setup(cfg)
|
||||
|
||||
def get_dataloader(self):
|
||||
return self.dataloader
|
||||
|
||||
def build_env_from_batch(self, batch: BatchSpec, **kwargs):
|
||||
return list(batch.payload or [])
|
||||
|
||||
def build_train_env(self, batch_size: int, seed: int, **kwargs):
|
||||
batch = self.dataloader.build_train_batch(batch_size=batch_size, seed=seed, **kwargs)
|
||||
return self.build_env_from_batch(batch, **kwargs)
|
||||
|
||||
def build_eval_env(self, env_num: int, split: str, seed: int, **kwargs):
|
||||
batch = self.dataloader.build_eval_batch(env_num=env_num, split=split, seed=seed, **kwargs)
|
||||
return self.build_env_from_batch(batch, **kwargs)
|
||||
|
||||
def rollout(
|
||||
self,
|
||||
env_manager,
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict]:
|
||||
items: list[dict] = env_manager
|
||||
return run_batch(
|
||||
items=items,
|
||||
out_root=out_dir,
|
||||
skill_content=skill_content,
|
||||
max_turns=self.max_turns,
|
||||
workers=self.workers,
|
||||
image_detail=self.image_detail,
|
||||
judge_model=self.judge_model,
|
||||
judge_max_completion_tokens=self.judge_max_completion_tokens,
|
||||
judge_retries=self.judge_retries,
|
||||
diagnostic_mode=kwargs.get("diagnostic_mode", False),
|
||||
diagnostic_instruction=kwargs.get("diagnostic_instruction", ""),
|
||||
diagnostic_trace_context_by_id=kwargs.get("diagnostic_trace_context_by_id"),
|
||||
)
|
||||
|
||||
def reflect(
|
||||
self,
|
||||
results: list[dict],
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict | None]:
|
||||
prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
|
||||
patches_dir = kwargs.get("patches_dir", os.path.join(out_dir, "patches"))
|
||||
random_seed = kwargs.get("random_seed")
|
||||
step_buffer_context = kwargs.get("step_buffer_context", "")
|
||||
meta_skill_context = kwargs.get("meta_skill_context", "")
|
||||
|
||||
return run_minibatch_reflect(
|
||||
results=results,
|
||||
skill_content=skill_content,
|
||||
prediction_dir=prediction_dir,
|
||||
patches_dir=patches_dir,
|
||||
workers=self.analyst_workers,
|
||||
failure_only=self.failure_only,
|
||||
minibatch_size=self.minibatch_size,
|
||||
edit_budget=self.edit_budget,
|
||||
random_seed=random_seed,
|
||||
error_system=self.get_error_minibatch_prompt(),
|
||||
success_system=self.get_success_minibatch_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
meta_skill_context=meta_skill_context,
|
||||
update_mode=getattr(self, "_cfg", {}).get("skill_update_mode", "patch"),
|
||||
)
|
||||
|
||||
def deep_reflect(
|
||||
self,
|
||||
results: list[dict],
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict | None]:
|
||||
if not self.use_deep_reflect:
|
||||
return []
|
||||
|
||||
env_manager = kwargs.get("env_manager")
|
||||
prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
|
||||
random_seed = kwargs.get("random_seed")
|
||||
step_buffer_context = kwargs.get("step_buffer_context", "")
|
||||
meta_skill_context = kwargs.get("meta_skill_context", "")
|
||||
codex_backend = get_student_backend() == "codex_exec"
|
||||
selected_items = self.select_representative_items(
|
||||
results,
|
||||
env_manager if isinstance(env_manager, list) else None,
|
||||
n_failures=self.deep_reflect_failures,
|
||||
n_successes=self.deep_reflect_successes,
|
||||
seed=random_seed,
|
||||
)
|
||||
if not selected_items:
|
||||
return []
|
||||
selected_ids = {str(item["id"]) for item in selected_items}
|
||||
selected_results = [row for row in results if str(row.get("id")) in selected_ids]
|
||||
selected_examples = self.attach_reference_context(selected_results, selected_items)
|
||||
if codex_backend:
|
||||
selected_examples = self.attach_codex_probe_context(selected_examples, prediction_dir)
|
||||
selected_metadata = []
|
||||
cot_count = 0
|
||||
for item in selected_items:
|
||||
meta = self.get_reference_metadata(item)
|
||||
if meta["fields"]:
|
||||
cot_count += 1
|
||||
selected_metadata.append({
|
||||
"id": str(item["id"]),
|
||||
"task_type": str(item.get("subtype") or item.get("task_type") or "babyvision"),
|
||||
"reference_fields": meta["fields"],
|
||||
"reference_preview": meta["preview"],
|
||||
})
|
||||
|
||||
deep_dir = os.path.join(out_dir, "deep_reflect")
|
||||
rollout_dir = os.path.join(deep_dir, "rollout")
|
||||
patches_dir = os.path.join(deep_dir, "patches")
|
||||
os.makedirs(deep_dir, exist_ok=True)
|
||||
print(
|
||||
f" [2b/6 DEEP REFLECT setup] selected={len(selected_items)} "
|
||||
f"reference_fields=cot({cot_count}/{len(selected_items)})"
|
||||
)
|
||||
probe = generate_deep_probe_instruction(
|
||||
skill_content=skill_content,
|
||||
items=selected_examples,
|
||||
prediction_dir=prediction_dir,
|
||||
system_prompt=self.get_codex_deep_probe_prompt() if codex_backend else self.get_deep_probe_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
meta_skill_context=meta_skill_context,
|
||||
)
|
||||
if not probe:
|
||||
return []
|
||||
diagnostic_trace_context_by_id = None
|
||||
if codex_backend:
|
||||
selected_items, diagnostic_trace_context_by_id, probe = self.resolve_codex_probe_target(
|
||||
selected_items=selected_items,
|
||||
selected_examples=selected_examples,
|
||||
prediction_dir=prediction_dir,
|
||||
probe=probe,
|
||||
)
|
||||
probe_record = {
|
||||
**probe,
|
||||
"reference_summary": {
|
||||
"selected_count": len(selected_items),
|
||||
"field_counts": {
|
||||
"cot": cot_count,
|
||||
},
|
||||
},
|
||||
"selected_examples": selected_metadata,
|
||||
}
|
||||
with open(os.path.join(deep_dir, "probe.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(probe_record, f, ensure_ascii=False, indent=2)
|
||||
deep_results = run_batch(
|
||||
items=selected_items,
|
||||
out_root=rollout_dir,
|
||||
skill_content=skill_content,
|
||||
max_turns=self.max_turns,
|
||||
workers=min(self.workers, max(len(selected_items), 1)),
|
||||
image_detail=self.image_detail,
|
||||
judge_model=self.judge_model,
|
||||
judge_max_completion_tokens=self.judge_max_completion_tokens,
|
||||
judge_retries=self.judge_retries,
|
||||
diagnostic_mode=True,
|
||||
diagnostic_instruction=probe["probe_instruction"],
|
||||
diagnostic_trace_context_by_id=diagnostic_trace_context_by_id,
|
||||
)
|
||||
deep_results = self.attach_reference_context(deep_results, selected_items)
|
||||
return run_minibatch_reflect(
|
||||
results=deep_results,
|
||||
skill_content=skill_content,
|
||||
prediction_dir=os.path.join(rollout_dir, "predictions"),
|
||||
patches_dir=patches_dir,
|
||||
workers=self.analyst_workers,
|
||||
failure_only=self.failure_only,
|
||||
minibatch_size=self.minibatch_size,
|
||||
edit_budget=self.edit_budget,
|
||||
random_seed=random_seed,
|
||||
error_system=self.get_error_minibatch_prompt(),
|
||||
success_system=self.get_success_minibatch_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
meta_skill_context=meta_skill_context,
|
||||
update_mode=getattr(self, "_cfg", {}).get("skill_update_mode", "patch"),
|
||||
)
|
||||
|
||||
def get_task_types(self) -> list[str]:
|
||||
return self.dataloader.get_task_types()
|
||||
@@ -0,0 +1,214 @@
|
||||
"""BabyVision task dataloader."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
from skillopt.datasets.base import SplitDataLoader
|
||||
|
||||
|
||||
# ── Raw data loading utilities (for preprocessing / standalone eval) ─────
|
||||
|
||||
_CHOICE_LABELS = ["A", "B", "C", "D", "E", "F", "G"]
|
||||
|
||||
|
||||
def _iter_jsonl(path: str) -> list[dict]:
|
||||
items: list[dict] = []
|
||||
with open(path, encoding="utf-8") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
items.append(json.loads(line))
|
||||
return items
|
||||
|
||||
|
||||
def _normalize_ans_type(raw: Any, options: list[dict], choice_answer: Any) -> str:
|
||||
text = str(raw or "").strip().lower()
|
||||
if text in {"choice", "multiple_choice", "mcq", "option"}:
|
||||
return "choice"
|
||||
if text in {"blank", "open", "open_ended", "fill_blank", "short_answer"}:
|
||||
return "blank"
|
||||
if options or choice_answer not in (None, "", []):
|
||||
return "choice"
|
||||
return "blank"
|
||||
|
||||
|
||||
def _coerce_options(raw: Any) -> list[dict]:
|
||||
options: list[dict] = []
|
||||
if isinstance(raw, list):
|
||||
for idx, item in enumerate(raw):
|
||||
if isinstance(item, dict):
|
||||
text = str(item.get("text") or item.get("content") or item.get("option") or "").strip()
|
||||
label = str(item.get("label") or _CHOICE_LABELS[idx]).strip()
|
||||
else:
|
||||
text = str(item).strip()
|
||||
label = _CHOICE_LABELS[idx]
|
||||
if text:
|
||||
options.append({"label": label, "text": text})
|
||||
elif isinstance(raw, dict):
|
||||
for idx, (key, value) in enumerate(raw.items()):
|
||||
text = str(value).strip()
|
||||
if text:
|
||||
options.append({"label": str(key).strip() or _CHOICE_LABELS[idx], "text": text})
|
||||
return options
|
||||
|
||||
|
||||
def _normalize_choice_answer(choice_answer: Any, options: list[dict]) -> dict[str, str]:
|
||||
if not options:
|
||||
return {"label": "", "text": ""}
|
||||
|
||||
if isinstance(choice_answer, dict):
|
||||
label = str(choice_answer.get("label") or "").strip().upper()
|
||||
text = str(choice_answer.get("text") or "").strip()
|
||||
for option in options:
|
||||
if label and option["label"].strip().upper() == label:
|
||||
return {"label": option["label"], "text": option["text"]}
|
||||
if text and option["text"] == text:
|
||||
return {"label": option["label"], "text": option["text"]}
|
||||
|
||||
if isinstance(choice_answer, int):
|
||||
idx = choice_answer
|
||||
if 0 <= idx < len(options):
|
||||
return dict(options[idx])
|
||||
if 1 <= idx <= len(options):
|
||||
return dict(options[idx - 1])
|
||||
|
||||
text = str(choice_answer or "").strip()
|
||||
label = text.upper().rstrip(".):")
|
||||
for option in options:
|
||||
if option["label"].strip().upper() == label:
|
||||
return dict(option)
|
||||
if option["text"] == text:
|
||||
return dict(option)
|
||||
|
||||
return {"label": "", "text": ""}
|
||||
|
||||
|
||||
def _coerce_blank_answers(raw: Any) -> list[str]:
|
||||
if isinstance(raw, list):
|
||||
return [str(item).strip() for item in raw if str(item).strip()]
|
||||
if raw is None:
|
||||
return []
|
||||
text = str(raw).strip()
|
||||
return [text] if text else []
|
||||
|
||||
|
||||
def load_items(data_path: str) -> list[dict]:
|
||||
"""Load and normalise BabyVision items from a directory or JSONL file."""
|
||||
if not data_path:
|
||||
raise ValueError("BabyVision requires data_path pointing to a local dataset directory or meta_data.jsonl.")
|
||||
|
||||
if os.path.isdir(data_path):
|
||||
meta_path = os.path.join(data_path, "meta_data.jsonl")
|
||||
image_root = os.path.join(data_path, "images")
|
||||
else:
|
||||
meta_path = data_path
|
||||
image_root = os.path.join(os.path.dirname(data_path), "images")
|
||||
|
||||
if not os.path.exists(meta_path):
|
||||
raise ValueError(
|
||||
"BabyVision expected a meta_data.jsonl file. "
|
||||
f"Could not find: {meta_path}"
|
||||
)
|
||||
|
||||
raw_items = _iter_jsonl(meta_path)
|
||||
items: list[dict] = []
|
||||
for idx, raw in enumerate(raw_items):
|
||||
options = _coerce_options(raw.get("options") or raw.get("choices") or raw.get("choiceOptions"))
|
||||
ans_type = _normalize_ans_type(raw.get("ansType"), options, raw.get("choiceAns"))
|
||||
correct_choice = _normalize_choice_answer(raw.get("choiceAns"), options)
|
||||
blank_answers = _coerce_blank_answers(raw.get("blankAns"))
|
||||
|
||||
image_name = str(
|
||||
raw.get("image")
|
||||
or raw.get("image_path")
|
||||
or raw.get("image_file")
|
||||
or raw.get("img")
|
||||
or ""
|
||||
).strip()
|
||||
if not image_name:
|
||||
continue
|
||||
image_path = image_name if os.path.isabs(image_name) else os.path.join(image_root, image_name)
|
||||
if not os.path.exists(image_path):
|
||||
alt = os.path.join(os.path.dirname(meta_path), image_name)
|
||||
if os.path.exists(alt):
|
||||
image_path = alt
|
||||
else:
|
||||
continue
|
||||
|
||||
task_id = str(raw.get("taskId") or raw.get("id") or idx + 1)
|
||||
task_type = str(raw.get("type") or raw.get("taskType") or "unknown").strip() or "unknown"
|
||||
subtype = str(raw.get("subtype") or raw.get("subType") or task_type).strip() or task_type
|
||||
question = str(raw.get("question") or raw.get("query") or "").strip()
|
||||
if not question:
|
||||
continue
|
||||
|
||||
if ans_type == "choice" and not correct_choice["label"]:
|
||||
continue
|
||||
if ans_type != "choice" and not blank_answers:
|
||||
continue
|
||||
|
||||
items.append({
|
||||
"id": task_id,
|
||||
"task_type": task_type,
|
||||
"subtype": subtype,
|
||||
"question": question,
|
||||
"image_path": os.path.abspath(image_path),
|
||||
"ans_type": ans_type,
|
||||
"choices": options,
|
||||
"correct_choice": correct_choice,
|
||||
"blank_answers": blank_answers,
|
||||
"cot": str(raw.get("coT") or raw.get("cot") or "").strip(),
|
||||
"source_path": os.path.abspath(meta_path),
|
||||
})
|
||||
|
||||
if not items:
|
||||
raise ValueError(f"No valid BabyVision items loaded from {data_path}")
|
||||
return items
|
||||
|
||||
|
||||
# ── Dataloader ───────────────────────────────────────────────────────────
|
||||
|
||||
class BabyVisionDataLoader(SplitDataLoader):
|
||||
"""BabyVision dataloader."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
split_dir: str = "",
|
||||
data_path: str = "",
|
||||
split_mode: str = "ratio",
|
||||
split_ratio: str = "2:1:7",
|
||||
split_seed: int = 42,
|
||||
split_output_dir: str = "",
|
||||
seed: int = 42,
|
||||
limit: int = 0,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
split_dir=split_dir,
|
||||
data_path=data_path,
|
||||
split_mode=split_mode,
|
||||
split_ratio=split_ratio,
|
||||
split_seed=split_seed,
|
||||
split_output_dir=split_output_dir,
|
||||
seed=seed,
|
||||
limit=limit,
|
||||
)
|
||||
self._task_types: list[str] = []
|
||||
|
||||
def load_raw_items(self, data_path: str) -> list[dict]:
|
||||
return load_items(data_path)
|
||||
|
||||
def setup(self, cfg: dict) -> None:
|
||||
super().setup(cfg)
|
||||
all_items = self.train_items + self.val_items + self.test_items
|
||||
task_types = {
|
||||
item.get("subtype") or item.get("task_type") or "unknown"
|
||||
for item in all_items
|
||||
}
|
||||
self._task_types = sorted(task_types)
|
||||
|
||||
def get_task_types(self) -> list[str]:
|
||||
return list(self._task_types)
|
||||
@@ -0,0 +1,160 @@
|
||||
"""BabyVision evaluation helpers using the official-style LLM judge."""
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import string
|
||||
|
||||
import regex
|
||||
|
||||
from skillopt.model import chat_with_deployment
|
||||
from skillopt.prompts import load_prompt
|
||||
|
||||
_EVAL_MODE = "babyvision_judge_v2_official_style"
|
||||
|
||||
def normalize_text(text: str) -> str:
|
||||
text = str(text).strip().lower()
|
||||
text = "".join(ch for ch in text if ch not in string.punctuation)
|
||||
return " ".join(text.split())
|
||||
|
||||
|
||||
def extract_boxed_answer(text: str | None) -> str | None:
|
||||
"""Extract the final answer using the official BabyVision rule."""
|
||||
if text is None:
|
||||
return None
|
||||
|
||||
pattern = r'\\boxed\{((?:[^{}]|{(?:[^{}]|{.*})*})*)\}'
|
||||
matches = regex.findall(pattern, text)
|
||||
if matches:
|
||||
return matches[-1]
|
||||
|
||||
pattern_alt = r'<\|begin_of_box\|>(.*?)<\|end_of_box\|>'
|
||||
matches_alt = regex.findall(pattern_alt, text)
|
||||
if matches_alt:
|
||||
return matches_alt[-1].strip()
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _token_f1(prediction: str, gold: str) -> float:
|
||||
pred_tokens = normalize_text(prediction).split()
|
||||
gold_tokens = normalize_text(gold).split()
|
||||
if not pred_tokens and not gold_tokens:
|
||||
return 1.0
|
||||
if not pred_tokens or not gold_tokens:
|
||||
return 0.0
|
||||
pred_set = {}
|
||||
gold_set = {}
|
||||
for tok in pred_tokens:
|
||||
pred_set[tok] = pred_set.get(tok, 0) + 1
|
||||
for tok in gold_tokens:
|
||||
gold_set[tok] = gold_set.get(tok, 0) + 1
|
||||
common = 0
|
||||
for tok, count in pred_set.items():
|
||||
common += min(count, gold_set.get(tok, 0))
|
||||
if common == 0:
|
||||
return 0.0
|
||||
precision = common / len(pred_tokens)
|
||||
recall = common / len(gold_tokens)
|
||||
return 2 * precision * recall / (precision + recall)
|
||||
|
||||
|
||||
def _format_choices(choices: list[dict]) -> str:
|
||||
return "\n".join(f"{choice['label']}. {choice['text']}" for choice in choices)
|
||||
|
||||
|
||||
def _judge_answer(
|
||||
*,
|
||||
item: dict,
|
||||
prediction_text: str,
|
||||
extracted_answer: str,
|
||||
judge_model: str,
|
||||
max_completion_tokens: int,
|
||||
retries: int,
|
||||
) -> dict:
|
||||
if item["ans_type"] == "choice":
|
||||
ground_truth = str(item["correct_choice"]["label"])
|
||||
else:
|
||||
if len(item["blank_answers"]) == 1:
|
||||
ground_truth = item["blank_answers"][0]
|
||||
else:
|
||||
ground_truth = " | ".join(item["blank_answers"])
|
||||
|
||||
question = str(item["question"])
|
||||
if item["ans_type"] == "choice" and item.get("choices"):
|
||||
question = f"{question}\nChoices:\n{_format_choices(item['choices'])}"
|
||||
|
||||
raw, _ = chat_with_deployment(
|
||||
deployment=judge_model,
|
||||
system="You are a careful and strict evaluator.",
|
||||
user=load_prompt("judge", env="babyvision").format(
|
||||
question=question,
|
||||
groundtruth=ground_truth,
|
||||
modeloutput=extracted_answer,
|
||||
),
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
retries=retries,
|
||||
stage="babyvision_judge",
|
||||
)
|
||||
judge_response_clean = str(raw).strip().lower()
|
||||
if "true" in judge_response_clean:
|
||||
correct = True
|
||||
elif "false" in judge_response_clean:
|
||||
correct = False
|
||||
else:
|
||||
correct = False
|
||||
return {
|
||||
"raw": raw,
|
||||
"correct": correct,
|
||||
"reason": judge_response_clean,
|
||||
"matched_gold": ground_truth if correct else "",
|
||||
}
|
||||
|
||||
|
||||
def evaluate_item(
|
||||
*,
|
||||
item: dict,
|
||||
prediction_text: str,
|
||||
judge_model: str,
|
||||
max_completion_tokens: int = 256,
|
||||
retries: int = 5,
|
||||
) -> dict:
|
||||
answer = extract_boxed_answer(prediction_text)
|
||||
judge = _judge_answer(
|
||||
item=item,
|
||||
prediction_text=prediction_text,
|
||||
extracted_answer=answer,
|
||||
judge_model=judge_model,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
retries=retries,
|
||||
)
|
||||
hard = 1.0 if judge["correct"] else 0.0
|
||||
|
||||
result = {
|
||||
"evaluation_mode": _EVAL_MODE,
|
||||
"predicted_answer": answer,
|
||||
"em": hard,
|
||||
"f1": hard,
|
||||
"sub_em": hard,
|
||||
"judge_model": judge_model,
|
||||
"judge_raw": judge["raw"],
|
||||
"judge_reason": judge["reason"],
|
||||
"matched_gold": judge["matched_gold"],
|
||||
}
|
||||
|
||||
if item["ans_type"] == "choice":
|
||||
result["predicted_label"] = str(answer or "").strip().upper().rstrip(".):")
|
||||
result["predicted_text"] = ""
|
||||
result["correct_label"] = str(item["correct_choice"].get("label") or "")
|
||||
result["correct_text"] = str(item["correct_choice"].get("text") or "")
|
||||
else:
|
||||
result["gold_answers"] = list(item["blank_answers"])
|
||||
best_f1 = 0.0
|
||||
for gold in item["blank_answers"]:
|
||||
best_f1 = max(best_f1, _token_f1(str(answer or ""), gold))
|
||||
result["string_f1"] = best_f1
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def evaluation_mode() -> str:
|
||||
return _EVAL_MODE
|
||||
@@ -0,0 +1,36 @@
|
||||
You are an expert failure-analysis agent for child-level visual reasoning tasks.
|
||||
|
||||
You will be given MULTIPLE failed BabyVision trajectories from a minibatch and the current skill document.
|
||||
Each trajectory includes the text prompt, the model answer, and the evaluation result.
|
||||
You do not have direct access to raw pixel content during reflection, so focus on general reasoning,
|
||||
option-selection, and visual-question-answering behaviors that can be improved through prompting.
|
||||
|
||||
## Failure Type Categories
|
||||
- **visual_detail_miss**: the agent likely overlooked a salient visual attribute, relation, count, or object state
|
||||
- **option_mismatch**: the agent selected the wrong option despite relevant evidence likely being present
|
||||
- **instruction_slip**: the agent ignored output format or answered too vaguely
|
||||
- **answer_granularity**: the agent gave an answer that was too broad, too narrow, or mismatched the expected specificity
|
||||
- **other**: none of the above
|
||||
|
||||
## Rules
|
||||
1. Focus on patterns recurring across the minibatch.
|
||||
2. Prefer reusable behaviors for inspecting images and grounding answers in visible evidence.
|
||||
3. Do not memorize dataset-specific answers.
|
||||
4. Only patch gaps not already covered by the current skill.
|
||||
|
||||
Respond ONLY with a valid JSON object:
|
||||
{
|
||||
"batch_size": <number>,
|
||||
"failure_summary": [
|
||||
{"failure_type": "<type>", "count": <int>, "description": "<one-line>"}
|
||||
],
|
||||
"patch": {
|
||||
"reasoning": "<why these edits address the common failures>",
|
||||
"edits": [
|
||||
{"op": "append", "content": "<markdown>"},
|
||||
{"op": "insert_after", "target": "<heading/text>", "content": "<markdown>"},
|
||||
{"op": "replace", "target": "<old text>", "content": "<new text>"},
|
||||
{"op": "delete", "target": "<exact text to remove>"}
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
You are an expert success-pattern analyst for child-level visual reasoning tasks.
|
||||
|
||||
You will be given MULTIPLE successful BabyVision trajectories from a minibatch and the current skill document.
|
||||
Identify generalizable behavior patterns that help the agent inspect the image carefully and answer at the right level of specificity.
|
||||
|
||||
## Rules
|
||||
- Focus on broadly useful visual QA behaviors.
|
||||
- Prefer patterns about systematic image inspection, comparing options, and concise grounded answers.
|
||||
- Do not add dataset-specific facts.
|
||||
- "edits" may be empty if the skill already captures the useful patterns.
|
||||
|
||||
Respond ONLY with a valid JSON object:
|
||||
{
|
||||
"batch_size": <number>,
|
||||
"success_patterns": ["<pattern 1>", "<pattern 2>"],
|
||||
"patch": {
|
||||
"reasoning": "<why these patterns matter>",
|
||||
"edits": [
|
||||
{"op": "append", "content": "<markdown>"},
|
||||
{"op": "insert_after", "target": "<heading/text>", "content": "<markdown>"},
|
||||
{"op": "replace", "target": "<old text>", "content": "<new text>"},
|
||||
{"op": "delete", "target": "<exact text to remove>"}
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
You are an expert diagnostic-probe designer for BabyVision-style visual reasoning tasks.
|
||||
|
||||
You will be shown representative trajectories, the current student skill, and the student's original prompt context.
|
||||
Design one SMALL diagnostic instruction that exposes the student's intermediate visual judgment without materially changing the original scaffold.
|
||||
|
||||
## Hard Constraints
|
||||
1. Do NOT substantially change the original scaffold.
|
||||
2. Do NOT prescribe a new step-by-step solving method.
|
||||
3. You MAY ask for a short structured list of a few intermediate conclusions, candidate cues, or counted units, as long as it stays close to the original scaffold.
|
||||
4. Do NOT ask for exhaustive listing of all cells, all objects, or a full chain-of-thought.
|
||||
5. Ask only for a short readout that reveals the student's current latent state.
|
||||
6. Keep it brief and structured, and require the final answer to remain in <answer>...</answer>.
|
||||
|
||||
## Good Probe Targets
|
||||
- top answer and runner-up
|
||||
- decisive visual cue
|
||||
- suspicious region or compared objects
|
||||
- counting unit or formatting interpretation
|
||||
- 2-4 short intermediate conclusions that directly support the final answer
|
||||
|
||||
Respond ONLY with a valid JSON object:
|
||||
{
|
||||
"reasoning": "<why this probe is informative>",
|
||||
"probe_instruction": "<the exact instruction text to append to the student prompt>"
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
You are a careful and strict evaluator. You will be given:
|
||||
|
||||
1. **Question**
|
||||
2. **Ground Truth Answer** (correct answer)
|
||||
3. **Model Output** (answer from another model)
|
||||
|
||||
**Your goal:** Determine if the Model Output **accurately matches** the Ground Truth Answer in meaning.
|
||||
|
||||
* Matching means: the facts, entities, and key details are equivalent, even if phrasing differs.
|
||||
* Not matching means: the Model Output is wrong, incomplete, contains extra incorrect facts, or changes the meaning.
|
||||
|
||||
**Process (internal reasoning):**
|
||||
|
||||
1. Read and understand the Question, Ground Truth Answer, and Model Output.
|
||||
2. Ignore small wording differences, formatting, or synonyms.
|
||||
3. If all factual content matches, conclude `1`. Otherwise, conclude `0`.
|
||||
|
||||
**Important:**
|
||||
|
||||
* Think through your decision step-by-step **internally** before responding.
|
||||
* In your final output, return **only** True or False, with no extra text or explanation.
|
||||
|
||||
**Output format:**
|
||||
|
||||
True
|
||||
|
||||
or
|
||||
|
||||
False
|
||||
|
||||
**Input:**
|
||||
|
||||
Question: {question},
|
||||
Ground Truth Answer: {groundtruth},
|
||||
Model Output: {modeloutput}
|
||||
@@ -0,0 +1,13 @@
|
||||
You are an expert visual reasoning agent solving child-level image understanding tasks.
|
||||
|
||||
{skill_section}## Task Format
|
||||
You will receive one image and one question about it.
|
||||
Inspect the image carefully before answering. Ground the answer in visible evidence.
|
||||
|
||||
## Answer Format
|
||||
Think step by step, then provide your final answer in \boxed{{Answer}} format.
|
||||
- For multiple-choice questions, output only the single choice label, such as \boxed{{A}}.
|
||||
- For open questions, output only a short final answer inside \boxed{{...}}.
|
||||
|
||||
Example:
|
||||
\boxed{{B}}
|
||||
@@ -0,0 +1,4 @@
|
||||
"""BabyVision Reflect stage.
|
||||
|
||||
Prompts are now loaded from .md files by the base adapter.
|
||||
"""
|
||||
@@ -0,0 +1,483 @@
|
||||
"""BabyVision rollout — multimodal visual QA with image input."""
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
from skillopt.envs.babyvision.evaluator import evaluate_item, evaluation_mode, extract_boxed_answer
|
||||
from skillopt.model import chat_student_messages, get_student_backend, is_student_exec_backend
|
||||
from skillopt.model.codex_harness import prepare_workspace, render_skill_md, run_student_exec
|
||||
from skillopt.prompts import load_prompt
|
||||
|
||||
def _build_system(skill_content: str) -> str:
|
||||
if skill_content.strip():
|
||||
skill_section = f"## Skill\n{skill_content.strip()}\n\n"
|
||||
else:
|
||||
skill_section = ""
|
||||
return load_prompt("rollout_system", env="babyvision").format(skill_section=skill_section)
|
||||
|
||||
|
||||
def _format_choices(choices: list[dict]) -> str:
|
||||
return "\n".join(f"{choice['label']}. {choice['text']}" for choice in choices)
|
||||
|
||||
|
||||
def _build_user_text(
|
||||
item: dict,
|
||||
*,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context: str = "",
|
||||
) -> str:
|
||||
parts = []
|
||||
if diagnostic_trace_context.strip():
|
||||
parts.append(
|
||||
"## Previous Codex Trace Snapshot\n"
|
||||
"This is a partial transcript from an earlier attempt. Use it as your current reasoning context.\n\n"
|
||||
f"{diagnostic_trace_context.strip()}"
|
||||
)
|
||||
parts.append(f"## Question\n{item['question']}")
|
||||
if item["ans_type"] == "choice":
|
||||
parts.append(f"## Choices\n{_format_choices(item['choices'])}")
|
||||
parts.append("Answer using the single correct option label in \\boxed{...}.")
|
||||
else:
|
||||
parts.append("Answer with a short phrase in \\boxed{...}.")
|
||||
if diagnostic_mode and diagnostic_instruction.strip():
|
||||
parts.append(f"## Training Readout\n{diagnostic_instruction.strip()}")
|
||||
return "\n\n".join(parts)
|
||||
|
||||
|
||||
def _image_to_data_uri(path: str) -> str:
|
||||
mime = mimetypes.guess_type(path)[0] or "image/png"
|
||||
with open(path, "rb") as f:
|
||||
encoded = base64.b64encode(f.read()).decode("ascii")
|
||||
return f"data:{mime};base64,{encoded}"
|
||||
|
||||
|
||||
def _build_messages(
|
||||
item: dict,
|
||||
skill_content: str,
|
||||
image_detail: str,
|
||||
*,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context: str = "",
|
||||
) -> tuple[list[dict], str, str]:
|
||||
system = _build_system(skill_content)
|
||||
user_text = _build_user_text(
|
||||
item,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
diagnostic_trace_context=diagnostic_trace_context,
|
||||
)
|
||||
image_url = {
|
||||
"url": _image_to_data_uri(item["image_path"]),
|
||||
}
|
||||
if image_detail and image_detail != "auto":
|
||||
image_url["detail"] = image_detail
|
||||
messages = [
|
||||
{"role": "system", "content": system},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": user_text},
|
||||
{"type": "image_url", "image_url": image_url},
|
||||
],
|
||||
},
|
||||
]
|
||||
return messages, system, user_text
|
||||
|
||||
|
||||
def _build_codex_skill(skill_content: str) -> str:
|
||||
return render_skill_md(
|
||||
skill_content,
|
||||
description="Dynamic ReflACT skill for solving the current BabyVision visual reasoning question.",
|
||||
preamble=(
|
||||
"Use this skill when answering the current visual reasoning question.\n"
|
||||
"Inspect the attached image carefully and return the final answer in \\boxed{...}."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _run_codex_once(
|
||||
*,
|
||||
pred_dir: str,
|
||||
item: dict,
|
||||
skill_content: str,
|
||||
model: str,
|
||||
timeout: int,
|
||||
image_detail: str,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context: str = "",
|
||||
previous_response: str = "",
|
||||
) -> tuple[str, str, str, str]:
|
||||
user_text = _build_user_text(
|
||||
item,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
diagnostic_trace_context=diagnostic_trace_context,
|
||||
)
|
||||
task_parts = [user_text]
|
||||
if previous_response:
|
||||
task_parts.append(
|
||||
"## Previous Attempt\n"
|
||||
f"{previous_response}\n\n"
|
||||
"Review the same image and question carefully. If needed, correct the answer."
|
||||
)
|
||||
task_text = "\n\n".join(task_parts)
|
||||
skill_md = _build_codex_skill(skill_content)
|
||||
work_dir = os.path.join(pred_dir, "codex_exec")
|
||||
prepare_workspace(
|
||||
work_dir=work_dir,
|
||||
skill_md=skill_md,
|
||||
task_text=task_text,
|
||||
images=[item["image_path"]],
|
||||
)
|
||||
prompt = (
|
||||
"Use the `skillopt-student` skill available in this workspace.\n"
|
||||
"Read `task.md`, inspect the attached image, and answer the question.\n"
|
||||
"Return the final answer in \\boxed{...}."
|
||||
)
|
||||
final_message, raw = run_student_exec(
|
||||
work_dir=work_dir,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
timeout=timeout,
|
||||
images=[item["image_path"]],
|
||||
)
|
||||
return final_message or raw, raw, skill_md, task_text
|
||||
|
||||
|
||||
def process_one(
|
||||
item: dict,
|
||||
out_root: str,
|
||||
skill_content: str,
|
||||
*,
|
||||
max_turns: int = 1,
|
||||
image_detail: str = "auto",
|
||||
judge_model: str = "gpt-5.4",
|
||||
judge_max_completion_tokens: int = 256,
|
||||
judge_retries: int = 5,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context: str = "",
|
||||
) -> dict:
|
||||
item_id = str(item["id"])
|
||||
result = {
|
||||
"id": item_id,
|
||||
"question": item["question"],
|
||||
"task_type": item.get("subtype") or item.get("task_type") or "babyvision",
|
||||
"task_description": item["question"],
|
||||
"hard": 0,
|
||||
"soft": 0.0,
|
||||
"predicted_answer": "",
|
||||
"predicted_label": "",
|
||||
"predicted_text": "",
|
||||
"response": "",
|
||||
"fail_reason": "",
|
||||
"agent_ok": False,
|
||||
"n_turns": 0,
|
||||
"image_path": item["image_path"],
|
||||
"ans_type": item["ans_type"],
|
||||
"evaluation_mode": evaluation_mode(),
|
||||
"judge_model": judge_model,
|
||||
}
|
||||
if item["ans_type"] == "choice":
|
||||
result["correct_label"] = item["correct_choice"]["label"]
|
||||
result["correct_text"] = item["correct_choice"]["text"]
|
||||
else:
|
||||
result["gold_answers"] = item["blank_answers"]
|
||||
|
||||
try:
|
||||
pred_dir = os.path.join(out_root, "predictions", item_id)
|
||||
os.makedirs(pred_dir, exist_ok=True)
|
||||
|
||||
if is_student_exec_backend():
|
||||
from skillopt.model import azure_openai as _llm
|
||||
|
||||
response = ""
|
||||
conversation: list[dict] = [
|
||||
{"role": "user", "content": f"{item['question']}\n\n[image] {os.path.basename(item['image_path'])}"}
|
||||
]
|
||||
system_prompt = ""
|
||||
user_text = ""
|
||||
for turn in range(max_turns):
|
||||
response, raw, system_prompt, user_text = _run_codex_once(
|
||||
pred_dir=pred_dir,
|
||||
item=item,
|
||||
skill_content=skill_content,
|
||||
model=_llm.STUDENT_DEPLOYMENT,
|
||||
timeout=120,
|
||||
image_detail=image_detail,
|
||||
diagnostic_mode=diagnostic_mode if turn == 0 else False,
|
||||
diagnostic_instruction=diagnostic_instruction if turn == 0 else "",
|
||||
diagnostic_trace_context=diagnostic_trace_context if turn == 0 else "",
|
||||
previous_response=response if turn > 0 else "",
|
||||
)
|
||||
conversation.append({"type": "message", "turn": turn + 1, "content": response})
|
||||
if extract_boxed_answer(response) is not None:
|
||||
break
|
||||
|
||||
result["response"] = response
|
||||
result["agent_ok"] = True
|
||||
result["n_turns"] = len(conversation) - 1
|
||||
with open(os.path.join(pred_dir, "student_system_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(system_prompt)
|
||||
with open(os.path.join(pred_dir, "student_user_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(user_text)
|
||||
|
||||
eval_result = evaluate_item(
|
||||
item=item,
|
||||
prediction_text=response,
|
||||
judge_model=judge_model,
|
||||
max_completion_tokens=judge_max_completion_tokens,
|
||||
retries=judge_retries,
|
||||
)
|
||||
result["evaluation_mode"] = eval_result["evaluation_mode"]
|
||||
result["judge_raw"] = eval_result["judge_raw"]
|
||||
result["judge_reason"] = eval_result["judge_reason"]
|
||||
result["matched_gold"] = eval_result["matched_gold"]
|
||||
if item["ans_type"] == "choice":
|
||||
result["predicted_label"] = eval_result["predicted_label"]
|
||||
result["predicted_text"] = eval_result["predicted_text"]
|
||||
result["predicted_answer"] = eval_result["predicted_answer"]
|
||||
result["hard"] = int(eval_result["em"])
|
||||
result["soft"] = eval_result["f1"]
|
||||
if not result["hard"]:
|
||||
result["fail_reason"] = (
|
||||
f"judge=0: predicted '{eval_result['predicted_label'] or eval_result['predicted_answer']}' "
|
||||
f"but expected '{eval_result['correct_label']}' ({eval_result['judge_reason']})"
|
||||
)
|
||||
eval_detail = (
|
||||
f"[EVALUATION RESULT]\n"
|
||||
f"Question: {item['question']}\n"
|
||||
f"Predicted label: {eval_result['predicted_label']!r}\n"
|
||||
f"Predicted text: {eval_result['predicted_text']!r}\n"
|
||||
f"Correct label: {eval_result['correct_label']!r}\n"
|
||||
f"Correct text: {eval_result['correct_text']!r}\n"
|
||||
f"Judge correct: {eval_result['em']}\n"
|
||||
f"Judge reason: {eval_result['judge_reason']}"
|
||||
)
|
||||
else:
|
||||
result["predicted_answer"] = eval_result["predicted_answer"]
|
||||
result["hard"] = int(eval_result["em"])
|
||||
result["soft"] = eval_result["f1"]
|
||||
if not result["hard"]:
|
||||
result["fail_reason"] = (
|
||||
f"judge=0: predicted '{eval_result['predicted_answer']}' "
|
||||
f"but expected {item['blank_answers']} ({eval_result['judge_reason']})"
|
||||
)
|
||||
eval_detail = (
|
||||
f"[EVALUATION RESULT]\n"
|
||||
f"Question: {item['question']}\n"
|
||||
f"Predicted answer: {eval_result['predicted_answer']!r}\n"
|
||||
f"Gold answers: {item['blank_answers']!r}\n"
|
||||
f"Judge correct: {eval_result['em']}\n"
|
||||
f"Judge reason: {eval_result['judge_reason']}\n"
|
||||
f"String F1: {eval_result.get('string_f1', 0.0):.4f}"
|
||||
)
|
||||
conversation.append({"role": "system", "content": eval_detail})
|
||||
with open(os.path.join(pred_dir, "conversation.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(conversation, f, ensure_ascii=False, indent=2)
|
||||
return result
|
||||
|
||||
messages, system_prompt, user_text = _build_messages(
|
||||
item,
|
||||
skill_content,
|
||||
image_detail,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
diagnostic_trace_context=diagnostic_trace_context,
|
||||
)
|
||||
response = ""
|
||||
conversation: list[dict] = [
|
||||
{"role": "user", "content": f"{user_text}\n\n[image] {os.path.basename(item['image_path'])}"}
|
||||
]
|
||||
|
||||
for turn in range(max_turns):
|
||||
if turn == 0:
|
||||
resp_text, _ = chat_student_messages(
|
||||
messages=messages,
|
||||
max_completion_tokens=768,
|
||||
retries=5,
|
||||
stage="rollout",
|
||||
)
|
||||
else:
|
||||
refinement_text = (
|
||||
f"Your previous answer was:\n{response}\n\n"
|
||||
"Review the same image and question carefully. "
|
||||
"If needed, correct your answer. Output the final answer in \\boxed{...}."
|
||||
)
|
||||
refinement_messages = [
|
||||
messages[0],
|
||||
messages[1],
|
||||
{"role": "assistant", "content": response},
|
||||
{"role": "user", "content": refinement_text},
|
||||
]
|
||||
resp_text, _ = chat_student_messages(
|
||||
messages=refinement_messages,
|
||||
max_completion_tokens=512,
|
||||
retries=5,
|
||||
stage="rollout",
|
||||
)
|
||||
response = resp_text
|
||||
conversation.append({"type": "message", "turn": turn + 1, "content": resp_text})
|
||||
if extract_boxed_answer(resp_text) is not None:
|
||||
break
|
||||
|
||||
result["response"] = response
|
||||
result["agent_ok"] = True
|
||||
result["n_turns"] = len(conversation) - 1
|
||||
|
||||
with open(os.path.join(pred_dir, "student_system_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(system_prompt)
|
||||
with open(os.path.join(pred_dir, "student_user_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(user_text)
|
||||
|
||||
eval_result = evaluate_item(
|
||||
item=item,
|
||||
prediction_text=response,
|
||||
judge_model=judge_model,
|
||||
max_completion_tokens=judge_max_completion_tokens,
|
||||
retries=judge_retries,
|
||||
)
|
||||
result["evaluation_mode"] = eval_result["evaluation_mode"]
|
||||
result["judge_raw"] = eval_result["judge_raw"]
|
||||
result["judge_reason"] = eval_result["judge_reason"]
|
||||
result["matched_gold"] = eval_result["matched_gold"]
|
||||
|
||||
if item["ans_type"] == "choice":
|
||||
result["predicted_label"] = eval_result["predicted_label"]
|
||||
result["predicted_text"] = eval_result["predicted_text"]
|
||||
result["predicted_answer"] = eval_result["predicted_answer"]
|
||||
result["hard"] = int(eval_result["em"])
|
||||
result["soft"] = eval_result["f1"]
|
||||
if not result["hard"]:
|
||||
result["fail_reason"] = (
|
||||
f"judge=0: predicted '{eval_result['predicted_label'] or eval_result['predicted_answer']}' "
|
||||
f"but expected '{eval_result['correct_label']}' ({eval_result['judge_reason']})"
|
||||
)
|
||||
eval_detail = (
|
||||
f"[EVALUATION RESULT]\n"
|
||||
f"Question: {item['question']}\n"
|
||||
f"Predicted label: {eval_result['predicted_label']!r}\n"
|
||||
f"Predicted text: {eval_result['predicted_text']!r}\n"
|
||||
f"Correct label: {eval_result['correct_label']!r}\n"
|
||||
f"Correct text: {eval_result['correct_text']!r}\n"
|
||||
f"Judge correct: {eval_result['em']}\n"
|
||||
f"Judge reason: {eval_result['judge_reason']}"
|
||||
)
|
||||
else:
|
||||
result["predicted_answer"] = eval_result["predicted_answer"]
|
||||
result["hard"] = int(eval_result["em"])
|
||||
result["soft"] = eval_result["f1"]
|
||||
if not result["hard"]:
|
||||
result["fail_reason"] = (
|
||||
f"judge=0: predicted '{eval_result['predicted_answer']}' "
|
||||
f"but expected {item['blank_answers']} ({eval_result['judge_reason']})"
|
||||
)
|
||||
eval_detail = (
|
||||
f"[EVALUATION RESULT]\n"
|
||||
f"Question: {item['question']}\n"
|
||||
f"Predicted answer: {eval_result['predicted_answer']!r}\n"
|
||||
f"Gold answers: {item['blank_answers']!r}\n"
|
||||
f"Judge correct: {eval_result['em']}\n"
|
||||
f"Judge reason: {eval_result['judge_reason']}\n"
|
||||
f"String F1: {eval_result.get('string_f1', 0.0):.4f}"
|
||||
)
|
||||
|
||||
conversation.append({"role": "system", "content": eval_detail})
|
||||
with open(os.path.join(pred_dir, "conversation.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(conversation, f, ensure_ascii=False, indent=2)
|
||||
except Exception as e: # noqa: BLE001
|
||||
result["fail_reason"] = f"error: {e}"
|
||||
return result
|
||||
|
||||
|
||||
def run_batch(
|
||||
items: list[dict],
|
||||
out_root: str,
|
||||
skill_content: str,
|
||||
*,
|
||||
max_turns: int = 1,
|
||||
workers: int = 32,
|
||||
image_detail: str = "auto",
|
||||
judge_model: str = "gpt-5.4",
|
||||
judge_max_completion_tokens: int = 256,
|
||||
judge_retries: int = 5,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context_by_id: dict[str, str] | None = None,
|
||||
) -> list[dict]:
|
||||
results_path = os.path.join(out_root, "results.jsonl")
|
||||
os.makedirs(out_root, exist_ok=True)
|
||||
|
||||
expected_eval_mode = evaluation_mode()
|
||||
done_ids: set[str] = set()
|
||||
existing: list[dict] = []
|
||||
rewrite_results = False
|
||||
if os.path.exists(results_path):
|
||||
with open(results_path, encoding="utf-8") as f:
|
||||
for line in f:
|
||||
try:
|
||||
row = json.loads(line)
|
||||
if row.get("evaluation_mode") != expected_eval_mode:
|
||||
rewrite_results = True
|
||||
continue
|
||||
done_ids.add(str(row["id"]))
|
||||
existing.append(row)
|
||||
except Exception:
|
||||
rewrite_results = True
|
||||
|
||||
pending = [item for item in items if str(item["id"]) not in done_ids]
|
||||
if not pending and not rewrite_results:
|
||||
return existing
|
||||
|
||||
total = len(existing) + len(pending)
|
||||
completed = len(existing)
|
||||
correct_count = sum(1 for r in existing if r.get("hard", 0))
|
||||
if existing:
|
||||
print(f" [rollout] resuming: {completed}/{total} already done", flush=True)
|
||||
|
||||
results = list(existing)
|
||||
file_mode = "w" if rewrite_results else "a"
|
||||
with open(results_path, file_mode, encoding="utf-8") as outf, ThreadPoolExecutor(max_workers=workers) as ex:
|
||||
if rewrite_results:
|
||||
for row in existing:
|
||||
outf.write(json.dumps(row, ensure_ascii=False) + "\n")
|
||||
futs = {
|
||||
ex.submit(
|
||||
process_one,
|
||||
item,
|
||||
out_root,
|
||||
skill_content,
|
||||
max_turns=max_turns,
|
||||
image_detail=image_detail,
|
||||
judge_model=judge_model,
|
||||
judge_max_completion_tokens=judge_max_completion_tokens,
|
||||
judge_retries=judge_retries,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
diagnostic_trace_context=(diagnostic_trace_context_by_id or {}).get(str(item["id"]), ""),
|
||||
): item
|
||||
for item in pending
|
||||
}
|
||||
for fut in as_completed(futs):
|
||||
row = fut.result()
|
||||
results.append(row)
|
||||
completed += 1
|
||||
if row.get("hard", 0):
|
||||
correct_count += 1
|
||||
acc = correct_count / completed if completed else 0
|
||||
print(
|
||||
f" [rollout] {completed}/{total} "
|
||||
f"(acc={acc:.3f}) id={row.get('id', '?')} "
|
||||
f"hard={row.get('hard', '?')}",
|
||||
flush=True,
|
||||
)
|
||||
outf.write(json.dumps(row, ensure_ascii=False) + "\n")
|
||||
outf.flush()
|
||||
return results
|
||||
@@ -0,0 +1,18 @@
|
||||
# BabyVision Visual QA Heuristics
|
||||
|
||||
## Image Inspection
|
||||
- First identify the main objects, their attributes, and their spatial relations before answering.
|
||||
- If the question involves counting, compare all relevant instances carefully instead of stopping after the first match.
|
||||
- If the question asks about color, size, position, or action, verify the specific visible evidence for that attribute.
|
||||
|
||||
## Multiple Choice
|
||||
- Compare every option against the visible image evidence before deciding.
|
||||
- Prefer the option that matches the image exactly; reject options that are only partially true or too vague.
|
||||
- When two options are close, check the smallest discriminating visual detail.
|
||||
|
||||
## Open Answers
|
||||
- Answer with the shortest phrase that is fully supported by the image.
|
||||
- Match the expected level of specificity: not broader than the image evidence, not narrower than the question asks.
|
||||
|
||||
## Final Answer
|
||||
- Output only the final answer inside <answer>...</answer>.
|
||||
@@ -0,0 +1,396 @@
|
||||
"""ReflACT environment adapter — abstract interface.
|
||||
|
||||
To connect ReflACT to a new environment (benchmark, simulator, etc.),
|
||||
implement a subclass of :class:`EnvAdapter` with environment-specific
|
||||
rollout and reflection logic.
|
||||
|
||||
Example::
|
||||
|
||||
class MyBenchAdapter(EnvAdapter):
|
||||
def build_train_env(self, batch_size, seed, **kw):
|
||||
return MyEnvManager(split="train", n=batch_size, seed=seed)
|
||||
|
||||
def build_eval_env(self, env_num, split, seed, **kw):
|
||||
return MyEnvManager(split=split, n=env_num, seed=seed)
|
||||
|
||||
def rollout(self, env_manager, skill_content, out_dir, **kw):
|
||||
# Run episodes, return [{"id": ..., "hard": 0/1, "soft": 0.0-1.0, ...}]
|
||||
...
|
||||
|
||||
def reflect(self, results, skill_content, out_dir, **kw):
|
||||
# Analyze trajectories, return list of patch dicts
|
||||
...
|
||||
|
||||
def get_task_types(self):
|
||||
return ["task_a", "task_b"]
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
import os
|
||||
import random
|
||||
|
||||
from skillopt.datasets.base import BaseDataLoader, BatchSpec
|
||||
from skillopt.model.codex_harness import extract_codex_trace_prefix, format_codex_trace_steps, parse_codex_raw
|
||||
from skillopt.prompts import load_prompt
|
||||
|
||||
|
||||
class EnvAdapter(ABC):
|
||||
"""Abstract adapter for connecting ReflACT to any environment.
|
||||
|
||||
Subclasses must implement all abstract methods. The ReflACT trainer
|
||||
calls these methods at the appropriate pipeline stages.
|
||||
"""
|
||||
|
||||
# ── Lifecycle hooks ────────────────────────────────────────────────────
|
||||
|
||||
def setup(self, cfg: dict) -> None:
|
||||
"""Called once by the trainer before the training loop begins.
|
||||
|
||||
Override to perform one-time initialization that requires the full
|
||||
config (e.g., data loading, split creation). Default is a no-op.
|
||||
"""
|
||||
self._cfg = dict(cfg)
|
||||
|
||||
def get_dataloader(self) -> BaseDataLoader | None:
|
||||
"""Return the task dataloader used by this adapter, if any."""
|
||||
return None
|
||||
|
||||
def requires_ray(self) -> bool:
|
||||
"""Return whether this adapter requires Ray runtime initialization."""
|
||||
return False
|
||||
|
||||
def deep_reflect(
|
||||
self,
|
||||
results: list[dict],
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict | None]:
|
||||
"""Optional deeper diagnostic reflection pass.
|
||||
|
||||
Default behavior is a no-op. Dataset-backed adapters may override this
|
||||
to re-query the student on a small representative subset of the current
|
||||
batch using minimally-perturbed diagnostic prompts that expose
|
||||
intermediate reasoning state.
|
||||
"""
|
||||
return []
|
||||
|
||||
def build_reference_text(self, item: dict) -> str:
|
||||
"""Return hidden reference material for deep reflection, if any."""
|
||||
return str(item.get("reference_text") or "").strip()
|
||||
|
||||
def get_reference_metadata(self, item: dict) -> dict:
|
||||
"""Return structured metadata about hidden reference material."""
|
||||
reference_text = self.build_reference_text(item)
|
||||
if not reference_text:
|
||||
return {"fields": [], "preview": ""}
|
||||
return {
|
||||
"fields": ["reference_text"],
|
||||
"preview": reference_text[:400],
|
||||
}
|
||||
|
||||
def get_codex_deep_probe_prompt(self) -> str | None:
|
||||
env_name = getattr(self, "_cfg", {}).get("env_name")
|
||||
return load_prompt("deep_probe_codex", env=env_name)
|
||||
|
||||
def attach_codex_probe_context(
|
||||
self,
|
||||
results: list[dict],
|
||||
prediction_dir: str,
|
||||
) -> list[dict]:
|
||||
"""Attach compact Codex step metadata for codex-aware deep reflection."""
|
||||
enriched: list[dict] = []
|
||||
for row in results:
|
||||
merged = dict(row)
|
||||
tid = str(row.get("id"))
|
||||
raw_path = os.path.join(prediction_dir, tid, "codex_raw.txt")
|
||||
if os.path.exists(raw_path):
|
||||
with open(raw_path, encoding="utf-8") as f:
|
||||
raw = f.read()
|
||||
parsed = parse_codex_raw(raw)
|
||||
merged["codex_probe_trace_steps"] = format_codex_trace_steps(raw)
|
||||
merged["codex_probe_step_count"] = len(parsed["steps"])
|
||||
enriched.append(merged)
|
||||
return enriched
|
||||
|
||||
def resolve_codex_probe_target(
|
||||
self,
|
||||
*,
|
||||
selected_items: list[dict],
|
||||
selected_examples: list[dict],
|
||||
prediction_dir: str,
|
||||
probe: dict,
|
||||
) -> tuple[list[dict], dict[str, str] | None, dict]:
|
||||
"""Resolve the teacher-selected codex probe target and raw trace prefix."""
|
||||
target_id = str(probe.get("probe_target_id", "")).strip()
|
||||
selected_id_set = {str(item["id"]) for item in selected_items}
|
||||
if target_id not in selected_id_set:
|
||||
target_id = str(selected_items[0]["id"])
|
||||
target_item = next(item for item in selected_items if str(item["id"]) == target_id)
|
||||
target_result = next(
|
||||
(row for row in selected_examples if str(row.get("id")) == target_id),
|
||||
None,
|
||||
)
|
||||
max_probe_step = int((target_result or {}).get("codex_probe_step_count", 0))
|
||||
default_probe_step = max_probe_step - 1 if max_probe_step > 1 else max_probe_step
|
||||
probe_after_step = int(probe.get("probe_after_step", default_probe_step))
|
||||
if max_probe_step > 0:
|
||||
probe_after_step = max(0, min(probe_after_step, max_probe_step))
|
||||
else:
|
||||
probe_after_step = 0
|
||||
raw_path = os.path.join(prediction_dir, target_id, "codex_raw.txt")
|
||||
trace_prefix = ""
|
||||
if os.path.exists(raw_path):
|
||||
with open(raw_path, encoding="utf-8") as f:
|
||||
trace_prefix = extract_codex_trace_prefix(f.read(), after_step=probe_after_step)
|
||||
updated_probe = dict(probe)
|
||||
updated_probe["probe_target_id"] = target_id
|
||||
updated_probe["probe_after_step"] = probe_after_step
|
||||
return [target_item], {target_id: trace_prefix}, updated_probe
|
||||
|
||||
def attach_reference_context(
|
||||
self,
|
||||
results: list[dict],
|
||||
items: list[dict] | None,
|
||||
) -> list[dict]:
|
||||
"""Attach environment-specific hidden reference text to result dicts."""
|
||||
if not results or not items:
|
||||
return list(results)
|
||||
|
||||
item_by_id = {
|
||||
str(item.get("id")): item
|
||||
for item in items
|
||||
if isinstance(item, dict) and item.get("id") is not None
|
||||
}
|
||||
enriched: list[dict] = []
|
||||
for row in results:
|
||||
merged = dict(row)
|
||||
item = item_by_id.get(str(row.get("id")))
|
||||
if item:
|
||||
reference_text = self.build_reference_text(item)
|
||||
if reference_text:
|
||||
merged["reference_text"] = reference_text
|
||||
enriched.append(merged)
|
||||
return enriched
|
||||
|
||||
def select_representative_items(
|
||||
self,
|
||||
results: list[dict],
|
||||
items: list[dict] | None,
|
||||
*,
|
||||
n_failures: int,
|
||||
n_successes: int,
|
||||
seed: int | None = None,
|
||||
) -> list[dict]:
|
||||
"""Select a small diverse subset of current-batch items by outcome."""
|
||||
if not items:
|
||||
return []
|
||||
|
||||
item_by_id = {
|
||||
str(item.get("id")): item
|
||||
for item in items
|
||||
if isinstance(item, dict) and item.get("id") is not None
|
||||
}
|
||||
failures = [
|
||||
(result, item_by_id[str(result.get("id"))])
|
||||
for result in results
|
||||
if not result.get("hard") and str(result.get("id")) in item_by_id
|
||||
]
|
||||
successes = [
|
||||
(result, item_by_id[str(result.get("id"))])
|
||||
for result in results
|
||||
if result.get("hard") and str(result.get("id")) in item_by_id
|
||||
]
|
||||
|
||||
rng = random.Random(seed)
|
||||
|
||||
def _pick(pool: list[tuple[dict, dict]], quota: int) -> list[dict]:
|
||||
if quota <= 0 or not pool:
|
||||
return []
|
||||
shuffled = list(pool)
|
||||
rng.shuffle(shuffled)
|
||||
|
||||
picked_ids: set[str] = set()
|
||||
picked: list[dict] = []
|
||||
seen_types: set[str] = set()
|
||||
|
||||
for result, item in shuffled:
|
||||
task_type = str(result.get("task_type") or item.get("task_type") or item.get("subtype") or "unknown")
|
||||
item_id = str(item["id"])
|
||||
if task_type in seen_types or item_id in picked_ids:
|
||||
continue
|
||||
picked.append(item)
|
||||
picked_ids.add(item_id)
|
||||
seen_types.add(task_type)
|
||||
if len(picked) >= quota:
|
||||
return picked
|
||||
|
||||
for _, item in shuffled:
|
||||
item_id = str(item["id"])
|
||||
if item_id in picked_ids:
|
||||
continue
|
||||
picked.append(item)
|
||||
picked_ids.add(item_id)
|
||||
if len(picked) >= quota:
|
||||
break
|
||||
return picked
|
||||
|
||||
selected = _pick(failures, n_failures)
|
||||
selected_ids = {str(item["id"]) for item in selected}
|
||||
selected.extend(
|
||||
item for item in _pick(successes, n_successes)
|
||||
if str(item["id"]) not in selected_ids
|
||||
)
|
||||
return selected
|
||||
|
||||
def build_env_from_batch(self, batch: BatchSpec, **kwargs):
|
||||
"""Build an environment manager or item list from a :class:`BatchSpec`.
|
||||
|
||||
Default behavior preserves the legacy adapter API by routing training
|
||||
batches through :meth:`build_train_env` and evaluation batches through
|
||||
:meth:`build_eval_env`.
|
||||
"""
|
||||
if batch.phase == "train":
|
||||
return self.build_train_env(batch_size=batch.batch_size, seed=batch.seed, **kwargs)
|
||||
return self.build_eval_env(
|
||||
env_num=batch.batch_size,
|
||||
split=batch.split,
|
||||
seed=batch.seed,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
def build_train_env(self, batch_size: int, seed: int, **kwargs):
|
||||
"""Build a training environment manager.
|
||||
|
||||
Returns
|
||||
-------
|
||||
object
|
||||
An environment manager that can be passed to :meth:`rollout`.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def build_eval_env(self, env_num: int, split: str, seed: int, **kwargs):
|
||||
"""Build an evaluation environment manager.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
env_num : int
|
||||
Number of evaluation environments.
|
||||
split : str
|
||||
Dataset split (e.g. ``"valid_seen"``, ``"valid_unseen"``).
|
||||
seed : int
|
||||
Random seed for reproducibility.
|
||||
|
||||
Returns
|
||||
-------
|
||||
object
|
||||
An environment manager that can be passed to :meth:`rollout`.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def rollout(
|
||||
self,
|
||||
env_manager,
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict]:
|
||||
"""Run a batch of episodes using the current skill.
|
||||
|
||||
Returns
|
||||
-------
|
||||
list[dict]
|
||||
Each dict conforms to :class:`~skillopt.types.RolloutResult`:
|
||||
must have ``"id"`` (str), ``"hard"`` (0/1), ``"soft"``
|
||||
(float 0-1). May include env-specific fields.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def reflect(
|
||||
self,
|
||||
results: list[dict],
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict | None]:
|
||||
"""Analyze rollout results and produce patches.
|
||||
|
||||
Each returned dict conforms to :class:`~skillopt.types.RawPatch`:
|
||||
``"patch"`` (with ``"edits"`` list) + ``"source_type"``
|
||||
(``"failure"`` or ``"success"``).
|
||||
|
||||
Returns
|
||||
-------
|
||||
list[dict | None]
|
||||
Raw analyst outputs; ``None`` entries are filtered out.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get_task_types(self) -> list[str]:
|
||||
"""Return the list of task type names for this environment."""
|
||||
|
||||
# ── Prompt configuration (two-level priority) ────────────────────────
|
||||
#
|
||||
# Priority: env-specific prompt file > generic default prompt file.
|
||||
#
|
||||
# Prompts are loaded from ``.md`` files via ``load_prompt(name, env)``:
|
||||
# 1. ``skillopt/envs/<env>/prompts/<name>.md`` (env-specific)
|
||||
# 2. ``skillopt/prompts/<name>.md`` (generic fallback)
|
||||
#
|
||||
# Subclasses can still override ``get_*_prompt()`` for full control.
|
||||
|
||||
@property
|
||||
def _env_name(self) -> str:
|
||||
"""Derive the env directory name from this adapter's module path."""
|
||||
# e.g. "skillopt.envs.searchqa.adapter" → "searchqa"
|
||||
module = type(self).__module__
|
||||
parts = module.split(".")
|
||||
if len(parts) >= 3 and parts[-3] == "envs":
|
||||
return parts[-2]
|
||||
return ""
|
||||
|
||||
def _load_env_prompt(self, name: str) -> str | None:
|
||||
"""Load a prompt with env-specific override. Returns None if not found."""
|
||||
try:
|
||||
return load_prompt(name, env=self._env_name)
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
|
||||
def get_error_minibatch_prompt(self) -> str | None:
|
||||
update_mode = getattr(self, "_cfg", {}).get("skill_update_mode", "patch")
|
||||
raw_mode = str(update_mode).strip().lower()
|
||||
if raw_mode in {"full_rewrite", "full_rewrite_minibatch", "minibatch_full_rewrite", "skill_rewrite_minibatch"}:
|
||||
prompt = self._load_env_prompt("analyst_error_full_rewrite")
|
||||
if prompt is not None:
|
||||
return prompt
|
||||
if raw_mode in {"rewrite", "rewrite_from_suggestions", "suggestions", "rewrite_suggestions"}:
|
||||
prompt = self._load_env_prompt("analyst_error_rewrite")
|
||||
if prompt is not None:
|
||||
return prompt
|
||||
return self._load_env_prompt("analyst_error")
|
||||
|
||||
def get_success_minibatch_prompt(self) -> str | None:
|
||||
update_mode = getattr(self, "_cfg", {}).get("skill_update_mode", "patch")
|
||||
raw_mode = str(update_mode).strip().lower()
|
||||
if raw_mode in {"full_rewrite", "full_rewrite_minibatch", "minibatch_full_rewrite", "skill_rewrite_minibatch"}:
|
||||
prompt = self._load_env_prompt("analyst_success_full_rewrite")
|
||||
if prompt is not None:
|
||||
return prompt
|
||||
if raw_mode in {"rewrite", "rewrite_from_suggestions", "suggestions", "rewrite_suggestions"}:
|
||||
prompt = self._load_env_prompt("analyst_success_rewrite")
|
||||
if prompt is not None:
|
||||
return prompt
|
||||
return self._load_env_prompt("analyst_success")
|
||||
|
||||
def get_deep_probe_prompt(self) -> str | None:
|
||||
return self._load_env_prompt("deep_probe")
|
||||
|
||||
def get_meta_reflect_prompt(self) -> str | None:
|
||||
update_mode = getattr(self, "_cfg", {}).get("skill_update_mode", "patch")
|
||||
if str(update_mode).strip().lower() == "rewrite_from_suggestions":
|
||||
prompt = self._load_env_prompt("meta_reflect_rewrite")
|
||||
if prompt is not None:
|
||||
return prompt
|
||||
return self._load_env_prompt("meta_reflect")
|
||||
@@ -0,0 +1,114 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Callable
|
||||
|
||||
from skillopt.gradient.deep_probe import generate_deep_probe_instruction
|
||||
from skillopt.gradient.reflect import run_minibatch_reflect
|
||||
|
||||
|
||||
def run_no_reference_deep_reflect(
|
||||
adapter: Any,
|
||||
results: list[dict],
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
*,
|
||||
env_manager: Any = None,
|
||||
prediction_dir: str | None = None,
|
||||
random_seed: int | None = None,
|
||||
step_buffer_context: str = "",
|
||||
output_requirements: list[str] | None = None,
|
||||
metadata_builder: Callable[[dict], dict] | None = None,
|
||||
) -> list[dict | None]:
|
||||
"""Run teacher-designed diagnostic probing without hidden references."""
|
||||
if not getattr(adapter, "use_deep_reflect", False):
|
||||
return []
|
||||
if not isinstance(env_manager, list):
|
||||
return []
|
||||
|
||||
prediction_dir = prediction_dir or os.path.join(out_dir, "predictions")
|
||||
selected_items = adapter.select_representative_items(
|
||||
results,
|
||||
env_manager,
|
||||
n_failures=getattr(adapter, "deep_reflect_failures", 4),
|
||||
n_successes=getattr(adapter, "deep_reflect_successes", 2),
|
||||
seed=random_seed,
|
||||
)
|
||||
if not selected_items:
|
||||
return []
|
||||
|
||||
selected_ids = {str(item["id"]) for item in selected_items}
|
||||
selected_results = [row for row in results if str(row.get("id")) in selected_ids]
|
||||
if metadata_builder is None:
|
||||
selected_metadata = [
|
||||
{
|
||||
"id": str(item.get("id")),
|
||||
"task_type": str(item.get("task_type") or item.get("topic") or "unknown"),
|
||||
"question_preview": str(item.get("question") or "")[:200],
|
||||
}
|
||||
for item in selected_items
|
||||
]
|
||||
else:
|
||||
selected_metadata = [metadata_builder(item) for item in selected_items]
|
||||
|
||||
deep_dir = os.path.join(out_dir, "deep_reflect")
|
||||
rollout_dir = os.path.join(deep_dir, "rollout")
|
||||
patches_dir = os.path.join(deep_dir, "patches")
|
||||
os.makedirs(deep_dir, exist_ok=True)
|
||||
print(
|
||||
f" [2b/6 DEEP REFLECT setup] selected={len(selected_items)} "
|
||||
"mode=no_reference_probe"
|
||||
)
|
||||
|
||||
probe = generate_deep_probe_instruction(
|
||||
skill_content=skill_content,
|
||||
items=selected_results,
|
||||
prediction_dir=prediction_dir,
|
||||
system_prompt=adapter.get_deep_probe_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
output_requirements=output_requirements,
|
||||
)
|
||||
if not probe:
|
||||
return []
|
||||
|
||||
with open(os.path.join(deep_dir, "probe.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(
|
||||
{
|
||||
**probe,
|
||||
"reference_summary": {
|
||||
"mode": "no_reference_probe",
|
||||
"selected_count": len(selected_items),
|
||||
},
|
||||
"selected_examples": selected_metadata,
|
||||
},
|
||||
f,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
|
||||
deep_results = adapter.rollout(
|
||||
selected_items,
|
||||
skill_content,
|
||||
rollout_dir,
|
||||
diagnostic_mode=True,
|
||||
diagnostic_instruction=probe["probe_instruction"],
|
||||
)
|
||||
return run_minibatch_reflect(
|
||||
results=deep_results,
|
||||
skill_content=skill_content,
|
||||
prediction_dir=os.path.join(rollout_dir, "predictions"),
|
||||
patches_dir=patches_dir,
|
||||
workers=getattr(adapter, "analyst_workers", 8),
|
||||
failure_only=getattr(adapter, "failure_only", False),
|
||||
minibatch_size=getattr(adapter, "minibatch_size", 8),
|
||||
edit_budget=getattr(adapter, "edit_budget", 4),
|
||||
random_seed=random_seed,
|
||||
error_system=adapter.get_error_minibatch_prompt(),
|
||||
success_system=adapter.get_success_minibatch_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
update_mode=getattr(getattr(adapter, "_cfg", {}), "get", lambda *_: "patch")(
|
||||
"skill_update_mode",
|
||||
"patch",
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1 @@
|
||||
"""DocVQA environment package for ReflACT."""
|
||||
@@ -0,0 +1,151 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
from skillopt.datasets.base import BatchSpec
|
||||
from skillopt.envs.base import EnvAdapter
|
||||
from skillopt.envs.deep_reflect import run_no_reference_deep_reflect
|
||||
from skillopt.envs.docvqa.dataloader import DocVQADataLoader
|
||||
from skillopt.envs.docvqa.rollout import run_batch
|
||||
from skillopt.gradient.reflect import run_minibatch_reflect
|
||||
|
||||
|
||||
class DocVQAAdapter(EnvAdapter):
|
||||
def __init__(
|
||||
self,
|
||||
split_dir: str = "",
|
||||
data_path: str = "",
|
||||
split_mode: str = "split_dir",
|
||||
split_ratio: str = "2:1:7",
|
||||
split_seed: int = 42,
|
||||
split_output_dir: str = "",
|
||||
max_turns: int = 1,
|
||||
exec_timeout: int = 120,
|
||||
workers: int = 16,
|
||||
analyst_workers: int = 16,
|
||||
failure_only: bool = False,
|
||||
minibatch_size: int = 8,
|
||||
edit_budget: int = 4,
|
||||
seed: int = 42,
|
||||
limit: int = 0,
|
||||
image_detail: str = "auto",
|
||||
use_deep_reflect: bool = False,
|
||||
deep_reflect_failures: int = 4,
|
||||
deep_reflect_successes: int = 2,
|
||||
) -> None:
|
||||
self.max_turns = max_turns
|
||||
self.exec_timeout = exec_timeout
|
||||
self.workers = workers
|
||||
self.analyst_workers = analyst_workers
|
||||
self.failure_only = failure_only
|
||||
self.minibatch_size = minibatch_size
|
||||
self.edit_budget = edit_budget
|
||||
self.image_detail = image_detail
|
||||
self.use_deep_reflect = use_deep_reflect
|
||||
self.deep_reflect_failures = deep_reflect_failures
|
||||
self.deep_reflect_successes = deep_reflect_successes
|
||||
self.dataloader = DocVQADataLoader(
|
||||
split_dir=split_dir,
|
||||
data_path=data_path,
|
||||
split_mode=split_mode,
|
||||
split_ratio=split_ratio,
|
||||
split_seed=split_seed,
|
||||
split_output_dir=split_output_dir,
|
||||
seed=seed,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
def setup(self, cfg: dict) -> None:
|
||||
super().setup(cfg)
|
||||
self.dataloader.setup(cfg)
|
||||
|
||||
def get_dataloader(self):
|
||||
return self.dataloader
|
||||
|
||||
def build_env_from_batch(self, batch: BatchSpec, **kwargs):
|
||||
return list(batch.payload or [])
|
||||
|
||||
def build_train_env(self, batch_size: int, seed: int, **kwargs):
|
||||
batch = self.dataloader.build_train_batch(batch_size=batch_size, seed=seed, **kwargs)
|
||||
return self.build_env_from_batch(batch, **kwargs)
|
||||
|
||||
def build_eval_env(self, env_num: int, split: str, seed: int, **kwargs):
|
||||
batch = self.dataloader.build_eval_batch(env_num=env_num, split=split, seed=seed, **kwargs)
|
||||
return self.build_env_from_batch(batch, **kwargs)
|
||||
|
||||
def rollout(self, env_manager, skill_content: str, out_dir: str, **kwargs) -> list[dict]:
|
||||
items: list[dict] = env_manager
|
||||
return run_batch(
|
||||
items=items,
|
||||
out_root=out_dir,
|
||||
skill_content=skill_content,
|
||||
max_turns=self.max_turns,
|
||||
exec_timeout=self.exec_timeout,
|
||||
workers=self.workers,
|
||||
image_detail=self.image_detail,
|
||||
diagnostic_mode=kwargs.get("diagnostic_mode", False),
|
||||
diagnostic_instruction=kwargs.get("diagnostic_instruction", ""),
|
||||
task_timeout=self.exec_timeout,
|
||||
)
|
||||
|
||||
def reflect(self, results: list[dict], skill_content: str, out_dir: str, **kwargs) -> list[dict | None]:
|
||||
prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
|
||||
patches_dir = kwargs.get("patches_dir", os.path.join(out_dir, "patches"))
|
||||
random_seed = kwargs.get("random_seed")
|
||||
step_buffer_context = kwargs.get("step_buffer_context", "")
|
||||
return run_minibatch_reflect(
|
||||
results=results,
|
||||
skill_content=skill_content,
|
||||
prediction_dir=prediction_dir,
|
||||
patches_dir=patches_dir,
|
||||
workers=self.analyst_workers,
|
||||
failure_only=self.failure_only,
|
||||
minibatch_size=self.minibatch_size,
|
||||
edit_budget=self.edit_budget,
|
||||
random_seed=random_seed,
|
||||
error_system=self.get_error_minibatch_prompt(),
|
||||
success_system=self.get_success_minibatch_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
update_mode=getattr(self, "_cfg", {}).get("skill_update_mode", "patch"),
|
||||
)
|
||||
|
||||
def deep_reflect(
|
||||
self,
|
||||
results: list[dict],
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict | None]:
|
||||
return run_no_reference_deep_reflect(
|
||||
self,
|
||||
results,
|
||||
skill_content,
|
||||
out_dir,
|
||||
env_manager=kwargs.get("env_manager"),
|
||||
prediction_dir=kwargs.get("prediction_dir"),
|
||||
random_seed=kwargs.get("random_seed"),
|
||||
step_buffer_context=kwargs.get("step_buffer_context", ""),
|
||||
output_requirements=[
|
||||
"- There is no hidden reference block. Use only the document image prompt, student output, and evaluation result to infer what intermediate state is worth probing.",
|
||||
"- The instruction must explicitly request a short <analysis>...</analysis> block before the final <answer>...</answer>.",
|
||||
"- The readout should focus on visual region, field/table/figure label, OCR text read, candidate answer, and answer-format normalization.",
|
||||
"- Do not ask for exhaustive transcription or a full chain-of-thought.",
|
||||
"- The instruction text should be ready to append directly to the student's prompt.",
|
||||
],
|
||||
metadata_builder=lambda item: {
|
||||
"id": str(item.get("id")),
|
||||
"task_type": str(item.get("task_type") or "docvqa"),
|
||||
"question_preview": str(item.get("question") or "")[:200],
|
||||
"image_path": item.get("image_path", ""),
|
||||
"docId": item.get("docId", ""),
|
||||
"page": item.get("ucsf_document_page_no", ""),
|
||||
},
|
||||
)
|
||||
|
||||
def get_task_types(self) -> list[str]:
|
||||
seen: list[str] = []
|
||||
for item in self.dataloader.train_items + self.dataloader.val_items + self.dataloader.test_items:
|
||||
task_type = str(item.get("task_type") or "docvqa")
|
||||
if task_type not in seen:
|
||||
seen.append(task_type)
|
||||
return seen or ["docvqa"]
|
||||
@@ -0,0 +1,61 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import csv
|
||||
from pathlib import Path
|
||||
|
||||
from skillopt.datasets.base import SplitDataLoader
|
||||
|
||||
|
||||
def _parse_answers(raw: str) -> list[str]:
|
||||
text = str(raw or "").strip()
|
||||
if not text:
|
||||
return []
|
||||
try:
|
||||
parsed = ast.literal_eval(text)
|
||||
except Exception:
|
||||
return [text]
|
||||
if isinstance(parsed, list):
|
||||
return [str(item).strip() for item in parsed if str(item).strip()]
|
||||
return [str(parsed).strip()]
|
||||
|
||||
|
||||
def _extract_document_path(question: str) -> tuple[str, str]:
|
||||
marker = "document_path:"
|
||||
if marker not in question:
|
||||
return question.strip(), ""
|
||||
main, tail = question.split(marker, 1)
|
||||
return main.strip(), tail.strip()
|
||||
|
||||
|
||||
def _normalize_row(row: dict[str, str]) -> dict:
|
||||
question_text, document_path = _extract_document_path(str(row.get("question") or ""))
|
||||
answers = _parse_answers(row.get("answer") or row.get("ground_truth") or "")
|
||||
image_path = str(row.get("image_path") or document_path or "").strip()
|
||||
task_type = str(row.get("topic") or row.get("category") or "docvqa").strip() or "docvqa"
|
||||
return {
|
||||
"id": str(row.get("questionId") or row.get("id") or "").strip(),
|
||||
"question": question_text,
|
||||
"answer": answers[0] if answers else "",
|
||||
"answers": answers,
|
||||
"task_type": task_type,
|
||||
"subtask": task_type,
|
||||
"image_paths": [image_path] if image_path else [],
|
||||
"image_path": image_path,
|
||||
"questionId": str(row.get("questionId") or "").strip(),
|
||||
"docId": str(row.get("docId") or "").strip(),
|
||||
"ucsf_document_id": str(row.get("ucsf_document_id") or "").strip(),
|
||||
"ucsf_document_page_no": str(row.get("ucsf_document_page_no") or "").strip(),
|
||||
"source_split": str(row.get("source_split") or "").strip(),
|
||||
}
|
||||
|
||||
|
||||
class DocVQADataLoader(SplitDataLoader):
|
||||
def load_split_items(self, split_path: str) -> list[dict]:
|
||||
path = Path(split_path)
|
||||
csv_files = sorted(path.glob("*.csv"))
|
||||
if not csv_files:
|
||||
raise FileNotFoundError(f"No .csv file found in {split_path}")
|
||||
with csv_files[0].open(encoding="utf-8", newline="") as f:
|
||||
reader = csv.DictReader(f)
|
||||
return [_normalize_row(row) for row in reader]
|
||||
@@ -0,0 +1,113 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import json
|
||||
from collections.abc import Iterable
|
||||
from typing import Any
|
||||
|
||||
DEFAULT_ANLS_THRESHOLD = 0.5
|
||||
|
||||
|
||||
def _normalize_text(value: Any) -> str:
|
||||
if value is None:
|
||||
return ""
|
||||
text = str(value).strip().lower()
|
||||
return " ".join(text.split())
|
||||
|
||||
|
||||
def _levenshtein_distance(a: str, b: str) -> int:
|
||||
if a == b:
|
||||
return 0
|
||||
if not a:
|
||||
return len(b)
|
||||
if not b:
|
||||
return len(a)
|
||||
if len(a) > len(b):
|
||||
a, b = b, a
|
||||
previous = list(range(len(b) + 1))
|
||||
for i, char_a in enumerate(a, start=1):
|
||||
current = [i]
|
||||
for j, char_b in enumerate(b, start=1):
|
||||
insert_cost = current[j - 1] + 1
|
||||
delete_cost = previous[j] + 1
|
||||
replace_cost = previous[j - 1] + (char_a != char_b)
|
||||
current.append(min(insert_cost, delete_cost, replace_cost))
|
||||
previous = current
|
||||
return previous[-1]
|
||||
|
||||
|
||||
def _score_single_answer(predicted: Any, target: Any, threshold: float) -> float:
|
||||
predicted_norm = _normalize_text(predicted)
|
||||
target_norm = _normalize_text(target)
|
||||
if not predicted_norm and not target_norm:
|
||||
return 1.0
|
||||
if not predicted_norm or not target_norm:
|
||||
return 0.0
|
||||
distance = _levenshtein_distance(predicted_norm, target_norm)
|
||||
normalized_distance = distance / max(len(predicted_norm), len(target_norm))
|
||||
if normalized_distance >= threshold:
|
||||
return 0.0
|
||||
return 1.0 - normalized_distance
|
||||
|
||||
|
||||
def _extract_answer_strings(raw: Any) -> list[str]:
|
||||
if raw is None:
|
||||
return [""]
|
||||
if isinstance(raw, str):
|
||||
text = raw.strip()
|
||||
if not text:
|
||||
return [""]
|
||||
parsed = None
|
||||
if text[0] in "[{":
|
||||
try:
|
||||
parsed = json.loads(text)
|
||||
except json.JSONDecodeError:
|
||||
try:
|
||||
parsed = ast.literal_eval(text)
|
||||
except (ValueError, SyntaxError):
|
||||
parsed = None
|
||||
if parsed is None:
|
||||
return [text]
|
||||
return _extract_answer_strings(parsed)
|
||||
if isinstance(raw, dict):
|
||||
for key in ("answers", "ground_truth", "answer"):
|
||||
if key in raw:
|
||||
return _extract_answer_strings(raw[key])
|
||||
return [str(raw)]
|
||||
if isinstance(raw, Iterable) and not isinstance(raw, (bytes, bytearray)):
|
||||
answers: list[str] = []
|
||||
for item in raw:
|
||||
if isinstance(item, dict):
|
||||
for key in ("text", "answer", "value"):
|
||||
if key in item:
|
||||
answers.extend(_extract_answer_strings(item[key]))
|
||||
break
|
||||
else:
|
||||
answers.append(str(item))
|
||||
continue
|
||||
answers.append(str(item))
|
||||
return answers or [""]
|
||||
return [str(raw)]
|
||||
|
||||
|
||||
def extract_answer(text: str) -> str:
|
||||
lower = text.lower()
|
||||
start = lower.rfind("<answer>")
|
||||
end = lower.rfind("</answer>")
|
||||
if start != -1 and end != -1 and end > start:
|
||||
return text[start + len("<answer>"):end].strip()
|
||||
lines = [line.strip() for line in text.splitlines() if line.strip()]
|
||||
return lines[-1] if lines else text.strip()
|
||||
|
||||
|
||||
def evaluate(prediction_text: str, gold_answers: Any) -> dict:
|
||||
answer = extract_answer(prediction_text)
|
||||
answers = _extract_answer_strings(gold_answers)
|
||||
score = 0.0
|
||||
for target in answers:
|
||||
score = max(score, _score_single_answer(answer, target, DEFAULT_ANLS_THRESHOLD))
|
||||
return {
|
||||
"anls": score,
|
||||
"predicted_answer": answer,
|
||||
"gold_answers": answers,
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
You are an expert failure-analysis agent for visual document question answering tasks.
|
||||
|
||||
You will be given MULTIPLE failed DocVQA trajectories from a single minibatch and the current skill document. Each trajectory includes the model response and an evaluation result scored with ANLS against one or more acceptable answers.
|
||||
|
||||
Your job is to identify the most important COMMON failure patterns across the batch and propose concise skill edits.
|
||||
|
||||
## Failure Type Categories
|
||||
- evidence_miss: the model overlooked the relevant visible region or line
|
||||
- near_match_confusion: the model selected a nearby but incorrect text span
|
||||
- normalization_error: the answer differed mainly in formatting, spacing, punctuation, or minor text normalization
|
||||
- reading_error: the model misread the document content
|
||||
- other: none of the above
|
||||
|
||||
## Rules
|
||||
- Focus on common, reusable reading and extraction behaviors.
|
||||
- Do not hardcode image-specific answers.
|
||||
- Prefer concise edits that improve evidence selection and exact span extraction.
|
||||
|
||||
Respond ONLY with a valid JSON object (no markdown fences, no extra text):
|
||||
{
|
||||
"batch_size": <number of trajectories analysed>,
|
||||
"failure_summary": [
|
||||
{"failure_type": "<type>", "count": <int>, "description": "<one-line>"}
|
||||
],
|
||||
"patch": {
|
||||
"reasoning": "<why these edits address the batch's common failures>",
|
||||
"edits": [
|
||||
{"op": "append", "content": "<markdown to add at end of skill>"},
|
||||
{"op": "insert_after", "target": "<exact heading/text to insert after>", "content": "<markdown>"},
|
||||
{"op": "replace", "target": "<exact text to replace>", "content": "<replacement>"},
|
||||
{"op": "delete", "target": "<exact text to remove>"}
|
||||
]
|
||||
}
|
||||
}
|
||||
Only include edits that are needed. "edits" can be an empty list if no patch is warranted.
|
||||
@@ -0,0 +1,24 @@
|
||||
You are an expert success-pattern analyst for visual document question answering tasks.
|
||||
|
||||
You will be given MULTIPLE successful DocVQA trajectories from a single minibatch and the current skill document. Your job is to identify common visual reading and exact-answer extraction behaviors worth encoding in the skill.
|
||||
|
||||
## Rules
|
||||
- Focus on patterns shared across multiple successful trajectories.
|
||||
- Reinforce reusable behaviors like locating the right region, copying exact spans, and preferring the shortest exact answer over paraphrase.
|
||||
- Only propose patches for patterns not already captured by the current skill.
|
||||
|
||||
Respond ONLY with a valid JSON object:
|
||||
{
|
||||
"batch_size": <number of trajectories analysed>,
|
||||
"success_patterns": ["<pattern 1>", "<pattern 2>"],
|
||||
"patch": {
|
||||
"reasoning": "<why these patterns are worth encoding>",
|
||||
"edits": [
|
||||
{"op": "append", "content": "<markdown>"},
|
||||
{"op": "insert_after", "target": "<heading/text>", "content": "<markdown>"},
|
||||
{"op": "replace", "target": "<old text>", "content": "<new text>"},
|
||||
{"op": "delete", "target": "<exact text to remove>"}
|
||||
]
|
||||
}
|
||||
}
|
||||
"edits" may be empty if the skill already covers all observed patterns.
|
||||
@@ -0,0 +1,12 @@
|
||||
You are an expert visual document question answering agent.
|
||||
|
||||
{skill_section}You will receive a document image and a question about the document.
|
||||
Read the visual evidence carefully and answer concisely.
|
||||
|
||||
Rules:
|
||||
- Ground the answer in the visible document content.
|
||||
- Prefer exact spans, numbers, dates, and names from the document.
|
||||
- Do not invent content that is not visible.
|
||||
- If multiple near-matches exist, choose the one best supported by the document.
|
||||
|
||||
Return the final answer inside <answer>...</answer>.
|
||||
@@ -0,0 +1,388 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
|
||||
|
||||
from skillopt.envs.docvqa.evaluator import evaluate
|
||||
from skillopt.model import chat_student_messages, get_student_backend, is_student_exec_backend
|
||||
from skillopt.model.codex_harness import prepare_workspace, render_skill_md, run_student_exec
|
||||
from skillopt.prompts import load_prompt
|
||||
|
||||
|
||||
def _build_system(skill_content: str) -> str:
|
||||
if skill_content.strip():
|
||||
skill_section = f"## Skill\n{skill_content.strip()}\n\n"
|
||||
else:
|
||||
skill_section = ""
|
||||
return load_prompt("rollout_system", env="docvqa").format(skill_section=skill_section)
|
||||
|
||||
|
||||
def _image_to_data_uri(path: str) -> str:
|
||||
import base64
|
||||
import mimetypes
|
||||
|
||||
mime = mimetypes.guess_type(path)[0] or "image/png"
|
||||
with open(path, "rb") as f:
|
||||
encoded = base64.b64encode(f.read()).decode("ascii")
|
||||
return f"data:{mime};base64,{encoded}"
|
||||
|
||||
|
||||
def _build_messages(
|
||||
item: dict,
|
||||
skill_content: str,
|
||||
image_detail: str,
|
||||
*,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
) -> tuple[list[dict], str, str]:
|
||||
system = _build_system(skill_content)
|
||||
user_text = item["question"] + "\n\nReturn the final answer inside <answer>...</answer>."
|
||||
if diagnostic_mode and diagnostic_instruction.strip():
|
||||
user_text += f"\n\n## Training Readout\n{diagnostic_instruction.strip()}"
|
||||
image_url = {"url": _image_to_data_uri(item["image_path"])}
|
||||
if image_detail and image_detail != "auto":
|
||||
image_url["detail"] = image_detail
|
||||
messages = [
|
||||
{"role": "system", "content": system},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": user_text},
|
||||
{"type": "image_url", "image_url": image_url},
|
||||
],
|
||||
},
|
||||
]
|
||||
return messages, system, user_text
|
||||
|
||||
|
||||
def _build_codex_skill(skill_content: str) -> str:
|
||||
return render_skill_md(
|
||||
skill_content,
|
||||
description="Dynamic ReflACT skill for solving the current DocVQA document-image question.",
|
||||
preamble=(
|
||||
"Use this skill when answering the current DocVQA question.\n"
|
||||
"Inspect the attached document image carefully and return the final answer inside <answer>...</answer>."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _run_codex_once(
|
||||
*,
|
||||
pred_dir: str,
|
||||
item: dict,
|
||||
skill_content: str,
|
||||
model: str,
|
||||
timeout: int,
|
||||
image_detail: str,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
previous_response: str = "",
|
||||
) -> tuple[str, str, str, str]:
|
||||
_ = image_detail
|
||||
_messages, _system, user_text = _build_messages(
|
||||
item,
|
||||
skill_content,
|
||||
image_detail,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
)
|
||||
task_parts = [user_text]
|
||||
image_abs = os.path.abspath(item["image_path"])
|
||||
task_parts.append(
|
||||
"## Document Image\n"
|
||||
"The document image is available in this workspace via `ATTACHMENTS.md`.\n"
|
||||
f"Original image path: `{image_abs}`\n"
|
||||
"Open or inspect that image before answering; do not answer from memory."
|
||||
)
|
||||
if previous_response:
|
||||
task_parts.append(
|
||||
"## Previous Attempt\n"
|
||||
f"{previous_response}\n\n"
|
||||
"Review the same document image carefully and correct the answer if needed."
|
||||
)
|
||||
task_text = "\n\n".join(task_parts)
|
||||
skill_md = _build_codex_skill(skill_content)
|
||||
work_dir = os.path.join(pred_dir, "codex_exec")
|
||||
prepare_workspace(
|
||||
work_dir=work_dir,
|
||||
skill_md=skill_md,
|
||||
task_text=task_text,
|
||||
images=[item["image_path"]],
|
||||
)
|
||||
prompt = (
|
||||
"Use the `skillopt-student` skill available in this workspace.\n"
|
||||
"Read `task.md`, inspect the attached document image, and answer the DocVQA question.\n"
|
||||
"Return the final answer inside <answer>...</answer>."
|
||||
)
|
||||
final_message, raw = run_student_exec(
|
||||
work_dir=work_dir,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
timeout=timeout,
|
||||
images=[item["image_path"]],
|
||||
)
|
||||
return final_message or raw, raw, skill_md, task_text
|
||||
|
||||
|
||||
def process_one(
|
||||
item: dict,
|
||||
out_root: str,
|
||||
skill_content: str,
|
||||
*,
|
||||
max_turns: int = 1,
|
||||
exec_timeout: int = 120,
|
||||
image_detail: str = "auto",
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
) -> dict:
|
||||
item_id = str(item["id"])
|
||||
result = {
|
||||
"id": item_id,
|
||||
"question": item["question"],
|
||||
"task_type": item.get("subtask") or item.get("task_type") or "docvqa",
|
||||
"task_description": item["question"],
|
||||
"hard": 0,
|
||||
"soft": 0.0,
|
||||
"predicted_answer": "",
|
||||
"response": "",
|
||||
"fail_reason": "",
|
||||
"agent_ok": False,
|
||||
"n_turns": 0,
|
||||
"image_paths": item.get("image_paths", []),
|
||||
"gold_answer": item.get("answers", []),
|
||||
}
|
||||
try:
|
||||
response = ""
|
||||
system_prompt = ""
|
||||
user_text = ""
|
||||
conversation: list[dict] = []
|
||||
if is_student_exec_backend():
|
||||
from skillopt.model import azure_openai as _llm
|
||||
|
||||
conversation = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": item["question"] + "\n\n" + f"[image] {os.path.basename(item['image_path'])}",
|
||||
}
|
||||
]
|
||||
for turn in range(max_turns):
|
||||
response, _raw, system_prompt, user_text = _run_codex_once(
|
||||
pred_dir=os.path.join(out_root, "predictions", item_id),
|
||||
item=item,
|
||||
skill_content=skill_content,
|
||||
model=_llm.STUDENT_DEPLOYMENT,
|
||||
timeout=exec_timeout,
|
||||
image_detail=image_detail,
|
||||
diagnostic_mode=diagnostic_mode if turn == 0 else False,
|
||||
diagnostic_instruction=diagnostic_instruction if turn == 0 else "",
|
||||
previous_response=response if turn > 0 else "",
|
||||
)
|
||||
conversation.append({"type": "message", "turn": turn + 1, "content": response})
|
||||
if "<answer>" in response.lower():
|
||||
break
|
||||
else:
|
||||
messages, system_prompt, user_text = _build_messages(
|
||||
item,
|
||||
skill_content,
|
||||
image_detail,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
)
|
||||
conversation = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": user_text + "\n\n" + f"[image] {os.path.basename(item['image_path'])}",
|
||||
}
|
||||
]
|
||||
for turn in range(max_turns):
|
||||
if turn == 0:
|
||||
resp_text, _ = chat_student_messages(
|
||||
messages=messages,
|
||||
max_completion_tokens=768,
|
||||
retries=5,
|
||||
stage="rollout",
|
||||
timeout=exec_timeout,
|
||||
)
|
||||
else:
|
||||
refinement_messages = [
|
||||
messages[0],
|
||||
messages[1],
|
||||
{"role": "assistant", "content": response},
|
||||
{"role": "user", "content": "Review the same image carefully and answer again. Keep the final answer inside <answer>...</answer>."},
|
||||
]
|
||||
resp_text, _ = chat_student_messages(
|
||||
messages=refinement_messages,
|
||||
max_completion_tokens=512,
|
||||
retries=5,
|
||||
stage="rollout",
|
||||
timeout=exec_timeout,
|
||||
)
|
||||
response = resp_text
|
||||
conversation.append({"type": "message", "turn": turn + 1, "content": resp_text})
|
||||
if "<answer>" in resp_text.lower():
|
||||
break
|
||||
|
||||
result["response"] = response
|
||||
result["agent_ok"] = True
|
||||
result["n_turns"] = len(conversation) - 1
|
||||
|
||||
pred_dir = os.path.join(out_root, "predictions", item_id)
|
||||
os.makedirs(pred_dir, exist_ok=True)
|
||||
with open(os.path.join(pred_dir, "student_system_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(system_prompt)
|
||||
with open(os.path.join(pred_dir, "student_user_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(user_text)
|
||||
|
||||
eval_result = evaluate(response, item.get("answers", []))
|
||||
result["predicted_answer"] = eval_result["predicted_answer"]
|
||||
result["hard"] = int(eval_result["anls"] >= 0.999)
|
||||
result["soft"] = eval_result["anls"]
|
||||
if result["soft"] <= 0.0:
|
||||
result["fail_reason"] = f"predicted '{eval_result['predicted_answer']}' but expected one of {item.get('answers', [])}"
|
||||
|
||||
eval_detail = (
|
||||
"[EVALUATION RESULT]\n"
|
||||
f"Question: {item['question']}\n"
|
||||
f"Predicted answer: {eval_result['predicted_answer']!r}\n"
|
||||
f"Gold answers: {item.get('answers', [])!r}\n"
|
||||
f"ANLS: {eval_result['anls']:.4f}"
|
||||
)
|
||||
conversation.append({"role": "system", "content": eval_detail})
|
||||
with open(os.path.join(pred_dir, "conversation.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(conversation, f, ensure_ascii=False, indent=2)
|
||||
except Exception as e: # noqa: BLE001
|
||||
result["fail_reason"] = f"error: {e}"
|
||||
return result
|
||||
|
||||
|
||||
def run_batch(
|
||||
items: list[dict],
|
||||
out_root: str,
|
||||
skill_content: str,
|
||||
*,
|
||||
max_turns: int = 1,
|
||||
exec_timeout: int = 120,
|
||||
workers: int = 16,
|
||||
image_detail: str = "auto",
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
task_timeout: int = 600,
|
||||
) -> list[dict]:
|
||||
task_timeout = max(int(task_timeout), int(exec_timeout) + 60)
|
||||
results_path = os.path.join(out_root, "results.jsonl")
|
||||
os.makedirs(out_root, exist_ok=True)
|
||||
|
||||
done_ids: set[str] = set()
|
||||
existing: list[dict] = []
|
||||
if os.path.exists(results_path):
|
||||
with open(results_path, encoding="utf-8") as f:
|
||||
for line in f:
|
||||
try:
|
||||
row = json.loads(line)
|
||||
except Exception:
|
||||
continue
|
||||
done_ids.add(str(row["id"]))
|
||||
existing.append(row)
|
||||
|
||||
pending = [item for item in items if str(item["id"]) not in done_ids]
|
||||
if not pending:
|
||||
return existing
|
||||
|
||||
def _timeout_result(item: dict) -> dict:
|
||||
return {
|
||||
"id": str(item["id"]),
|
||||
"question": item.get("question", ""),
|
||||
"task_type": item.get("subtask") or item.get("task_type") or "docvqa",
|
||||
"task_description": item.get("question", ""),
|
||||
"hard": 0,
|
||||
"soft": 0.0,
|
||||
"predicted_answer": "",
|
||||
"response": "",
|
||||
"fail_reason": f"task-timeout-{task_timeout}s",
|
||||
"agent_ok": False,
|
||||
"n_turns": 0,
|
||||
"image_paths": item.get("image_paths", []),
|
||||
"gold_answer": item.get("answers", []),
|
||||
"phase": "timeout",
|
||||
}
|
||||
|
||||
def _error_result(item: dict, exc: Exception) -> dict:
|
||||
row = _timeout_result(item)
|
||||
row["phase"] = "error"
|
||||
row["fail_reason"] = f"unexpected: {type(exc).__name__}: {exc}"
|
||||
return row
|
||||
|
||||
started_at: dict[str, float] = {}
|
||||
|
||||
def _run_one(item: dict) -> dict:
|
||||
started_at[str(item["id"])] = time.time()
|
||||
return process_one(
|
||||
item,
|
||||
out_root,
|
||||
skill_content,
|
||||
max_turns=max_turns,
|
||||
exec_timeout=exec_timeout,
|
||||
image_detail=image_detail,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
)
|
||||
|
||||
total = len(existing) + len(pending)
|
||||
completed = len(existing)
|
||||
correct = sum(1 for r in existing if r.get("hard", 0))
|
||||
if existing:
|
||||
print(f" [rollout] resuming: {completed}/{total} already done", flush=True)
|
||||
|
||||
results = list(existing)
|
||||
with open(results_path, "a", encoding="utf-8") as outf:
|
||||
ex = ThreadPoolExecutor(max_workers=workers)
|
||||
try:
|
||||
futs = {ex.submit(_run_one, item): item for item in pending}
|
||||
pending_futs = set(futs)
|
||||
while pending_futs:
|
||||
done, _ = wait(pending_futs, timeout=5, return_when=FIRST_COMPLETED)
|
||||
now = time.time()
|
||||
timed_out = [
|
||||
fut for fut in pending_futs - done
|
||||
if str(futs[fut]["id"]) in started_at
|
||||
and now - started_at[str(futs[fut]["id"])] >= task_timeout
|
||||
]
|
||||
for fut in done:
|
||||
pending_futs.remove(fut)
|
||||
item = futs[fut]
|
||||
try:
|
||||
res = fut.result()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
res = _error_result(item, exc)
|
||||
results.append(res)
|
||||
completed += 1
|
||||
if res.get("hard", 0):
|
||||
correct += 1
|
||||
acc = correct / completed if completed else 0
|
||||
print(
|
||||
f" [rollout] {completed}/{total} "
|
||||
f"(acc={acc:.3f}) id={res['id']} "
|
||||
f"hard={res.get('hard', '?')}",
|
||||
flush=True,
|
||||
)
|
||||
outf.write(json.dumps(res, ensure_ascii=False) + "\n")
|
||||
outf.flush()
|
||||
for fut in timed_out:
|
||||
pending_futs.remove(fut)
|
||||
fut.cancel()
|
||||
res = _timeout_result(futs[fut])
|
||||
results.append(res)
|
||||
completed += 1
|
||||
acc = correct / completed if completed else 0
|
||||
print(
|
||||
f" [rollout] {completed}/{total} "
|
||||
f"(acc={acc:.3f}) id={res['id']} TIMEOUT",
|
||||
flush=True,
|
||||
)
|
||||
outf.write(json.dumps(res, ensure_ascii=False) + "\n")
|
||||
outf.flush()
|
||||
finally:
|
||||
ex.shutdown(wait=False, cancel_futures=True)
|
||||
return results
|
||||
@@ -0,0 +1,11 @@
|
||||
# DocVQA Skill
|
||||
|
||||
## Visual Evidence Discipline
|
||||
- Read the document carefully before answering.
|
||||
- Prefer the smallest exact text span that answers the question.
|
||||
- When several nearby strings look similar, choose the one whose surrounding labels or layout best match the question.
|
||||
|
||||
## Exact Answer Discipline
|
||||
- Copy names, numbers, and dates exactly from the document whenever possible.
|
||||
- Prefer direct extraction over paraphrase.
|
||||
- Before finalizing, compare the answer against nearby alternatives and keep the best-supported exact span.
|
||||
@@ -0,0 +1 @@
|
||||
"""LiveMathematicianBench environment package for ReflACT."""
|
||||
@@ -0,0 +1,282 @@
|
||||
"""LiveMathematicianBench environment adapter for ReflACT."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
from skillopt.gradient.deep_probe import generate_deep_probe_instruction
|
||||
from skillopt.datasets.base import BatchSpec
|
||||
from skillopt.gradient.reflect import run_minibatch_reflect
|
||||
from skillopt.envs.base import EnvAdapter
|
||||
from skillopt.envs.livemathematicianbench.dataloader import LiveMathematicianBenchDataLoader
|
||||
from skillopt.envs.livemathematicianbench.rollout import run_batch
|
||||
from skillopt.model import get_student_backend
|
||||
|
||||
|
||||
class LiveMathematicianBenchAdapter(EnvAdapter):
|
||||
"""LiveMathematicianBench adapter."""
|
||||
|
||||
def build_reference_text(self, item: dict) -> str:
|
||||
parts: list[str] = []
|
||||
theorem = str(item.get("theorem") or "").strip()
|
||||
sketch = str(item.get("sketch") or "").strip()
|
||||
if theorem:
|
||||
parts.append(f"## Reference Theorem\n{theorem}")
|
||||
if sketch:
|
||||
parts.append(f"## Reference Sketch\n{sketch}")
|
||||
return "\n\n".join(parts)
|
||||
|
||||
def get_reference_metadata(self, item: dict) -> dict:
|
||||
fields: list[str] = []
|
||||
previews: list[str] = []
|
||||
theorem = str(item.get("theorem") or "").strip()
|
||||
sketch = str(item.get("sketch") or "").strip()
|
||||
if theorem:
|
||||
fields.append("theorem")
|
||||
previews.append(f"[theorem]\n{theorem[:220]}")
|
||||
if sketch:
|
||||
fields.append("sketch")
|
||||
previews.append(f"[sketch]\n{sketch[:220]}")
|
||||
return {
|
||||
"fields": fields,
|
||||
"preview": "\n\n".join(previews)[:500],
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
split_dir: str = "",
|
||||
data_path: str = "",
|
||||
split_mode: str = "ratio",
|
||||
split_ratio: str = "2:1:7",
|
||||
split_seed: int = 42,
|
||||
split_output_dir: str = "",
|
||||
max_turns: int = 1,
|
||||
exec_timeout: int = 600,
|
||||
workers: int = 64,
|
||||
analyst_workers: int = 16,
|
||||
failure_only: bool = False,
|
||||
minibatch_size: int = 8,
|
||||
edit_budget: int = 4,
|
||||
seed: int = 42,
|
||||
limit: int = 0,
|
||||
shuffle_choices: bool = True,
|
||||
use_theorem: bool = False,
|
||||
use_sketch: bool = False,
|
||||
use_deep_reflect: bool = False,
|
||||
deep_reflect_failures: int = 4,
|
||||
deep_reflect_successes: int = 2,
|
||||
) -> None:
|
||||
self.max_turns = max_turns
|
||||
self.exec_timeout = exec_timeout
|
||||
self.workers = workers
|
||||
self.analyst_workers = analyst_workers
|
||||
self.failure_only = failure_only
|
||||
self.minibatch_size = minibatch_size
|
||||
self.edit_budget = edit_budget
|
||||
self.use_theorem = use_theorem
|
||||
self.use_sketch = use_sketch
|
||||
self.use_deep_reflect = use_deep_reflect
|
||||
self.deep_reflect_failures = deep_reflect_failures
|
||||
self.deep_reflect_successes = deep_reflect_successes
|
||||
self.dataloader = LiveMathematicianBenchDataLoader(
|
||||
split_dir=split_dir,
|
||||
data_path=data_path,
|
||||
split_mode=split_mode,
|
||||
split_ratio=split_ratio,
|
||||
split_seed=split_seed,
|
||||
split_output_dir=split_output_dir,
|
||||
seed=seed,
|
||||
limit=limit,
|
||||
shuffle_choices=shuffle_choices,
|
||||
)
|
||||
|
||||
def setup(self, cfg: dict) -> None:
|
||||
super().setup(cfg)
|
||||
self.dataloader.setup(cfg)
|
||||
|
||||
def get_dataloader(self):
|
||||
return self.dataloader
|
||||
|
||||
def build_env_from_batch(self, batch: BatchSpec, **kwargs):
|
||||
return list(batch.payload or [])
|
||||
|
||||
def build_train_env(self, batch_size: int, seed: int, **kwargs):
|
||||
batch = self.dataloader.build_train_batch(batch_size=batch_size, seed=seed, **kwargs)
|
||||
return self.build_env_from_batch(batch, **kwargs)
|
||||
|
||||
def build_eval_env(self, env_num: int, split: str, seed: int, **kwargs):
|
||||
batch = self.dataloader.build_eval_batch(env_num=env_num, split=split, seed=seed, **kwargs)
|
||||
return self.build_env_from_batch(batch, **kwargs)
|
||||
|
||||
def rollout(
|
||||
self,
|
||||
env_manager,
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict]:
|
||||
items: list[dict] = env_manager
|
||||
return run_batch(
|
||||
items=items,
|
||||
out_root=out_dir,
|
||||
skill_content=skill_content,
|
||||
max_turns=self.max_turns,
|
||||
exec_timeout=self.exec_timeout,
|
||||
workers=self.workers,
|
||||
use_theorem=self.use_theorem,
|
||||
use_sketch=self.use_sketch,
|
||||
diagnostic_mode=kwargs.get("diagnostic_mode", False),
|
||||
diagnostic_instruction=kwargs.get("diagnostic_instruction", ""),
|
||||
diagnostic_trace_context_by_id=kwargs.get("diagnostic_trace_context_by_id"),
|
||||
task_timeout=self.exec_timeout,
|
||||
)
|
||||
|
||||
def reflect(
|
||||
self,
|
||||
results: list[dict],
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict | None]:
|
||||
prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
|
||||
patches_dir = kwargs.get("patches_dir", os.path.join(out_dir, "patches"))
|
||||
random_seed = kwargs.get("random_seed")
|
||||
step_buffer_context = kwargs.get("step_buffer_context", "")
|
||||
meta_skill_context = kwargs.get("meta_skill_context", "")
|
||||
|
||||
return run_minibatch_reflect(
|
||||
results=results,
|
||||
skill_content=skill_content,
|
||||
prediction_dir=prediction_dir,
|
||||
patches_dir=patches_dir,
|
||||
workers=self.analyst_workers,
|
||||
failure_only=self.failure_only,
|
||||
minibatch_size=self.minibatch_size,
|
||||
edit_budget=self.edit_budget,
|
||||
random_seed=random_seed,
|
||||
error_system=self.get_error_minibatch_prompt(),
|
||||
success_system=self.get_success_minibatch_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
meta_skill_context=meta_skill_context,
|
||||
update_mode=getattr(self, "_cfg", {}).get("skill_update_mode", "patch"),
|
||||
)
|
||||
|
||||
def deep_reflect(
|
||||
self,
|
||||
results: list[dict],
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict | None]:
|
||||
if not self.use_deep_reflect:
|
||||
return []
|
||||
|
||||
env_manager = kwargs.get("env_manager")
|
||||
prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
|
||||
random_seed = kwargs.get("random_seed")
|
||||
step_buffer_context = kwargs.get("step_buffer_context", "")
|
||||
meta_skill_context = kwargs.get("meta_skill_context", "")
|
||||
codex_backend = get_student_backend() == "codex_exec"
|
||||
selected_items = self.select_representative_items(
|
||||
results,
|
||||
env_manager if isinstance(env_manager, list) else None,
|
||||
n_failures=self.deep_reflect_failures,
|
||||
n_successes=self.deep_reflect_successes,
|
||||
seed=random_seed,
|
||||
)
|
||||
if not selected_items:
|
||||
return []
|
||||
selected_ids = {str(item["id"]) for item in selected_items}
|
||||
selected_results = [row for row in results if str(row.get("id")) in selected_ids]
|
||||
selected_examples = self.attach_reference_context(selected_results, selected_items)
|
||||
if codex_backend:
|
||||
selected_examples = self.attach_codex_probe_context(selected_examples, prediction_dir)
|
||||
selected_metadata = []
|
||||
theorem_count = 0
|
||||
sketch_count = 0
|
||||
for item in selected_items:
|
||||
meta = self.get_reference_metadata(item)
|
||||
if "theorem" in meta["fields"]:
|
||||
theorem_count += 1
|
||||
if "sketch" in meta["fields"]:
|
||||
sketch_count += 1
|
||||
selected_metadata.append({
|
||||
"id": str(item["id"]),
|
||||
"task_type": str(item.get("theorem_type", ["math_mcq"])[0] if item.get("theorem_type") else "math_mcq"),
|
||||
"reference_fields": meta["fields"],
|
||||
"reference_preview": meta["preview"],
|
||||
})
|
||||
|
||||
deep_dir = os.path.join(out_dir, "deep_reflect")
|
||||
rollout_dir = os.path.join(deep_dir, "rollout")
|
||||
patches_dir = os.path.join(deep_dir, "patches")
|
||||
os.makedirs(deep_dir, exist_ok=True)
|
||||
print(
|
||||
f" [2b/6 DEEP REFLECT setup] selected={len(selected_items)} "
|
||||
f"reference_fields=theorem({theorem_count}/{len(selected_items)}),"
|
||||
f"sketch({sketch_count}/{len(selected_items)})"
|
||||
)
|
||||
probe = generate_deep_probe_instruction(
|
||||
skill_content=skill_content,
|
||||
items=selected_examples,
|
||||
prediction_dir=prediction_dir,
|
||||
system_prompt=self.get_codex_deep_probe_prompt() if codex_backend else self.get_deep_probe_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
meta_skill_context=meta_skill_context,
|
||||
)
|
||||
if not probe:
|
||||
return []
|
||||
diagnostic_trace_context_by_id = None
|
||||
if codex_backend:
|
||||
selected_items, diagnostic_trace_context_by_id, probe = self.resolve_codex_probe_target(
|
||||
selected_items=selected_items,
|
||||
selected_examples=selected_examples,
|
||||
prediction_dir=prediction_dir,
|
||||
probe=probe,
|
||||
)
|
||||
probe_record = {
|
||||
**probe,
|
||||
"reference_summary": {
|
||||
"selected_count": len(selected_items),
|
||||
"field_counts": {
|
||||
"theorem": theorem_count,
|
||||
"sketch": sketch_count,
|
||||
},
|
||||
},
|
||||
"selected_examples": selected_metadata,
|
||||
}
|
||||
with open(os.path.join(deep_dir, "probe.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(probe_record, f, ensure_ascii=False, indent=2)
|
||||
deep_results = run_batch(
|
||||
items=selected_items,
|
||||
out_root=rollout_dir,
|
||||
skill_content=skill_content,
|
||||
max_turns=self.max_turns,
|
||||
workers=min(self.workers, max(len(selected_items), 1)),
|
||||
use_theorem=self.use_theorem,
|
||||
use_sketch=self.use_sketch,
|
||||
diagnostic_mode=True,
|
||||
diagnostic_instruction=probe["probe_instruction"],
|
||||
diagnostic_trace_context_by_id=diagnostic_trace_context_by_id,
|
||||
task_timeout=self.exec_timeout,
|
||||
)
|
||||
deep_results = self.attach_reference_context(deep_results, selected_items)
|
||||
return run_minibatch_reflect(
|
||||
results=deep_results,
|
||||
skill_content=skill_content,
|
||||
prediction_dir=os.path.join(rollout_dir, "predictions"),
|
||||
patches_dir=patches_dir,
|
||||
workers=self.analyst_workers,
|
||||
failure_only=self.failure_only,
|
||||
minibatch_size=self.minibatch_size,
|
||||
edit_budget=self.edit_budget,
|
||||
random_seed=random_seed,
|
||||
error_system=self.get_error_minibatch_prompt(),
|
||||
success_system=self.get_success_minibatch_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
meta_skill_context=meta_skill_context,
|
||||
update_mode=getattr(self, "_cfg", {}).get("skill_update_mode", "patch"),
|
||||
)
|
||||
|
||||
def get_task_types(self) -> list[str]:
|
||||
return self.dataloader.get_task_types()
|
||||
@@ -0,0 +1,308 @@
|
||||
"""LiveMathematicianBench task dataloader."""
|
||||
from __future__ import annotations
|
||||
|
||||
import glob
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
from typing import Any
|
||||
|
||||
from skillopt.datasets.base import BatchSpec, SplitDataLoader
|
||||
|
||||
|
||||
# ── Raw data loading utilities (for preprocessing / standalone eval) ─────
|
||||
|
||||
_CHOICE_LABELS = ["A", "B", "C", "D", "E", "F", "G"]
|
||||
|
||||
|
||||
def _load_json(path: str) -> Any:
|
||||
with open(path) as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def _iter_monthly_files(data_path: str) -> list[str]:
|
||||
if not data_path:
|
||||
return []
|
||||
if os.path.isfile(data_path):
|
||||
return [data_path]
|
||||
if os.path.isdir(data_path):
|
||||
nested = glob.glob(
|
||||
os.path.join(data_path, "**", "qa_*_final.json"),
|
||||
recursive=True,
|
||||
)
|
||||
flat = glob.glob(os.path.join(data_path, "qa_*_final.json"))
|
||||
return sorted(set(nested + flat))
|
||||
return []
|
||||
|
||||
|
||||
def _coerce_choices(raw_choices: Any) -> list[dict]:
|
||||
if isinstance(raw_choices, list):
|
||||
choices: list[dict] = []
|
||||
for idx, item in enumerate(raw_choices):
|
||||
if isinstance(item, dict):
|
||||
label = str(item.get("label") or _CHOICE_LABELS[idx]).strip()
|
||||
text = str(item.get("text") or item.get("content") or "").strip()
|
||||
else:
|
||||
label = _CHOICE_LABELS[idx]
|
||||
text = str(item).strip()
|
||||
if text:
|
||||
choices.append({"label": label, "text": text})
|
||||
return choices
|
||||
|
||||
if isinstance(raw_choices, dict):
|
||||
labels = sorted(raw_choices.keys())
|
||||
return [
|
||||
{"label": str(label).strip(), "text": str(raw_choices[label]).strip()}
|
||||
for label in labels
|
||||
if str(raw_choices[label]).strip()
|
||||
]
|
||||
|
||||
return []
|
||||
|
||||
|
||||
def _coerce_theorem_types(raw: Any) -> list[str]:
|
||||
if isinstance(raw, list):
|
||||
return [str(x).strip() for x in raw if str(x).strip()]
|
||||
if raw is None:
|
||||
return []
|
||||
text = str(raw).strip()
|
||||
return [text] if text else []
|
||||
|
||||
|
||||
def _normalize_label(text: str) -> str:
|
||||
return str(text).strip().upper().rstrip(".):")
|
||||
|
||||
|
||||
def _normalize_item(item: dict, row_idx: int, source_path: str) -> dict:
|
||||
mcq = item.get("mcq", {}) if isinstance(item.get("mcq"), dict) else {}
|
||||
question = str(mcq.get("question") or item.get("question") or "").strip()
|
||||
choices = _coerce_choices(mcq.get("choices") or item.get("choices") or [])
|
||||
correct = mcq.get("correct_choice") or item.get("correct_choice") or {}
|
||||
|
||||
if isinstance(correct, dict):
|
||||
correct_label = _normalize_label(correct.get("label", ""))
|
||||
correct_text = str(correct.get("text") or "").strip()
|
||||
else:
|
||||
correct_label = _normalize_label(correct)
|
||||
correct_text = ""
|
||||
|
||||
choice_by_label = {
|
||||
_normalize_label(choice["label"]): choice["text"]
|
||||
for choice in choices
|
||||
}
|
||||
if correct_label and not correct_text:
|
||||
correct_text = choice_by_label.get(correct_label, "")
|
||||
if correct_label and correct_text and correct_label not in choice_by_label:
|
||||
choices.append({"label": correct_label, "text": correct_text})
|
||||
choices.sort(key=lambda choice: _CHOICE_LABELS.index(choice["label"]) if choice["label"] in _CHOICE_LABELS else len(_CHOICE_LABELS))
|
||||
choice_by_label[correct_label] = correct_text
|
||||
|
||||
month = str(item.get("month") or "").strip()
|
||||
item_no = item.get("no", row_idx + 1)
|
||||
item_id = f"{month}:{item_no}" if month else str(item_no)
|
||||
|
||||
return {
|
||||
"id": item_id,
|
||||
"month": month,
|
||||
"no": item_no,
|
||||
"paper_link": str(item.get("paper_link") or "").strip(),
|
||||
"theorem": str(item.get("theorem") or "").strip(),
|
||||
"sketch": str(item.get("sketch") or "").strip(),
|
||||
"theorem_type": _coerce_theorem_types(item.get("theorem_type")),
|
||||
"question": question,
|
||||
"choices": choices,
|
||||
"correct_choice": {
|
||||
"label": correct_label,
|
||||
"text": correct_text,
|
||||
},
|
||||
"source_path": source_path,
|
||||
}
|
||||
|
||||
|
||||
def load_items(data_path: str) -> list[dict]:
|
||||
"""Load and normalise LiveMathematicianBench items from JSON files."""
|
||||
files = _iter_monthly_files(data_path)
|
||||
if not files:
|
||||
raise ValueError(
|
||||
"LiveMathematicianBench requires data_path to be a qa_*_final.json file "
|
||||
"or a directory containing monthly qa_*_final.json files."
|
||||
)
|
||||
|
||||
items: list[dict] = []
|
||||
for path in files:
|
||||
raw = _load_json(path)
|
||||
if not isinstance(raw, list):
|
||||
raise ValueError(f"Expected JSON array in {path}, got {type(raw).__name__}")
|
||||
for row_idx, item in enumerate(raw):
|
||||
norm = _normalize_item(item, row_idx=row_idx, source_path=path)
|
||||
if norm["question"] and norm["choices"] and norm["correct_choice"]["label"]:
|
||||
items.append(norm)
|
||||
if not items:
|
||||
raise ValueError(f"No valid LiveMathematicianBench items loaded from {data_path}")
|
||||
return items
|
||||
|
||||
|
||||
# ── Dataloader ───────────────────────────────────────────────────────────
|
||||
|
||||
class LiveMathematicianBenchDataLoader(SplitDataLoader):
|
||||
"""LiveMathematicianBench dataloader with per-seed choice shuffling."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
split_dir: str = "",
|
||||
data_path: str = "",
|
||||
split_mode: str = "ratio",
|
||||
split_ratio: str = "2:1:7",
|
||||
split_seed: int = 42,
|
||||
split_output_dir: str = "",
|
||||
seed: int = 42,
|
||||
limit: int = 0,
|
||||
shuffle_choices: bool = True,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
split_dir=split_dir,
|
||||
data_path=data_path,
|
||||
split_mode=split_mode,
|
||||
split_ratio=split_ratio,
|
||||
split_seed=split_seed,
|
||||
split_output_dir=split_output_dir,
|
||||
seed=seed,
|
||||
limit=limit,
|
||||
)
|
||||
self.shuffle_choices = shuffle_choices
|
||||
self._task_types: list[str] = []
|
||||
|
||||
def load_raw_items(self, data_path: str) -> list[dict]:
|
||||
return load_items(data_path)
|
||||
|
||||
def setup(self, cfg: dict) -> None:
|
||||
super().setup(cfg)
|
||||
all_items = self.train_items + self.val_items + self.test_items
|
||||
task_types: set[str] = set()
|
||||
for item in all_items:
|
||||
for name in item.get("theorem_type", []):
|
||||
if name:
|
||||
task_types.add(name)
|
||||
self._task_types = sorted(task_types)
|
||||
|
||||
def get_task_types(self) -> list[str]:
|
||||
return list(self._task_types)
|
||||
|
||||
# ── Choice shuffling ─────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _item_shuffle_seed(item_id: str, seed: int) -> int:
|
||||
digest = hashlib.sha256(f"{seed}:{item_id}".encode("utf-8")).hexdigest()
|
||||
return int(digest[:16], 16)
|
||||
|
||||
def _shuffle_item_choices(self, item: dict, seed: int) -> dict:
|
||||
if not self.shuffle_choices:
|
||||
return {
|
||||
**item,
|
||||
"choices": [dict(c) for c in item["choices"]],
|
||||
"correct_choice": dict(item["correct_choice"]),
|
||||
}
|
||||
|
||||
shuffled_choices = [dict(c) for c in item["choices"]]
|
||||
rng = random.Random(self._item_shuffle_seed(str(item["id"]), seed))
|
||||
rng.shuffle(shuffled_choices)
|
||||
|
||||
original_correct = _normalize_label(item["correct_choice"]["label"])
|
||||
remapped_choices: list[dict] = []
|
||||
new_correct_choice = dict(item["correct_choice"])
|
||||
|
||||
for idx, choice in enumerate(shuffled_choices):
|
||||
new_label = _CHOICE_LABELS[idx]
|
||||
old_label = _normalize_label(choice["label"])
|
||||
remapped_choices.append({"label": new_label, "text": choice["text"]})
|
||||
if old_label == original_correct:
|
||||
new_correct_choice = {"label": new_label, "text": choice["text"]}
|
||||
|
||||
transformed = dict(item)
|
||||
transformed["choices"] = remapped_choices
|
||||
transformed["correct_choice"] = new_correct_choice
|
||||
return transformed
|
||||
|
||||
def _materialize_batch(self, items: list[dict], seed: int) -> list[dict]:
|
||||
return [self._shuffle_item_choices(item, seed) for item in items]
|
||||
|
||||
# ── Batch construction (override for choice shuffling) ───────────────
|
||||
|
||||
def plan_train_epoch(
|
||||
self,
|
||||
*,
|
||||
epoch: int,
|
||||
steps_per_epoch: int,
|
||||
accumulation: int,
|
||||
batch_size: int,
|
||||
seed: int,
|
||||
**kwargs,
|
||||
) -> list[BatchSpec]:
|
||||
"""Build a shuffled epoch while preserving per-batch choice shuffling."""
|
||||
epoch_rng = random.Random(seed + epoch * 1000)
|
||||
items = list(self.train_items)
|
||||
epoch_rng.shuffle(items)
|
||||
|
||||
total_batches = steps_per_epoch * accumulation
|
||||
if total_batches <= 0:
|
||||
return []
|
||||
|
||||
batches: list[BatchSpec] = []
|
||||
cursor = 0
|
||||
for batch_idx in range(total_batches):
|
||||
batch_seed = seed + epoch * 1000 + batch_idx + 1
|
||||
batch_items = items[cursor: cursor + batch_size]
|
||||
cursor += len(batch_items)
|
||||
|
||||
if not batch_items and items:
|
||||
refill_rng = random.Random(batch_seed)
|
||||
batch_items = list(items)
|
||||
refill_rng.shuffle(batch_items)
|
||||
batch_items = batch_items[:batch_size]
|
||||
|
||||
batch_items = self._materialize_batch(batch_items, batch_seed)
|
||||
batches.append(
|
||||
BatchSpec(
|
||||
phase="train",
|
||||
split="train",
|
||||
seed=batch_seed,
|
||||
batch_size=len(batch_items),
|
||||
payload=batch_items,
|
||||
)
|
||||
)
|
||||
|
||||
return batches
|
||||
|
||||
def build_train_batch(self, batch_size: int, seed: int, **kwargs) -> BatchSpec:
|
||||
rng = random.Random(seed)
|
||||
items = list(self.train_items)
|
||||
rng.shuffle(items)
|
||||
items = self._materialize_batch(items[:batch_size], seed)
|
||||
return BatchSpec(
|
||||
phase="train",
|
||||
split="train",
|
||||
seed=seed,
|
||||
batch_size=len(items),
|
||||
payload=items,
|
||||
)
|
||||
|
||||
def build_eval_batch(
|
||||
self,
|
||||
env_num: int,
|
||||
split: str,
|
||||
seed: int,
|
||||
**kwargs,
|
||||
) -> BatchSpec:
|
||||
items = self.get_split_items(split)
|
||||
if env_num and env_num < len(items):
|
||||
items = items[:env_num]
|
||||
items = self._materialize_batch(items, seed)
|
||||
return BatchSpec(
|
||||
phase="eval",
|
||||
split=split,
|
||||
seed=seed,
|
||||
batch_size=len(items),
|
||||
payload=items,
|
||||
)
|
||||
@@ -0,0 +1,62 @@
|
||||
"""LiveMathematicianBench evaluation helpers."""
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
|
||||
def extract_answer(text: str) -> str:
|
||||
matches = re.findall(r"<answer>(.*?)</answer>", text, re.DOTALL | re.IGNORECASE)
|
||||
if matches:
|
||||
return matches[-1].strip()
|
||||
lines = [ln.strip() for ln in text.strip().splitlines() if ln.strip()]
|
||||
if lines:
|
||||
return lines[-1]
|
||||
return text.strip()
|
||||
|
||||
|
||||
def normalize_label(text: str) -> str:
|
||||
return str(text).strip().upper().rstrip(".):")
|
||||
|
||||
|
||||
def parse_choice_label(prediction_text: str, choices: list[dict]) -> str:
|
||||
answer = extract_answer(prediction_text)
|
||||
label = normalize_label(answer)
|
||||
valid_labels = {normalize_label(choice.get("label", "")) for choice in choices}
|
||||
if label in valid_labels:
|
||||
return label
|
||||
|
||||
answer_lower = answer.lower()
|
||||
for choice in choices:
|
||||
choice_label = normalize_label(choice.get("label", ""))
|
||||
choice_text = str(choice.get("text", "")).strip()
|
||||
if choice_text and choice_text.lower() == answer_lower:
|
||||
return choice_label
|
||||
|
||||
first_token = normalize_label(answer.split()[0]) if answer.split() else ""
|
||||
if first_token in valid_labels:
|
||||
return first_token
|
||||
return label
|
||||
|
||||
|
||||
def evaluate(prediction_text: str, correct_choice: dict, choices: list[dict]) -> dict:
|
||||
predicted_label = parse_choice_label(prediction_text, choices)
|
||||
correct_label = normalize_label(correct_choice.get("label", ""))
|
||||
predicted_text = ""
|
||||
correct_text = str(correct_choice.get("text", "")).strip()
|
||||
|
||||
for choice in choices:
|
||||
if normalize_label(choice.get("label", "")) == predicted_label:
|
||||
predicted_text = str(choice.get("text", "")).strip()
|
||||
break
|
||||
|
||||
is_correct = float(predicted_label == correct_label)
|
||||
return {
|
||||
"em": is_correct,
|
||||
"f1": is_correct,
|
||||
"sub_em": is_correct,
|
||||
"predicted_answer": predicted_label or extract_answer(prediction_text),
|
||||
"predicted_label": predicted_label,
|
||||
"predicted_text": predicted_text,
|
||||
"correct_label": correct_label,
|
||||
"correct_text": correct_text,
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
You are an expert failure-analysis agent for theorem-grounded mathematical multiple-choice questions.
|
||||
|
||||
You will be given MULTIPLE failed trajectories from a single minibatch and the current skill document.
|
||||
Each trajectory includes the student's response and an evaluation result showing the predicted option
|
||||
versus the correct option.
|
||||
|
||||
Your job is to identify COMMON reasoning failures across the batch and propose concise skill edits.
|
||||
|
||||
## Failure Type Categories
|
||||
- **quantifier_miss**: the agent missed exact quantifiers, scope, or existence/uniqueness conditions
|
||||
- **strength_mismatch**: the agent preferred a weaker or stronger statement than what was proved
|
||||
- **condition_miss**: the agent ignored hypotheses, equality cases, or domain restrictions
|
||||
- **option_confusion**: the agent confused similar answer choices or failed to compare them exactly
|
||||
- **other**: none of the above
|
||||
|
||||
## Rules
|
||||
1. Focus on patterns that recur across the minibatch.
|
||||
2. Prefer edits that improve exact choice discrimination, not theorem-specific memorization.
|
||||
3. Do not hardcode paper-specific content.
|
||||
4. Only patch gaps not already covered by the skill.
|
||||
|
||||
Respond ONLY with a valid JSON object:
|
||||
{
|
||||
"batch_size": <number>,
|
||||
"failure_summary": [
|
||||
{"failure_type": "<type>", "count": <int>, "description": "<one-line>"}
|
||||
],
|
||||
"patch": {
|
||||
"reasoning": "<why these edits address the common failures>",
|
||||
"edits": [
|
||||
{"op": "append", "content": "<markdown>"},
|
||||
{"op": "insert_after", "target": "<heading/text>", "content": "<markdown>"},
|
||||
{"op": "replace", "target": "<old text>", "content": "<new text>"},
|
||||
{"op": "delete", "target": "<exact text to remove>"}
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
You are an expert success-pattern analyst for theorem-grounded mathematical multiple-choice questions.
|
||||
|
||||
You will be given MULTIPLE successful trajectories from a minibatch and the current skill document.
|
||||
Identify generalizable behavior patterns that are genuinely helping the agent choose the exact correct option.
|
||||
|
||||
## Rules
|
||||
- Focus on broadly useful reasoning behaviors.
|
||||
- Prefer patterns about exact comparison of options, quantifiers, and equality conditions.
|
||||
- Do not add theorem-specific facts.
|
||||
- "edits" may be empty if the skill already captures the useful patterns.
|
||||
|
||||
Respond ONLY with a valid JSON object:
|
||||
{
|
||||
"batch_size": <number>,
|
||||
"success_patterns": ["<pattern 1>", "<pattern 2>"],
|
||||
"patch": {
|
||||
"reasoning": "<why these patterns matter>",
|
||||
"edits": [
|
||||
{"op": "append", "content": "<markdown>"},
|
||||
{"op": "insert_after", "target": "<heading/text>", "content": "<markdown>"},
|
||||
{"op": "replace", "target": "<old text>", "content": "<new text>"},
|
||||
{"op": "delete", "target": "<exact text to remove>"}
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
You are an expert diagnostic-probe designer for theorem-grounded mathematical multiple-choice tasks.
|
||||
|
||||
You will be shown representative trajectories, the current student skill, and the student's original prompt context.
|
||||
Design one SMALL diagnostic instruction that exposes the student's intermediate judgment without materially changing the original scaffold.
|
||||
|
||||
## Hard Constraints
|
||||
1. Do NOT substantially change the original scaffold.
|
||||
2. Do NOT prescribe a new multi-step theorem-solving procedure.
|
||||
3. Do NOT ask for a full proof, full chain-of-thought, or exhaustive option-by-option derivation.
|
||||
4. Ask only for a short readout of the signals already behind the student's current answer.
|
||||
5. Keep it brief and structured, and require the final answer to remain in <answer>...</answer>.
|
||||
|
||||
## Good Probe Targets
|
||||
- top choice and runner-up
|
||||
- decisive constraint
|
||||
- why the runner-up was rejected
|
||||
- strongest-vs-weaker discrimination signal
|
||||
|
||||
Respond ONLY with a valid JSON object:
|
||||
{
|
||||
"reasoning": "<why this probe is informative>",
|
||||
"probe_instruction": "<the exact instruction text to append to the student prompt>"
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
You are an expert diagnostic-probe designer for theorem-grounded mathematical multiple-choice tasks executed through a Codex trace.
|
||||
|
||||
You will be shown representative trajectories, the current student skill, the student's original prompt context, hidden reference fields, and numbered Codex trace steps.
|
||||
Choose exactly one trajectory and one probe point. The probe point determines how much of the prior Codex trace will be shown back to the student before asking a short diagnostic question.
|
||||
|
||||
## Hard Constraints
|
||||
1. Do NOT reveal or paraphrase the hidden reference directly to the student.
|
||||
2. Do NOT prescribe a new full solving procedure.
|
||||
3. Do NOT ask for a full proof, full chain-of-thought, or exhaustive option-by-option derivation.
|
||||
4. Ask only for a short readout of the signal that should already exist at that point in the student's process.
|
||||
5. The probe instruction must explicitly request a short <analysis>...</analysis> block before the final <answer>...</answer>.
|
||||
6. Select a probe point that is informative about theorem choice, decisive constraint, option elimination, or why a stronger/weaker option should be rejected.
|
||||
|
||||
## Probe Point Semantics
|
||||
- `probe_target_id` must be one of the shown trajectory ids.
|
||||
- `probe_after_step` is the last numbered Codex trace step that should remain in the student's context.
|
||||
- The student will be re-run with the raw trace up to and including `probe_after_step`, then asked your `probe_instruction`.
|
||||
- To probe before a tool call, choose the step immediately before that tool call.
|
||||
|
||||
Respond ONLY with a valid JSON object:
|
||||
{
|
||||
"reasoning": "<why this trajectory and probe point expose the student's intermediate state>",
|
||||
"probe_target_id": "<trajectory id>",
|
||||
"probe_after_step": <integer step number>,
|
||||
"probe_instruction": "<the exact instruction text to append to the student's prompt>"
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
You are an expert mathematical reasoning agent solving multiple-choice questions.
|
||||
|
||||
{skill_section}## Task Format
|
||||
You will receive one mathematics multiple-choice question and its answer choices.
|
||||
Reason carefully about quantifiers, hypotheses, extremal wording, and exact equality conditions.
|
||||
|
||||
## Answer Format
|
||||
Think step by step, then provide your final answer inside <answer>...</answer> tags.
|
||||
Inside the tags, output only the single choice label, such as A or C.
|
||||
|
||||
Example:
|
||||
<answer>B</answer>
|
||||
@@ -0,0 +1,4 @@
|
||||
"""LiveMathematicianBench Reflect stage.
|
||||
|
||||
Prompts are now loaded from .md files by the base adapter.
|
||||
"""
|
||||
@@ -0,0 +1,424 @@
|
||||
"""LiveMathematicianBench rollout — theorem-grounded math MCQ agent."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
|
||||
|
||||
from skillopt.envs.livemathematicianbench.evaluator import evaluate
|
||||
from skillopt.model import chat_student, get_student_backend, is_student_exec_backend
|
||||
from skillopt.model.codex_harness import prepare_workspace, render_skill_md, run_student_exec
|
||||
from skillopt.prompts import load_prompt
|
||||
|
||||
def _build_system(skill_content: str) -> str:
|
||||
if skill_content.strip():
|
||||
skill_section = f"## Skill\n{skill_content.strip()}\n\n"
|
||||
else:
|
||||
skill_section = ""
|
||||
return load_prompt("rollout_system", env="livemathematicianbench").format(skill_section=skill_section)
|
||||
|
||||
|
||||
def _format_choices(choices: list[dict]) -> str:
|
||||
return "\n".join(
|
||||
f"{choice['label']}. {choice['text']}"
|
||||
for choice in choices
|
||||
)
|
||||
|
||||
|
||||
def _build_user(
|
||||
item: dict,
|
||||
*,
|
||||
use_theorem: bool = False,
|
||||
use_sketch: bool = False,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context: str = "",
|
||||
) -> str:
|
||||
parts = [f"## Question\n{item['question']}", f"## Choices\n{_format_choices(item['choices'])}"]
|
||||
if use_theorem and item.get("theorem"):
|
||||
parts.append(f"## Theorem\n{item['theorem']}")
|
||||
if use_sketch and item.get("sketch"):
|
||||
parts.append(f"## Proof Sketch\n{item['sketch']}")
|
||||
if diagnostic_trace_context.strip():
|
||||
parts.append(
|
||||
"## Previous Codex Trace Snapshot\n"
|
||||
"This is a partial transcript from an earlier attempt. Use it as your current reasoning context.\n\n"
|
||||
f"{diagnostic_trace_context.strip()}"
|
||||
)
|
||||
if diagnostic_mode and diagnostic_instruction.strip():
|
||||
parts.append(f"## Training Readout\n{diagnostic_instruction.strip()}")
|
||||
return "\n\n".join(parts)
|
||||
|
||||
|
||||
def _build_codex_skill(skill_content: str) -> str:
|
||||
return render_skill_md(
|
||||
skill_content,
|
||||
description="Dynamic ReflACT skill for solving the current LiveMathematicianBench multiple-choice question.",
|
||||
preamble=(
|
||||
"Use this skill when solving the current math multiple-choice question.\n"
|
||||
"Inspect the option wording carefully and output only the final choice label inside <answer>...</answer>."
|
||||
),
|
||||
)
|
||||
|
||||
def _run_codex_once(
|
||||
*,
|
||||
pred_dir: str,
|
||||
skill_content: str,
|
||||
item: dict,
|
||||
model: str,
|
||||
timeout: int,
|
||||
use_theorem: bool = False,
|
||||
use_sketch: bool = False,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context: str = "",
|
||||
previous_response: str = "",
|
||||
) -> tuple[str, str, str, str]:
|
||||
user = _build_user(
|
||||
item,
|
||||
use_theorem=use_theorem,
|
||||
use_sketch=use_sketch,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
diagnostic_trace_context=diagnostic_trace_context,
|
||||
)
|
||||
task_parts = [user]
|
||||
if previous_response:
|
||||
task_parts.append(
|
||||
"## Previous Attempt\n"
|
||||
f"{previous_response}\n\n"
|
||||
"Re-evaluate the exact option wording. If needed, correct it."
|
||||
)
|
||||
task_text = "\n\n".join(task_parts)
|
||||
skill_md = _build_codex_skill(skill_content)
|
||||
work_dir = os.path.join(pred_dir, "codex_exec")
|
||||
prepare_workspace(work_dir=work_dir, skill_md=skill_md, task_text=task_text)
|
||||
prompt = (
|
||||
"Use the `skillopt-student` skill available in this workspace.\n"
|
||||
"Read `task.md` and solve the multiple-choice problem.\n"
|
||||
"Output only the final choice label inside <answer>...</answer>."
|
||||
)
|
||||
final_message, raw = run_student_exec(
|
||||
work_dir=work_dir,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
timeout=timeout,
|
||||
)
|
||||
return final_message or raw, raw, skill_md, task_text
|
||||
|
||||
|
||||
def process_one(
|
||||
item: dict,
|
||||
out_root: str,
|
||||
skill_content: str,
|
||||
*,
|
||||
max_turns: int = 1,
|
||||
use_theorem: bool = False,
|
||||
use_sketch: bool = False,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context: str = "",
|
||||
exec_timeout: int = 300,
|
||||
) -> dict:
|
||||
item_id = str(item["id"])
|
||||
result = {
|
||||
"id": item_id,
|
||||
"question": item["question"],
|
||||
"task_type": item.get("theorem_type", ["math_mcq"])[0] if item.get("theorem_type") else "math_mcq",
|
||||
"hard": 0,
|
||||
"soft": 0.0,
|
||||
"predicted_answer": "",
|
||||
"predicted_label": "",
|
||||
"predicted_text": "",
|
||||
"correct_label": item["correct_choice"]["label"],
|
||||
"correct_text": item["correct_choice"]["text"],
|
||||
"response": "",
|
||||
"fail_reason": "",
|
||||
"agent_ok": False,
|
||||
"n_turns": 0,
|
||||
}
|
||||
|
||||
try:
|
||||
pred_dir = os.path.join(out_root, "predictions", item_id)
|
||||
os.makedirs(pred_dir, exist_ok=True)
|
||||
|
||||
if is_student_exec_backend():
|
||||
from skillopt.model import azure_openai as _llm
|
||||
|
||||
conversation: list[dict] = []
|
||||
response = ""
|
||||
system = ""
|
||||
user = ""
|
||||
for turn in range(max_turns):
|
||||
response, raw, system, user = _run_codex_once(
|
||||
pred_dir=pred_dir,
|
||||
skill_content=skill_content,
|
||||
item=item,
|
||||
model=_llm.STUDENT_DEPLOYMENT,
|
||||
timeout=exec_timeout,
|
||||
use_theorem=use_theorem,
|
||||
use_sketch=use_sketch,
|
||||
diagnostic_mode=diagnostic_mode if turn == 0 else False,
|
||||
diagnostic_instruction=diagnostic_instruction if turn == 0 else "",
|
||||
diagnostic_trace_context=diagnostic_trace_context if turn == 0 else "",
|
||||
previous_response=response if turn > 0 else "",
|
||||
)
|
||||
conversation.append({"type": "message", "turn": turn + 1, "content": response})
|
||||
if "<answer>" in response.lower():
|
||||
break
|
||||
|
||||
result["response"] = response
|
||||
result["agent_ok"] = True
|
||||
result["n_turns"] = len(conversation)
|
||||
|
||||
with open(os.path.join(pred_dir, "student_system_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(system)
|
||||
with open(os.path.join(pred_dir, "student_user_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(user)
|
||||
|
||||
eval_result = evaluate(response, item["correct_choice"], item["choices"])
|
||||
result["hard"] = int(eval_result["em"])
|
||||
result["soft"] = eval_result["f1"]
|
||||
result["predicted_answer"] = eval_result["predicted_answer"]
|
||||
result["predicted_label"] = eval_result["predicted_label"]
|
||||
result["predicted_text"] = eval_result["predicted_text"]
|
||||
if not result["hard"]:
|
||||
result["fail_reason"] = (
|
||||
f"MCQ=0: predicted '{eval_result['predicted_label'] or eval_result['predicted_answer']}' "
|
||||
f"but expected '{eval_result['correct_label']}'"
|
||||
)
|
||||
eval_detail = (
|
||||
f"[EVALUATION RESULT]\n"
|
||||
f"Question: {item['question']}\n"
|
||||
f"Predicted label: {eval_result['predicted_label']!r}\n"
|
||||
f"Predicted text: {eval_result['predicted_text']!r}\n"
|
||||
f"Correct label: {eval_result['correct_label']!r}\n"
|
||||
f"Correct text: {eval_result['correct_text']!r}\n"
|
||||
f"Exact Match: {eval_result['em']}"
|
||||
)
|
||||
conversation.append({"role": "system", "content": eval_detail})
|
||||
with open(os.path.join(pred_dir, "conversation.json"), "w") as f:
|
||||
json.dump(conversation, f, ensure_ascii=False, indent=2)
|
||||
return result
|
||||
|
||||
system = _build_system(skill_content)
|
||||
user = _build_user(
|
||||
item,
|
||||
use_theorem=use_theorem,
|
||||
use_sketch=use_sketch,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
diagnostic_trace_context=diagnostic_trace_context,
|
||||
)
|
||||
conversation: list[dict] = []
|
||||
response = ""
|
||||
|
||||
for turn in range(max_turns):
|
||||
if turn == 0:
|
||||
resp_text, _ = chat_student(
|
||||
system=system,
|
||||
user=user,
|
||||
max_completion_tokens=16384,
|
||||
retries=5,
|
||||
stage="rollout",
|
||||
timeout=exec_timeout,
|
||||
)
|
||||
else:
|
||||
refinement = (
|
||||
f"Your previous answer was:\n{response}\n\n"
|
||||
"Re-evaluate the exact option wording. If needed, correct it. "
|
||||
"Output only the final choice label inside <answer>...</answer>."
|
||||
)
|
||||
resp_text, _ = chat_student(
|
||||
system=system,
|
||||
user=refinement,
|
||||
max_completion_tokens=16384,
|
||||
retries=5,
|
||||
stage="rollout",
|
||||
timeout=exec_timeout,
|
||||
)
|
||||
response = resp_text
|
||||
conversation.append({"type": "message", "turn": turn + 1, "content": resp_text})
|
||||
if "<answer>" in resp_text.lower():
|
||||
break
|
||||
|
||||
result["response"] = response
|
||||
result["agent_ok"] = True
|
||||
result["n_turns"] = len(conversation)
|
||||
|
||||
with open(os.path.join(pred_dir, "student_system_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(system)
|
||||
with open(os.path.join(pred_dir, "student_user_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(user)
|
||||
|
||||
eval_result = evaluate(response, item["correct_choice"], item["choices"])
|
||||
result["hard"] = int(eval_result["em"])
|
||||
result["soft"] = eval_result["f1"]
|
||||
result["predicted_answer"] = eval_result["predicted_answer"]
|
||||
result["predicted_label"] = eval_result["predicted_label"]
|
||||
result["predicted_text"] = eval_result["predicted_text"]
|
||||
|
||||
if not result["hard"]:
|
||||
result["fail_reason"] = (
|
||||
f"MCQ=0: predicted '{eval_result['predicted_label'] or eval_result['predicted_answer']}' "
|
||||
f"but expected '{eval_result['correct_label']}'"
|
||||
)
|
||||
|
||||
eval_detail = (
|
||||
f"[EVALUATION RESULT]\n"
|
||||
f"Question: {item['question']}\n"
|
||||
f"Predicted label: {eval_result['predicted_label']!r}\n"
|
||||
f"Predicted text: {eval_result['predicted_text']!r}\n"
|
||||
f"Correct label: {eval_result['correct_label']!r}\n"
|
||||
f"Correct text: {eval_result['correct_text']!r}\n"
|
||||
f"Exact Match: {eval_result['em']}"
|
||||
)
|
||||
conversation.append({"role": "system", "content": eval_detail})
|
||||
|
||||
with open(os.path.join(pred_dir, "conversation.json"), "w") as f:
|
||||
json.dump(conversation, f, ensure_ascii=False, indent=2)
|
||||
|
||||
except Exception as e: # noqa: BLE001
|
||||
result["fail_reason"] = f"error: {e}"
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def run_batch(
|
||||
items: list[dict],
|
||||
out_root: str,
|
||||
skill_content: str,
|
||||
*,
|
||||
max_turns: int = 1,
|
||||
exec_timeout: int = 300,
|
||||
workers: int = 64,
|
||||
use_theorem: bool = False,
|
||||
use_sketch: bool = False,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context_by_id: dict[str, str] | None = None,
|
||||
task_timeout: int = 600,
|
||||
) -> list[dict]:
|
||||
task_timeout = max(int(task_timeout), int(exec_timeout) + 60)
|
||||
results_path = os.path.join(out_root, "results.jsonl")
|
||||
os.makedirs(out_root, exist_ok=True)
|
||||
|
||||
done_ids: set[str] = set()
|
||||
existing: list[dict] = []
|
||||
if os.path.exists(results_path):
|
||||
with open(results_path) as f:
|
||||
for line in f:
|
||||
try:
|
||||
r = json.loads(line)
|
||||
done_ids.add(str(r["id"]))
|
||||
existing.append(r)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
pending = [it for it in items if str(it["id"]) not in done_ids]
|
||||
if not pending:
|
||||
return existing
|
||||
|
||||
total = len(existing) + len(pending)
|
||||
completed = len(existing)
|
||||
correct_count = sum(1 for r in existing if r.get("hard", 0))
|
||||
if existing:
|
||||
print(f" [rollout] resuming: {completed}/{total} already done", flush=True)
|
||||
|
||||
results = list(existing)
|
||||
|
||||
started_at: dict[str, float] = {}
|
||||
|
||||
def _run_one(it: dict) -> dict:
|
||||
started_at[str(it["id"])] = time.time()
|
||||
return process_one(
|
||||
it,
|
||||
out_root,
|
||||
skill_content,
|
||||
max_turns=max_turns,
|
||||
exec_timeout=exec_timeout,
|
||||
use_theorem=use_theorem,
|
||||
use_sketch=use_sketch,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
diagnostic_trace_context=(diagnostic_trace_context_by_id or {}).get(str(it["id"]), ""),
|
||||
)
|
||||
|
||||
def _timeout_result(it: dict) -> dict:
|
||||
correct = it.get("correct_choice") or {}
|
||||
return {
|
||||
"id": str(it["id"]),
|
||||
"question": it.get("question", ""),
|
||||
"task_type": it.get("theorem_type", ["math_mcq"])[0] if it.get("theorem_type") else "math_mcq",
|
||||
"hard": 0,
|
||||
"soft": 0.0,
|
||||
"predicted_answer": "",
|
||||
"predicted_label": "",
|
||||
"predicted_text": "",
|
||||
"correct_label": correct.get("label", ""),
|
||||
"correct_text": correct.get("text", ""),
|
||||
"response": "",
|
||||
"fail_reason": f"task-timeout-{task_timeout}s",
|
||||
"agent_ok": False,
|
||||
"n_turns": 0,
|
||||
}
|
||||
|
||||
def _error_result(it: dict, exc: Exception) -> dict:
|
||||
res = _timeout_result(it)
|
||||
res["fail_reason"] = f"error: {type(exc).__name__}: {exc}"
|
||||
return res
|
||||
|
||||
with open(results_path, "a") as outf:
|
||||
ex = ThreadPoolExecutor(max_workers=workers)
|
||||
try:
|
||||
futs = {
|
||||
ex.submit(_run_one, it): it
|
||||
for it in pending
|
||||
}
|
||||
pending_futs = set(futs)
|
||||
while pending_futs:
|
||||
done, _ = wait(pending_futs, timeout=5, return_when=FIRST_COMPLETED)
|
||||
now = time.time()
|
||||
timed_out = [
|
||||
fut for fut in pending_futs - done
|
||||
if str(futs[fut]["id"]) in started_at
|
||||
and now - started_at[str(futs[fut]["id"])] >= task_timeout
|
||||
]
|
||||
for fut in done:
|
||||
pending_futs.remove(fut)
|
||||
item = futs[fut]
|
||||
try:
|
||||
res = fut.result()
|
||||
except Exception as e: # noqa: BLE001
|
||||
res = _error_result(item, e)
|
||||
results.append(res)
|
||||
completed += 1
|
||||
if res.get("hard", 0):
|
||||
correct_count += 1
|
||||
acc = correct_count / completed if completed else 0
|
||||
print(
|
||||
f" [rollout] {completed}/{total} "
|
||||
f"(acc={acc:.3f}) id={res['id']} "
|
||||
f"hard={res.get('hard', '?')}",
|
||||
flush=True,
|
||||
)
|
||||
outf.write(json.dumps(res, ensure_ascii=False) + "\n")
|
||||
outf.flush()
|
||||
for fut in timed_out:
|
||||
pending_futs.remove(fut)
|
||||
res = _timeout_result(futs[fut])
|
||||
results.append(res)
|
||||
completed += 1
|
||||
acc = correct_count / completed if completed else 0
|
||||
print(
|
||||
f" [rollout] {completed}/{total} "
|
||||
f"(acc={acc:.3f}) id={res['id']} TIMEOUT",
|
||||
flush=True,
|
||||
)
|
||||
outf.write(json.dumps(res, ensure_ascii=False) + "\n")
|
||||
outf.flush()
|
||||
finally:
|
||||
ex.shutdown(wait=False, cancel_futures=True)
|
||||
|
||||
return results
|
||||
@@ -0,0 +1,16 @@
|
||||
# Live Mathematical MCQ Heuristics
|
||||
|
||||
## Option Comparison
|
||||
- Compare all options before committing. The correct choice is often the strongest statement justified by the question, while nearby distractors are weaker, overstrong, or miss an equality case.
|
||||
- Track exact quantifiers such as "there exists", "for every", "if and only if", and "exactly when".
|
||||
|
||||
## Theorem-Level Precision
|
||||
- Check whether an option weakens the conclusion by dropping a characterization, equality clause, or full equivalence.
|
||||
- Check whether an option overstates the theorem by upgrading regularity, removing scale restrictions, or changing an existential statement into a universal one.
|
||||
|
||||
## Hypotheses
|
||||
- Verify the hypotheses and domain carefully. Distractors often keep the theorem shape but alter the required assumptions.
|
||||
- Pay close attention to equality cases, extremal conditions, and whether a result applies to the full family or only a restricted subfamily.
|
||||
|
||||
## Final Answer
|
||||
- Output the final answer as the single option label only.
|
||||
@@ -0,0 +1,5 @@
|
||||
"""MathVerse environment package."""
|
||||
|
||||
from skillopt.envs.mathverse.adapter import MathVerseAdapter
|
||||
|
||||
__all__ = ["MathVerseAdapter"]
|
||||
@@ -0,0 +1,280 @@
|
||||
"""MathVerse environment adapter for ReflACT."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
from skillopt.datasets.base import BatchSpec
|
||||
from skillopt.envs.base import EnvAdapter
|
||||
from skillopt.envs.mathverse.dataloader import MathVerseDataLoader
|
||||
from skillopt.envs.mathverse.rollout import run_batch
|
||||
from skillopt.gradient.deep_probe import generate_deep_probe_instruction
|
||||
from skillopt.gradient.reflect import run_minibatch_reflect
|
||||
from skillopt.model import get_student_backend
|
||||
|
||||
|
||||
class MathVerseAdapter(EnvAdapter):
|
||||
"""MathVerse adapter."""
|
||||
|
||||
def build_reference_text(self, item: dict) -> str:
|
||||
if not self.use_text_dominant_reference:
|
||||
return ""
|
||||
question = str(item.get("text_dominant_question") or "").strip()
|
||||
if not question:
|
||||
return ""
|
||||
return f"## Reference Full Question\n{question}"
|
||||
|
||||
def get_reference_metadata(self, item: dict) -> dict:
|
||||
if not self.use_text_dominant_reference:
|
||||
return {"fields": [], "preview": ""}
|
||||
question = str(item.get("text_dominant_question") or "").strip()
|
||||
if not question:
|
||||
return {"fields": [], "preview": ""}
|
||||
return {
|
||||
"fields": ["text_dominant_question"],
|
||||
"preview": question[:400],
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
split_dir: str = "",
|
||||
data_root: str = "",
|
||||
problem_version: str = "Text Lite",
|
||||
use_text_dominant_reference: bool = False,
|
||||
max_turns: int = 1,
|
||||
workers: int = 16,
|
||||
analyst_workers: int = 16,
|
||||
failure_only: bool = False,
|
||||
minibatch_size: int = 8,
|
||||
edit_budget: int = 4,
|
||||
seed: int = 42,
|
||||
limit: int = 0,
|
||||
image_detail: str = "auto",
|
||||
judge_model: str = "gpt-5.4",
|
||||
judge_max_completion_tokens: int = 256,
|
||||
judge_retries: int = 5,
|
||||
use_deep_reflect: bool = False,
|
||||
deep_reflect_failures: int = 4,
|
||||
deep_reflect_successes: int = 2,
|
||||
) -> None:
|
||||
self.max_turns = max_turns
|
||||
self.workers = workers
|
||||
self.analyst_workers = analyst_workers
|
||||
self.failure_only = failure_only
|
||||
self.minibatch_size = minibatch_size
|
||||
self.edit_budget = edit_budget
|
||||
self.image_detail = image_detail
|
||||
self.judge_model = judge_model
|
||||
self.judge_max_completion_tokens = judge_max_completion_tokens
|
||||
self.judge_retries = judge_retries
|
||||
self.problem_version = problem_version
|
||||
self.use_text_dominant_reference = use_text_dominant_reference
|
||||
self.use_deep_reflect = use_deep_reflect
|
||||
self.deep_reflect_failures = deep_reflect_failures
|
||||
self.deep_reflect_successes = deep_reflect_successes
|
||||
self.dataloader = MathVerseDataLoader(
|
||||
split_dir=split_dir,
|
||||
seed=seed,
|
||||
limit=limit,
|
||||
data_root=data_root,
|
||||
problem_version=problem_version,
|
||||
)
|
||||
|
||||
def setup(self, cfg: dict) -> None:
|
||||
super().setup(cfg)
|
||||
self.dataloader.setup(cfg)
|
||||
|
||||
def get_dataloader(self):
|
||||
return self.dataloader
|
||||
|
||||
def build_env_from_batch(self, batch: BatchSpec, **kwargs):
|
||||
return list(batch.payload or [])
|
||||
|
||||
def build_train_env(self, batch_size: int, seed: int, **kwargs):
|
||||
batch = self.dataloader.build_train_batch(batch_size=batch_size, seed=seed, **kwargs)
|
||||
return self.build_env_from_batch(batch, **kwargs)
|
||||
|
||||
def build_eval_env(self, env_num: int, split: str, seed: int, **kwargs):
|
||||
batch = self.dataloader.build_eval_batch(env_num=env_num, split=split, seed=seed, **kwargs)
|
||||
return self.build_env_from_batch(batch, **kwargs)
|
||||
|
||||
def rollout(
|
||||
self,
|
||||
env_manager,
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict]:
|
||||
items: list[dict] = env_manager
|
||||
return run_batch(
|
||||
items=items,
|
||||
out_root=out_dir,
|
||||
skill_content=skill_content,
|
||||
max_turns=self.max_turns,
|
||||
workers=self.workers,
|
||||
image_detail=self.image_detail,
|
||||
judge_model=self.judge_model,
|
||||
judge_max_completion_tokens=self.judge_max_completion_tokens,
|
||||
judge_retries=self.judge_retries,
|
||||
diagnostic_mode=kwargs.get("diagnostic_mode", False),
|
||||
diagnostic_instruction=kwargs.get("diagnostic_instruction", ""),
|
||||
diagnostic_trace_context_by_id=kwargs.get("diagnostic_trace_context_by_id"),
|
||||
)
|
||||
|
||||
def reflect(
|
||||
self,
|
||||
results: list[dict],
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict | None]:
|
||||
prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
|
||||
patches_dir = kwargs.get("patches_dir", os.path.join(out_dir, "patches"))
|
||||
random_seed = kwargs.get("random_seed")
|
||||
step_buffer_context = kwargs.get("step_buffer_context", "")
|
||||
|
||||
return run_minibatch_reflect(
|
||||
results=results,
|
||||
skill_content=skill_content,
|
||||
prediction_dir=prediction_dir,
|
||||
patches_dir=patches_dir,
|
||||
workers=self.analyst_workers,
|
||||
failure_only=self.failure_only,
|
||||
minibatch_size=self.minibatch_size,
|
||||
edit_budget=self.edit_budget,
|
||||
random_seed=random_seed,
|
||||
error_system=self.get_error_minibatch_prompt(),
|
||||
success_system=self.get_success_minibatch_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
update_mode=getattr(self, "_cfg", {}).get("skill_update_mode", "patch"),
|
||||
)
|
||||
|
||||
def deep_reflect(
|
||||
self,
|
||||
results: list[dict],
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict | None]:
|
||||
if not self.use_deep_reflect:
|
||||
return []
|
||||
|
||||
env_manager = kwargs.get("env_manager")
|
||||
prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
|
||||
random_seed = kwargs.get("random_seed")
|
||||
step_buffer_context = kwargs.get("step_buffer_context", "")
|
||||
selected_items = self.select_representative_items(
|
||||
results,
|
||||
env_manager if isinstance(env_manager, list) else None,
|
||||
n_failures=self.deep_reflect_failures,
|
||||
n_successes=self.deep_reflect_successes,
|
||||
seed=random_seed,
|
||||
)
|
||||
if not selected_items:
|
||||
return []
|
||||
|
||||
selected_ids = {str(item["id"]) for item in selected_items}
|
||||
selected_results = [row for row in results if str(row.get("id")) in selected_ids]
|
||||
selected_examples = self.attach_reference_context(selected_results, selected_items)
|
||||
codex_backend = get_student_backend() == "codex_exec"
|
||||
if codex_backend:
|
||||
selected_examples = self.attach_codex_probe_context(selected_examples, prediction_dir)
|
||||
selected_metadata = []
|
||||
ref_count = 0
|
||||
for item in selected_items:
|
||||
meta = self.get_reference_metadata(item)
|
||||
if meta["fields"]:
|
||||
ref_count += 1
|
||||
record = {
|
||||
"id": str(item["id"]),
|
||||
"task_type": str(item.get("task_type") or item.get("question_type") or "mathverse"),
|
||||
"reference_fields": meta["fields"],
|
||||
"reference_preview": meta["preview"],
|
||||
}
|
||||
if codex_backend:
|
||||
record["codex_probe_step_count"] = int(
|
||||
next(
|
||||
(row.get("codex_probe_step_count", 0) for row in selected_examples if str(row.get("id")) == str(item["id"])),
|
||||
0,
|
||||
)
|
||||
)
|
||||
selected_metadata.append(record)
|
||||
|
||||
deep_dir = os.path.join(out_dir, "deep_reflect")
|
||||
rollout_dir = os.path.join(deep_dir, "rollout")
|
||||
patches_dir = os.path.join(deep_dir, "patches")
|
||||
os.makedirs(deep_dir, exist_ok=True)
|
||||
print(
|
||||
f" [2b/6 DEEP REFLECT setup] selected={len(selected_items)} "
|
||||
f"reference_fields=text_dominant_question({ref_count}/{len(selected_items)})"
|
||||
)
|
||||
probe = generate_deep_probe_instruction(
|
||||
skill_content=skill_content,
|
||||
items=selected_examples,
|
||||
prediction_dir=prediction_dir,
|
||||
system_prompt=self.get_codex_deep_probe_prompt() if codex_backend else self.get_deep_probe_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
)
|
||||
if not probe:
|
||||
return []
|
||||
|
||||
targeted_items = selected_items
|
||||
diagnostic_trace_context_by_id: dict[str, str] | None = None
|
||||
if codex_backend:
|
||||
targeted_items, diagnostic_trace_context_by_id, probe = self.resolve_codex_probe_target(
|
||||
selected_items=selected_items,
|
||||
selected_examples=selected_examples,
|
||||
prediction_dir=prediction_dir,
|
||||
probe=probe,
|
||||
)
|
||||
|
||||
with open(os.path.join(deep_dir, "probe.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(
|
||||
{
|
||||
**probe,
|
||||
"reference_summary": {
|
||||
"selected_count": len(selected_items),
|
||||
"field_counts": {
|
||||
"text_dominant_question": ref_count,
|
||||
},
|
||||
},
|
||||
"selected_examples": selected_metadata,
|
||||
},
|
||||
f,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
|
||||
deep_results = run_batch(
|
||||
items=targeted_items,
|
||||
out_root=rollout_dir,
|
||||
skill_content=skill_content,
|
||||
max_turns=self.max_turns,
|
||||
workers=min(self.workers, max(len(targeted_items), 1)),
|
||||
image_detail=self.image_detail,
|
||||
judge_model=self.judge_model,
|
||||
judge_max_completion_tokens=self.judge_max_completion_tokens,
|
||||
judge_retries=self.judge_retries,
|
||||
diagnostic_mode=True,
|
||||
diagnostic_instruction=probe["probe_instruction"],
|
||||
diagnostic_trace_context_by_id=diagnostic_trace_context_by_id,
|
||||
)
|
||||
deep_results = self.attach_reference_context(deep_results, targeted_items)
|
||||
return run_minibatch_reflect(
|
||||
results=deep_results,
|
||||
skill_content=skill_content,
|
||||
prediction_dir=os.path.join(rollout_dir, "predictions"),
|
||||
patches_dir=patches_dir,
|
||||
workers=self.analyst_workers,
|
||||
failure_only=self.failure_only,
|
||||
minibatch_size=self.minibatch_size,
|
||||
edit_budget=self.edit_budget,
|
||||
random_seed=random_seed,
|
||||
error_system=self.get_error_minibatch_prompt(),
|
||||
success_system=self.get_success_minibatch_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
update_mode=getattr(self, "_cfg", {}).get("skill_update_mode", "patch"),
|
||||
)
|
||||
|
||||
def get_task_types(self) -> list[str]:
|
||||
return self.dataloader.get_task_types()
|
||||
@@ -0,0 +1,228 @@
|
||||
"""MathVerse task dataloader."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from skillopt.datasets.base import SplitDataLoader
|
||||
|
||||
|
||||
_CHOICE_LABELS = ["A", "B", "C", "D", "E", "F", "G"]
|
||||
_CHOICE_BLOCK_RE = re.compile(r"\bChoices?\s*:\s*", re.IGNORECASE)
|
||||
_CHOICE_ITEM_RE = re.compile(r"([A-G])\s*[:.)]\s*(.*?)(?=(?:\s+[A-G]\s*[:.)])|$)", re.DOTALL)
|
||||
|
||||
|
||||
def _load_json(path: str) -> Any:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def _normalize_space(text: Any) -> str:
|
||||
return re.sub(r"\s+", " ", str(text or "").strip())
|
||||
|
||||
|
||||
def _resolve_image_path(raw_path: str, *, data_root: str, source_path: str) -> str:
|
||||
candidates = []
|
||||
if raw_path:
|
||||
if os.path.isabs(raw_path):
|
||||
candidates.append(raw_path)
|
||||
else:
|
||||
if data_root:
|
||||
candidates.append(os.path.join(data_root, raw_path))
|
||||
candidates.append(os.path.join(data_root, "images", raw_path))
|
||||
candidates.append(os.path.join(os.path.dirname(source_path), raw_path))
|
||||
for candidate in candidates:
|
||||
if candidate and os.path.exists(candidate):
|
||||
return os.path.abspath(candidate)
|
||||
return ""
|
||||
|
||||
|
||||
def _split_question_and_choices(question: str) -> tuple[str, list[dict]]:
|
||||
text = str(question or "").strip()
|
||||
match = _CHOICE_BLOCK_RE.search(text)
|
||||
if not match:
|
||||
return text, []
|
||||
|
||||
stem = text[:match.start()].strip()
|
||||
choice_block = text[match.end():].strip()
|
||||
choices: list[dict] = []
|
||||
for idx, m in enumerate(_CHOICE_ITEM_RE.finditer(choice_block)):
|
||||
label = (m.group(1) or _CHOICE_LABELS[idx]).strip().upper()
|
||||
choice_text = _normalize_space(m.group(2))
|
||||
if choice_text:
|
||||
choices.append({"label": label, "text": choice_text})
|
||||
return stem or text, choices
|
||||
|
||||
|
||||
def _build_text_dominant_map(data_root: str) -> dict[str, str]:
|
||||
if not data_root:
|
||||
return {}
|
||||
candidates = [
|
||||
os.path.join(data_root, "testmini.json"),
|
||||
os.path.join(data_root, "data", "testmini.json"),
|
||||
]
|
||||
source_path = next((path for path in candidates if os.path.exists(path)), "")
|
||||
if not source_path:
|
||||
return {}
|
||||
|
||||
raw = _load_json(source_path)
|
||||
if not isinstance(raw, list):
|
||||
return {}
|
||||
|
||||
mapping: dict[str, str] = {}
|
||||
for item in raw:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
if str(item.get("problem_version") or "").strip() != "Text Dominant":
|
||||
continue
|
||||
problem_index = str(item.get("problem_index") or "").strip()
|
||||
question = str(item.get("question") or "").strip()
|
||||
if problem_index and question:
|
||||
mapping[problem_index] = question
|
||||
return mapping
|
||||
|
||||
|
||||
def _normalize_item(
|
||||
item: dict,
|
||||
*,
|
||||
row_idx: int,
|
||||
source_path: str,
|
||||
data_root: str,
|
||||
problem_version: str,
|
||||
text_dominant_map: dict[str, str],
|
||||
) -> dict | None:
|
||||
raw_problem_version = str(item.get("problem_version") or "").strip()
|
||||
if problem_version and raw_problem_version and raw_problem_version != problem_version:
|
||||
return None
|
||||
|
||||
question = str(item.get("question") or "").strip()
|
||||
question_type = str(item.get("question_type") or "").strip()
|
||||
answer = str(item.get("answer") or "").strip()
|
||||
image_rel = str(item.get("image") or "").strip()
|
||||
image_path = _resolve_image_path(image_rel, data_root=data_root, source_path=source_path)
|
||||
if not answer or not image_path:
|
||||
return None
|
||||
|
||||
metadata = item.get("metadata") if isinstance(item.get("metadata"), dict) else {}
|
||||
subject = str(metadata.get("subject") or "").strip()
|
||||
subfield = str(metadata.get("subfield") or "").strip()
|
||||
source = str(metadata.get("source") or "").strip()
|
||||
|
||||
question_stem, choices = _split_question_and_choices(question)
|
||||
is_choice = question_type == "multi-choice" or bool(choices)
|
||||
|
||||
correct_choice = {"label": "", "text": ""}
|
||||
if is_choice:
|
||||
label = str(answer).strip().upper().rstrip(".):")
|
||||
choice_text = ""
|
||||
for choice in choices:
|
||||
if choice["label"].upper() == label:
|
||||
choice_text = choice["text"]
|
||||
break
|
||||
correct_choice = {"label": label, "text": choice_text}
|
||||
|
||||
problem_index = str(item.get("problem_index") or "").strip()
|
||||
sample_index = str(item.get("sample_index") or row_idx + 1).strip()
|
||||
item_id = problem_index or sample_index
|
||||
task_type = subfield or subject or question_type or "mathverse"
|
||||
|
||||
return {
|
||||
"id": item_id,
|
||||
"sample_index": sample_index,
|
||||
"problem_index": problem_index,
|
||||
"problem_version": raw_problem_version or problem_version,
|
||||
"question": question,
|
||||
"question_stem": question_stem,
|
||||
"question_for_eval": str(item.get("question_for_eval") or question).strip(),
|
||||
"question_type": question_type or ("multi-choice" if is_choice else "free-form"),
|
||||
"is_choice": is_choice,
|
||||
"choices": choices,
|
||||
"correct_choice": correct_choice,
|
||||
"answer": answer,
|
||||
"gold_answers": [answer] if answer else [],
|
||||
"image_rel": image_rel,
|
||||
"image_path": image_path,
|
||||
"query_wo": str(item.get("query_wo") or "").strip(),
|
||||
"query_cot": str(item.get("query_cot") or "").strip(),
|
||||
"metadata": {
|
||||
"split": str(metadata.get("split") or "").strip(),
|
||||
"source": source,
|
||||
"subject": subject,
|
||||
"subfield": subfield,
|
||||
},
|
||||
"task_type": task_type,
|
||||
"source_path": os.path.abspath(source_path),
|
||||
"text_dominant_question": str(
|
||||
item.get("text_dominant_question")
|
||||
or text_dominant_map.get(problem_index, "")
|
||||
).strip(),
|
||||
}
|
||||
|
||||
|
||||
class MathVerseDataLoader(SplitDataLoader):
|
||||
"""MathVerse dataloader."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
split_dir: str = "",
|
||||
seed: int = 42,
|
||||
limit: int = 0,
|
||||
data_root: str = "",
|
||||
problem_version: str = "Text Lite",
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__(split_dir=split_dir, seed=seed, limit=limit)
|
||||
self.data_root = data_root
|
||||
self.problem_version = problem_version
|
||||
self._task_types: list[str] = []
|
||||
self._text_dominant_map = _build_text_dominant_map(data_root)
|
||||
|
||||
def setup(self, cfg: dict) -> None:
|
||||
if not self.data_root:
|
||||
self.data_root = str(cfg.get("data_root") or "")
|
||||
if not self.problem_version:
|
||||
self.problem_version = str(cfg.get("problem_version") or "Text Lite")
|
||||
self._text_dominant_map = _build_text_dominant_map(self.data_root)
|
||||
super().setup(cfg)
|
||||
all_items = self.train_items + self.val_items + self.test_items
|
||||
task_types = {
|
||||
item.get("task_type") or item.get("question_type") or "mathverse"
|
||||
for item in all_items
|
||||
}
|
||||
self._task_types = sorted(str(x) for x in task_types if str(x).strip())
|
||||
|
||||
def get_task_types(self) -> list[str]:
|
||||
return list(self._task_types)
|
||||
|
||||
def load_split_items(self, split_path: str) -> list[dict]:
|
||||
raw_items = super().load_split_items(split_path)
|
||||
source_path = next(
|
||||
(
|
||||
os.path.join(split_path, name)
|
||||
for name in sorted(os.listdir(split_path))
|
||||
if name.endswith(".json")
|
||||
),
|
||||
split_path,
|
||||
)
|
||||
items: list[dict] = []
|
||||
for row_idx, item in enumerate(raw_items):
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
norm = _normalize_item(
|
||||
item,
|
||||
row_idx=row_idx,
|
||||
source_path=source_path,
|
||||
data_root=self.data_root,
|
||||
problem_version=self.problem_version,
|
||||
text_dominant_map=self._text_dominant_map,
|
||||
)
|
||||
if norm is not None:
|
||||
items.append(norm)
|
||||
if not items:
|
||||
raise ValueError(
|
||||
f"No valid MathVerse items loaded from {split_path} "
|
||||
f"for problem_version={self.problem_version!r}"
|
||||
)
|
||||
return items
|
||||
@@ -0,0 +1,180 @@
|
||||
"""MathVerse evaluation helpers."""
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import string
|
||||
|
||||
from skillopt.model import chat_with_deployment
|
||||
from skillopt.prompts import load_prompt
|
||||
|
||||
|
||||
_EVAL_MODE = "mathverse_choice_or_judge_v1"
|
||||
|
||||
|
||||
def normalize_text(text: str) -> str:
|
||||
text = str(text or "").strip().lower()
|
||||
text = text.replace("\\,", " ")
|
||||
text = text.replace("\\ ", " ")
|
||||
text = "".join(ch for ch in text if ch not in string.punctuation)
|
||||
return " ".join(text.split())
|
||||
|
||||
|
||||
def normalize_math_text(text: str) -> str:
|
||||
text = str(text or "").strip()
|
||||
text = text.replace("$", "")
|
||||
text = text.replace("\\mathrm", "")
|
||||
text = text.replace("{", "")
|
||||
text = text.replace("}", "")
|
||||
text = text.replace("~", " ")
|
||||
text = text.replace("\\,", " ")
|
||||
text = text.replace("\\ ", " ")
|
||||
return " ".join(text.split()).lower()
|
||||
|
||||
|
||||
def extract_answer(text: str | None) -> str:
|
||||
raw = str(text or "").strip()
|
||||
if not raw:
|
||||
return ""
|
||||
|
||||
tags = re.findall(r"<answer>\s*(.*?)\s*</answer>", raw, re.IGNORECASE | re.DOTALL)
|
||||
if tags:
|
||||
return tags[-1].strip()
|
||||
|
||||
boxed = re.findall(r"\\boxed\{(.*?)\}", raw, re.IGNORECASE | re.DOTALL)
|
||||
if boxed:
|
||||
return boxed[-1].strip()
|
||||
|
||||
lines = [ln.strip() for ln in raw.splitlines() if ln.strip()]
|
||||
if lines:
|
||||
return lines[-1]
|
||||
return raw
|
||||
|
||||
|
||||
def _judge_answer(
|
||||
*,
|
||||
item: dict,
|
||||
extracted_answer: str,
|
||||
judge_model: str,
|
||||
max_completion_tokens: int,
|
||||
retries: int,
|
||||
) -> dict:
|
||||
question = str(item.get("question_for_eval") or item.get("question") or "").strip()
|
||||
ground_truth = str(item.get("answer") or "").strip()
|
||||
raw, _ = chat_with_deployment(
|
||||
deployment=judge_model,
|
||||
system="You are a careful and strict mathematical answer evaluator.",
|
||||
user=load_prompt("judge", env="mathverse").format(
|
||||
question=question,
|
||||
groundtruth=ground_truth,
|
||||
modeloutput=extracted_answer,
|
||||
),
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
retries=retries,
|
||||
stage="mathverse_judge",
|
||||
)
|
||||
response = str(raw).strip().lower()
|
||||
if "true" in response:
|
||||
correct = True
|
||||
elif "false" in response:
|
||||
correct = False
|
||||
else:
|
||||
correct = False
|
||||
return {
|
||||
"raw": raw,
|
||||
"correct": correct,
|
||||
"reason": response,
|
||||
"matched_gold": ground_truth if correct else "",
|
||||
}
|
||||
|
||||
|
||||
def evaluate_item(
|
||||
*,
|
||||
item: dict,
|
||||
prediction_text: str,
|
||||
judge_model: str,
|
||||
max_completion_tokens: int = 256,
|
||||
retries: int = 5,
|
||||
) -> dict:
|
||||
extracted = extract_answer(prediction_text)
|
||||
|
||||
if item.get("is_choice"):
|
||||
predicted_label = str(extracted).strip().upper().rstrip(".):")
|
||||
correct_label = str(item["correct_choice"].get("label") or "").strip().upper()
|
||||
predicted_text = ""
|
||||
for choice in item.get("choices") or []:
|
||||
if str(choice.get("label") or "").strip().upper() == predicted_label:
|
||||
predicted_text = str(choice.get("text") or "").strip()
|
||||
break
|
||||
hard = 1.0 if predicted_label == correct_label else 0.0
|
||||
return {
|
||||
"evaluation_mode": _EVAL_MODE,
|
||||
"predicted_answer": extracted,
|
||||
"predicted_label": predicted_label,
|
||||
"predicted_text": predicted_text,
|
||||
"correct_label": correct_label,
|
||||
"correct_text": str(item["correct_choice"].get("text") or "").strip(),
|
||||
"em": hard,
|
||||
"f1": hard,
|
||||
"sub_em": hard,
|
||||
"judge_raw": "",
|
||||
"judge_reason": "exact_label_match" if hard else "label_mismatch",
|
||||
"matched_gold": correct_label if hard else "",
|
||||
}
|
||||
|
||||
gold_answer = str(item.get("answer") or "").strip()
|
||||
pred_norm = normalize_math_text(extracted)
|
||||
gold_norm = normalize_math_text(gold_answer)
|
||||
if pred_norm and gold_norm and pred_norm == gold_norm:
|
||||
return {
|
||||
"evaluation_mode": _EVAL_MODE,
|
||||
"predicted_answer": extracted,
|
||||
"em": 1.0,
|
||||
"f1": 1.0,
|
||||
"sub_em": 1.0,
|
||||
"judge_raw": "",
|
||||
"judge_reason": "normalized_exact_match",
|
||||
"matched_gold": gold_answer,
|
||||
"string_f1": 1.0,
|
||||
}
|
||||
|
||||
judge = _judge_answer(
|
||||
item=item,
|
||||
extracted_answer=extracted,
|
||||
judge_model=judge_model,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
retries=retries,
|
||||
)
|
||||
hard = 1.0 if judge["correct"] else 0.0
|
||||
pred_tokens = normalize_text(extracted).split()
|
||||
gold_tokens = normalize_text(gold_answer).split()
|
||||
overlap = 0
|
||||
gold_counts: dict[str, int] = {}
|
||||
for tok in gold_tokens:
|
||||
gold_counts[tok] = gold_counts.get(tok, 0) + 1
|
||||
for tok in pred_tokens:
|
||||
count = gold_counts.get(tok, 0)
|
||||
if count > 0:
|
||||
overlap += 1
|
||||
gold_counts[tok] = count - 1
|
||||
if pred_tokens and gold_tokens and overlap:
|
||||
precision = overlap / len(pred_tokens)
|
||||
recall = overlap / len(gold_tokens)
|
||||
string_f1 = 2 * precision * recall / (precision + recall)
|
||||
else:
|
||||
string_f1 = 0.0
|
||||
|
||||
return {
|
||||
"evaluation_mode": _EVAL_MODE,
|
||||
"predicted_answer": extracted,
|
||||
"em": hard,
|
||||
"f1": hard,
|
||||
"sub_em": hard,
|
||||
"judge_raw": judge["raw"],
|
||||
"judge_reason": judge["reason"],
|
||||
"matched_gold": judge["matched_gold"],
|
||||
"string_f1": string_f1,
|
||||
}
|
||||
|
||||
|
||||
def evaluation_mode() -> str:
|
||||
return _EVAL_MODE
|
||||
@@ -0,0 +1,37 @@
|
||||
You are an expert failure-analysis agent for visual mathematical reasoning problems.
|
||||
|
||||
You will be given MULTIPLE failed trajectories from a single minibatch and the current skill document.
|
||||
Each trajectory includes the student's response, the evaluation result, and sometimes a hidden reference
|
||||
containing the fuller Text Dominant version of the same problem.
|
||||
|
||||
Your job is to identify COMMON reasoning failures across the batch and propose concise skill edits.
|
||||
|
||||
## Failure Type Categories
|
||||
- **diagram_underuse**: the agent did not recover key constraints from the image
|
||||
- **constraint_drop**: the agent ignored a condition or relation that should guide the solution
|
||||
- **option_confusion**: the agent failed to discriminate between close answer choices
|
||||
- **format_miss**: the agent solved roughly correctly but returned the wrong final form, unit, or expression
|
||||
- **other**: none of the above
|
||||
|
||||
## Rules
|
||||
1. Focus on patterns that recur across the minibatch.
|
||||
2. Prefer edits that improve visual grounding and exact answer selection.
|
||||
3. Do not hardcode problem-specific formulas or answers.
|
||||
4. If hidden reference text is present, use it only to infer what information the student failed to recover from the Text Lite version.
|
||||
|
||||
Respond ONLY with a valid JSON object:
|
||||
{
|
||||
"batch_size": <number>,
|
||||
"failure_summary": [
|
||||
{"failure_type": "<type>", "count": <int>, "description": "<one-line>"}
|
||||
],
|
||||
"patch": {
|
||||
"reasoning": "<why these edits address the common failures>",
|
||||
"edits": [
|
||||
{"op": "append", "content": "<markdown>"},
|
||||
{"op": "insert_after", "target": "<heading/text>", "content": "<markdown>"},
|
||||
{"op": "replace", "target": "<old text>", "content": "<new text>"},
|
||||
{"op": "delete", "target": "<exact text to remove>"}
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
You are an expert success-pattern analyst for visual mathematical reasoning problems.
|
||||
|
||||
You will be given MULTIPLE successful trajectories from a minibatch and the current skill document.
|
||||
Identify generalizable behavior patterns that genuinely help the agent recover the right constraints
|
||||
from the image and convert them into the exact final answer.
|
||||
|
||||
## Rules
|
||||
- Focus on broadly useful visual-math reasoning behaviors.
|
||||
- Prefer patterns about reading decisive diagram cues, checking hidden assumptions, and matching the final answer format exactly.
|
||||
- Do not add benchmark-specific facts or formulas.
|
||||
- "edits" may be empty if the skill already captures the useful patterns.
|
||||
|
||||
Respond ONLY with a valid JSON object:
|
||||
{
|
||||
"batch_size": <number>,
|
||||
"success_patterns": ["<pattern 1>", "<pattern 2>"],
|
||||
"patch": {
|
||||
"reasoning": "<why these patterns matter>",
|
||||
"edits": [
|
||||
{"op": "append", "content": "<markdown>"},
|
||||
{"op": "insert_after", "target": "<heading/text>", "content": "<markdown>"},
|
||||
{"op": "replace", "target": "<old text>", "content": "<new text>"},
|
||||
{"op": "delete", "target": "<exact text to remove>"}
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
You are an expert diagnostic-probe designer for visual mathematical reasoning tasks.
|
||||
|
||||
You will be shown representative trajectories, the current student skill, and the student's original prompt context.
|
||||
Some trajectories may also include a hidden reference containing the fuller Text Dominant wording of the same problem.
|
||||
Design one SMALL diagnostic instruction that exposes the student's intermediate judgment without materially changing the original scaffold.
|
||||
|
||||
## Hard Constraints
|
||||
1. Do NOT substantially change the original scaffold.
|
||||
2. Do NOT prescribe a new long multi-step solving procedure.
|
||||
3. Do NOT ask for a full proof or full chain-of-thought.
|
||||
4. Ask only for a short readout of the signals already behind the student's current answer.
|
||||
5. Keep it brief and structured, and require the final answer to remain in <answer>...</answer>.
|
||||
6. If hidden reference text is present, use it only to target what visual or textual constraint the student likely missed.
|
||||
|
||||
## Good Probe Targets
|
||||
- decisive diagram cue
|
||||
- top candidate and runner-up
|
||||
- missing relation or quantity
|
||||
- why a near-miss option was rejected
|
||||
|
||||
Respond ONLY with a valid JSON object:
|
||||
{
|
||||
"reasoning": "<why this probe is informative>",
|
||||
"probe_instruction": "<the exact instruction text to append to the student prompt>"
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
You are a careful and strict evaluator for visual math problems.
|
||||
|
||||
You will be given:
|
||||
1. The original question
|
||||
2. The ground-truth answer
|
||||
3. A model output
|
||||
|
||||
Decide whether the model output is mathematically equivalent to the ground-truth answer.
|
||||
|
||||
Rules:
|
||||
- Ignore harmless formatting differences.
|
||||
- Accept mathematically equivalent expressions, equations, and values.
|
||||
- Reject answers that are numerically wrong, symbolically different in meaning, missing required units when the unit changes meaning, or correspond to a different choice.
|
||||
- Do not reward partially correct reasoning if the final answer is wrong.
|
||||
|
||||
Return only:
|
||||
True
|
||||
|
||||
or
|
||||
|
||||
False
|
||||
|
||||
Question: {question}
|
||||
Ground Truth Answer: {groundtruth}
|
||||
Model Output: {modeloutput}
|
||||
@@ -0,0 +1,11 @@
|
||||
You are an expert visual mathematical reasoning agent.
|
||||
|
||||
{skill_section}## Task Format
|
||||
You will receive one math problem with an image or diagram.
|
||||
Use the visible diagram as evidence, not just the text.
|
||||
If some information is abbreviated in the text, recover it from the image before answering.
|
||||
|
||||
## Answer Format
|
||||
Think step by step, then provide your final answer inside <answer>...</answer>.
|
||||
- For multiple-choice questions, output only the single option label, such as <answer>B</answer>.
|
||||
- For free-form questions, output only the final mathematical answer, such as <answer>14</answer>.
|
||||
@@ -0,0 +1,4 @@
|
||||
"""MathVerse Reflect stage.
|
||||
|
||||
Prompts are loaded from .md files by the base adapter.
|
||||
"""
|
||||
@@ -0,0 +1,431 @@
|
||||
"""MathVerse rollout — single-image multimodal math reasoning."""
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
from skillopt.envs.mathverse.evaluator import evaluate_item, evaluation_mode, extract_answer
|
||||
from skillopt.model import chat_student_messages, get_student_backend, is_student_exec_backend
|
||||
from skillopt.model.codex_harness import prepare_workspace, render_skill_md, run_student_exec
|
||||
from skillopt.prompts import load_prompt
|
||||
|
||||
|
||||
def _build_system(skill_content: str) -> str:
|
||||
if skill_content.strip():
|
||||
skill_section = f"## Skill\n{skill_content.strip()}\n\n"
|
||||
else:
|
||||
skill_section = ""
|
||||
return load_prompt("rollout_system", env="mathverse").format(skill_section=skill_section)
|
||||
|
||||
|
||||
def _format_choices(choices: list[dict]) -> str:
|
||||
return "\n".join(f"{choice['label']}. {choice['text']}" for choice in choices)
|
||||
|
||||
|
||||
def _build_user_text(
|
||||
item: dict,
|
||||
*,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context: str = "",
|
||||
) -> str:
|
||||
parts = []
|
||||
if diagnostic_trace_context.strip():
|
||||
parts.append(
|
||||
"## Previous Codex Trace Snapshot\n"
|
||||
"This is a partial transcript from an earlier attempt. Use it as your current reasoning context.\n\n"
|
||||
f"{diagnostic_trace_context.strip()}"
|
||||
)
|
||||
question = str(item.get("question_stem") or item.get("question") or "").strip()
|
||||
if question:
|
||||
parts.append(f"## Question\n{question}")
|
||||
else:
|
||||
parts.append("## Question\nRead the full problem statement from the image.")
|
||||
|
||||
if item.get("is_choice"):
|
||||
choices = item.get("choices") or []
|
||||
if choices:
|
||||
parts.append(f"## Choices\n{_format_choices(choices)}")
|
||||
parts.append("Return only the final option label inside <answer>...</answer>.")
|
||||
else:
|
||||
parts.append("Return only the final mathematical answer inside <answer>...</answer>.")
|
||||
|
||||
if diagnostic_mode and diagnostic_instruction.strip():
|
||||
parts.append(f"## Training Readout\n{diagnostic_instruction.strip()}")
|
||||
return "\n\n".join(parts)
|
||||
|
||||
|
||||
def _image_to_data_uri(path: str) -> str:
|
||||
mime = mimetypes.guess_type(path)[0] or "image/png"
|
||||
with open(path, "rb") as f:
|
||||
encoded = base64.b64encode(f.read()).decode("ascii")
|
||||
return f"data:{mime};base64,{encoded}"
|
||||
|
||||
|
||||
def _build_messages(
|
||||
item: dict,
|
||||
skill_content: str,
|
||||
image_detail: str,
|
||||
*,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context: str = "",
|
||||
) -> tuple[list[dict], str, str]:
|
||||
system = _build_system(skill_content)
|
||||
user_text = _build_user_text(
|
||||
item,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
diagnostic_trace_context=diagnostic_trace_context,
|
||||
)
|
||||
image_url = {"url": _image_to_data_uri(item["image_path"])}
|
||||
if image_detail and image_detail != "auto":
|
||||
image_url["detail"] = image_detail
|
||||
messages = [
|
||||
{"role": "system", "content": system},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": user_text},
|
||||
{"type": "image_url", "image_url": image_url},
|
||||
],
|
||||
},
|
||||
]
|
||||
return messages, system, user_text
|
||||
|
||||
|
||||
def _build_codex_skill(skill_content: str) -> str:
|
||||
return render_skill_md(
|
||||
skill_content,
|
||||
description="Dynamic ReflACT skill for solving the current MathVerse visual math problem.",
|
||||
preamble=(
|
||||
"Use this skill when solving the current MathVerse problem.\n"
|
||||
"Read the image carefully and return the final answer inside <answer>...</answer>."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _run_codex_once(
|
||||
*,
|
||||
pred_dir: str,
|
||||
item: dict,
|
||||
skill_content: str,
|
||||
model: str,
|
||||
timeout: int,
|
||||
image_detail: str,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context: str = "",
|
||||
previous_response: str = "",
|
||||
) -> tuple[str, str, str, str]:
|
||||
user_text = _build_user_text(
|
||||
item,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
diagnostic_trace_context=diagnostic_trace_context,
|
||||
)
|
||||
task_parts = [user_text]
|
||||
if previous_response:
|
||||
task_parts.append(
|
||||
"## Previous Attempt\n"
|
||||
f"{previous_response}\n\n"
|
||||
"Re-check the diagram and the mathematical constraints. Correct the final answer if needed."
|
||||
)
|
||||
task_text = "\n\n".join(task_parts)
|
||||
skill_md = _build_codex_skill(skill_content)
|
||||
work_dir = os.path.join(pred_dir, "codex_exec")
|
||||
prepare_workspace(
|
||||
work_dir=work_dir,
|
||||
skill_md=skill_md,
|
||||
task_text=task_text,
|
||||
images=[item["image_path"]],
|
||||
)
|
||||
prompt = (
|
||||
"Use the `skillopt-student` skill available in this workspace.\n"
|
||||
"Read `task.md`, inspect the attached image, solve the problem, and return only the final answer inside <answer>...</answer>."
|
||||
)
|
||||
final_message, raw = run_student_exec(
|
||||
work_dir=work_dir,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
timeout=timeout,
|
||||
images=[item["image_path"]],
|
||||
)
|
||||
return final_message or raw, raw, skill_md, task_text
|
||||
|
||||
|
||||
def process_one(
|
||||
item: dict,
|
||||
out_root: str,
|
||||
skill_content: str,
|
||||
*,
|
||||
max_turns: int = 1,
|
||||
image_detail: str = "auto",
|
||||
judge_model: str = "gpt-5.4",
|
||||
judge_max_completion_tokens: int = 256,
|
||||
judge_retries: int = 5,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context: str = "",
|
||||
) -> dict:
|
||||
item_id = str(item["id"])
|
||||
result = {
|
||||
"id": item_id,
|
||||
"question": item["question"],
|
||||
"task_type": item.get("task_type") or item.get("question_type") or "mathverse",
|
||||
"task_description": item.get("question_stem") or item["question"],
|
||||
"hard": 0,
|
||||
"soft": 0.0,
|
||||
"predicted_answer": "",
|
||||
"predicted_label": "",
|
||||
"predicted_text": "",
|
||||
"response": "",
|
||||
"fail_reason": "",
|
||||
"agent_ok": False,
|
||||
"n_turns": 0,
|
||||
"image_path": item["image_path"],
|
||||
"question_type": item["question_type"],
|
||||
"evaluation_mode": evaluation_mode(),
|
||||
"judge_model": judge_model,
|
||||
}
|
||||
if item.get("is_choice"):
|
||||
result["correct_label"] = item["correct_choice"]["label"]
|
||||
result["correct_text"] = item["correct_choice"]["text"]
|
||||
else:
|
||||
result["gold_answers"] = item.get("gold_answers") or [item["answer"]]
|
||||
|
||||
try:
|
||||
pred_dir = os.path.join(out_root, "predictions", item_id)
|
||||
os.makedirs(pred_dir, exist_ok=True)
|
||||
|
||||
if is_student_exec_backend():
|
||||
from skillopt.model import azure_openai as _llm
|
||||
|
||||
response = ""
|
||||
conversation: list[dict] = [
|
||||
{"role": "user", "content": f"{item['question']}\n\n[image] {os.path.basename(item['image_path'])}"}
|
||||
]
|
||||
system_prompt = ""
|
||||
user_text = ""
|
||||
for turn in range(max_turns):
|
||||
response, raw, system_prompt, user_text = _run_codex_once(
|
||||
pred_dir=pred_dir,
|
||||
item=item,
|
||||
skill_content=skill_content,
|
||||
model=_llm.STUDENT_DEPLOYMENT,
|
||||
timeout=120,
|
||||
image_detail=image_detail,
|
||||
diagnostic_mode=diagnostic_mode if turn == 0 else False,
|
||||
diagnostic_instruction=diagnostic_instruction if turn == 0 else "",
|
||||
diagnostic_trace_context=diagnostic_trace_context if turn == 0 else "",
|
||||
previous_response=response if turn > 0 else "",
|
||||
)
|
||||
conversation.append({"type": "message", "turn": turn + 1, "content": response})
|
||||
if extract_answer(response):
|
||||
break
|
||||
|
||||
result["response"] = response
|
||||
result["agent_ok"] = True
|
||||
result["n_turns"] = len(conversation) - 1
|
||||
with open(os.path.join(pred_dir, "student_system_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(system_prompt)
|
||||
with open(os.path.join(pred_dir, "student_user_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(user_text)
|
||||
else:
|
||||
messages, system_prompt, user_text = _build_messages(
|
||||
item,
|
||||
skill_content,
|
||||
image_detail,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
diagnostic_trace_context=diagnostic_trace_context,
|
||||
)
|
||||
response = ""
|
||||
conversation = [
|
||||
{"role": "user", "content": f"{user_text}\n\n[image] {os.path.basename(item['image_path'])}"}
|
||||
]
|
||||
for turn in range(max_turns):
|
||||
if turn == 0:
|
||||
resp_text, _ = chat_student_messages(
|
||||
messages=messages,
|
||||
max_completion_tokens=1024,
|
||||
retries=5,
|
||||
stage="rollout",
|
||||
)
|
||||
else:
|
||||
refinement_text = (
|
||||
f"Your previous answer was:\n{response}\n\n"
|
||||
"Re-check the diagram and the mathematical constraints. "
|
||||
"If needed, correct your answer. Output only the final answer inside <answer>...</answer>."
|
||||
)
|
||||
refinement_messages = [
|
||||
messages[0],
|
||||
messages[1],
|
||||
{"role": "assistant", "content": response},
|
||||
{"role": "user", "content": refinement_text},
|
||||
]
|
||||
resp_text, _ = chat_student_messages(
|
||||
messages=refinement_messages,
|
||||
max_completion_tokens=768,
|
||||
retries=5,
|
||||
stage="rollout",
|
||||
)
|
||||
response = resp_text
|
||||
conversation.append({"type": "message", "turn": turn + 1, "content": resp_text})
|
||||
if extract_answer(resp_text):
|
||||
break
|
||||
|
||||
result["response"] = response
|
||||
result["agent_ok"] = True
|
||||
result["n_turns"] = len(conversation) - 1
|
||||
with open(os.path.join(pred_dir, "student_system_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(system_prompt)
|
||||
with open(os.path.join(pred_dir, "student_user_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(user_text)
|
||||
|
||||
eval_result = evaluate_item(
|
||||
item=item,
|
||||
prediction_text=result["response"],
|
||||
judge_model=judge_model,
|
||||
max_completion_tokens=judge_max_completion_tokens,
|
||||
retries=judge_retries,
|
||||
)
|
||||
result["evaluation_mode"] = eval_result["evaluation_mode"]
|
||||
result["judge_raw"] = eval_result.get("judge_raw", "")
|
||||
result["judge_reason"] = eval_result.get("judge_reason", "")
|
||||
result["matched_gold"] = eval_result.get("matched_gold", "")
|
||||
|
||||
if item.get("is_choice"):
|
||||
result["predicted_label"] = eval_result["predicted_label"]
|
||||
result["predicted_text"] = eval_result["predicted_text"]
|
||||
result["predicted_answer"] = eval_result["predicted_answer"]
|
||||
result["hard"] = int(eval_result["em"])
|
||||
result["soft"] = eval_result["f1"]
|
||||
if not result["hard"]:
|
||||
result["fail_reason"] = (
|
||||
f"choice=0: predicted '{eval_result['predicted_label'] or eval_result['predicted_answer']}' "
|
||||
f"but expected '{eval_result['correct_label']}'"
|
||||
)
|
||||
eval_detail = (
|
||||
f"[EVALUATION RESULT]\n"
|
||||
f"Question: {item['question_for_eval']}\n"
|
||||
f"Predicted label: {eval_result['predicted_label']!r}\n"
|
||||
f"Predicted text: {eval_result['predicted_text']!r}\n"
|
||||
f"Correct label: {eval_result['correct_label']!r}\n"
|
||||
f"Correct text: {eval_result['correct_text']!r}\n"
|
||||
f"Exact Match: {eval_result['em']}"
|
||||
)
|
||||
else:
|
||||
result["predicted_answer"] = eval_result["predicted_answer"]
|
||||
result["hard"] = int(eval_result["em"])
|
||||
result["soft"] = eval_result["f1"]
|
||||
if not result["hard"]:
|
||||
result["fail_reason"] = (
|
||||
f"judge=0: predicted '{eval_result['predicted_answer']}' "
|
||||
f"but expected '{item['answer']}' ({eval_result.get('judge_reason', '')})"
|
||||
)
|
||||
eval_detail = (
|
||||
f"[EVALUATION RESULT]\n"
|
||||
f"Question: {item['question_for_eval']}\n"
|
||||
f"Predicted answer: {eval_result['predicted_answer']!r}\n"
|
||||
f"Gold answer: {item['answer']!r}\n"
|
||||
f"Judge correct: {eval_result['em']}\n"
|
||||
f"Judge reason: {eval_result.get('judge_reason', '')}\n"
|
||||
f"String F1: {eval_result.get('string_f1', 0.0):.4f}"
|
||||
)
|
||||
|
||||
conversation.append({"role": "system", "content": eval_detail})
|
||||
with open(os.path.join(pred_dir, "conversation.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(conversation, f, ensure_ascii=False, indent=2)
|
||||
except Exception as e: # noqa: BLE001
|
||||
result["fail_reason"] = f"error: {e}"
|
||||
return result
|
||||
|
||||
|
||||
def run_batch(
|
||||
items: list[dict],
|
||||
out_root: str,
|
||||
skill_content: str,
|
||||
*,
|
||||
max_turns: int = 1,
|
||||
workers: int = 32,
|
||||
image_detail: str = "auto",
|
||||
judge_model: str = "gpt-5.4",
|
||||
judge_max_completion_tokens: int = 256,
|
||||
judge_retries: int = 5,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context_by_id: dict[str, str] | None = None,
|
||||
) -> list[dict]:
|
||||
results_path = os.path.join(out_root, "results.jsonl")
|
||||
os.makedirs(out_root, exist_ok=True)
|
||||
|
||||
expected_eval_mode = evaluation_mode()
|
||||
done_ids: set[str] = set()
|
||||
existing: list[dict] = []
|
||||
rewrite_results = False
|
||||
if os.path.exists(results_path):
|
||||
with open(results_path, encoding="utf-8") as f:
|
||||
for line in f:
|
||||
try:
|
||||
row = json.loads(line)
|
||||
if row.get("evaluation_mode") != expected_eval_mode:
|
||||
rewrite_results = True
|
||||
continue
|
||||
done_ids.add(str(row["id"]))
|
||||
existing.append(row)
|
||||
except Exception:
|
||||
rewrite_results = True
|
||||
|
||||
pending = [item for item in items if str(item["id"]) not in done_ids]
|
||||
if not pending and not rewrite_results:
|
||||
return existing
|
||||
|
||||
total = len(existing) + len(pending)
|
||||
completed = len(existing)
|
||||
correct_count = sum(1 for r in existing if r.get("hard", 0))
|
||||
if existing:
|
||||
print(f" [rollout] resuming: {completed}/{total} already done", flush=True)
|
||||
|
||||
results = list(existing)
|
||||
file_mode = "w" if rewrite_results else "a"
|
||||
with open(results_path, file_mode, encoding="utf-8") as outf, ThreadPoolExecutor(max_workers=workers) as ex:
|
||||
if rewrite_results:
|
||||
for row in existing:
|
||||
outf.write(json.dumps(row, ensure_ascii=False) + "\n")
|
||||
futs = {
|
||||
ex.submit(
|
||||
process_one,
|
||||
item,
|
||||
out_root,
|
||||
skill_content,
|
||||
max_turns=max_turns,
|
||||
image_detail=image_detail,
|
||||
judge_model=judge_model,
|
||||
judge_max_completion_tokens=judge_max_completion_tokens,
|
||||
judge_retries=judge_retries,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
diagnostic_trace_context=(diagnostic_trace_context_by_id or {}).get(str(item["id"]), ""),
|
||||
): item
|
||||
for item in pending
|
||||
}
|
||||
for fut in as_completed(futs):
|
||||
row = fut.result()
|
||||
results.append(row)
|
||||
completed += 1
|
||||
if row.get("hard", 0):
|
||||
correct_count += 1
|
||||
acc = correct_count / completed if completed else 0
|
||||
print(
|
||||
f" [rollout] {completed}/{total} "
|
||||
f"(acc={acc:.3f}) id={row.get('id', '?')} "
|
||||
f"hard={row.get('hard', '?')}",
|
||||
flush=True,
|
||||
)
|
||||
outf.write(json.dumps(row, ensure_ascii=False) + "\n")
|
||||
outf.flush()
|
||||
return results
|
||||
@@ -0,0 +1,15 @@
|
||||
# MathVerse Visual Math Heuristics
|
||||
|
||||
## Diagram First
|
||||
- Read the diagram before locking onto an equation or option.
|
||||
- Recover missing labels, lengths, angles, axes, or object relations from the image when the text is abbreviated.
|
||||
- If the text seems underspecified, assume the image may contain the decisive constraint.
|
||||
|
||||
## Constraint Tracking
|
||||
- Write down the few constraints that actually determine the answer instead of solving from vague intuition.
|
||||
- Prefer geometric or functional relations that are directly supported by the figure.
|
||||
- For multiple-choice questions, compare the final candidate against every option exactly.
|
||||
|
||||
## Final Answer
|
||||
- Use the image and the text consistently.
|
||||
- Return only the final answer inside <answer>...</answer>.
|
||||
@@ -0,0 +1,2 @@
|
||||
"""MMRB environment package."""
|
||||
|
||||
@@ -0,0 +1,283 @@
|
||||
"""MMRB environment adapter for ReflACT."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
from skillopt.gradient.deep_probe import generate_deep_probe_instruction
|
||||
from skillopt.datasets.base import BatchSpec
|
||||
from skillopt.gradient.reflect import run_minibatch_reflect
|
||||
from skillopt.envs.base import EnvAdapter
|
||||
from skillopt.envs.mmrb.dataloader import MMRBDataLoader
|
||||
from skillopt.envs.mmrb.rollout import run_batch
|
||||
from skillopt.model import get_student_backend
|
||||
|
||||
|
||||
class MMRBAdapter(EnvAdapter):
|
||||
"""MMRB adapter."""
|
||||
|
||||
def build_reference_text(self, item: dict) -> str:
|
||||
reasoning_steps = item.get("reasoning_steps") or []
|
||||
if not reasoning_steps:
|
||||
return ""
|
||||
|
||||
blocks: list[str] = []
|
||||
for path_idx, path in enumerate(reasoning_steps, 1):
|
||||
if not isinstance(path, list) or not path:
|
||||
continue
|
||||
lines = [f"### Reasoning Path {path_idx}"]
|
||||
for step in path:
|
||||
if not isinstance(step, dict):
|
||||
continue
|
||||
step_no = step.get("reasoning step", "?")
|
||||
step_type = str(step.get("reasoning type") or "").strip()
|
||||
rationale = str(step.get("rationale") or "").strip()
|
||||
if rationale:
|
||||
prefix = f"{step_no}. [{step_type}] " if step_type else f"{step_no}. "
|
||||
lines.append(prefix + rationale)
|
||||
if len(lines) > 1:
|
||||
blocks.append("\n".join(lines))
|
||||
if not blocks:
|
||||
return ""
|
||||
return "## Reference Reasoning Steps\n" + "\n\n".join(blocks[:3])
|
||||
|
||||
def get_reference_metadata(self, item: dict) -> dict:
|
||||
reasoning_steps = item.get("reasoning_steps") or []
|
||||
path_count = 0
|
||||
preview_parts: list[str] = []
|
||||
for path in reasoning_steps:
|
||||
if not isinstance(path, list) or not path:
|
||||
continue
|
||||
path_count += 1
|
||||
first = path[0] if isinstance(path[0], dict) else {}
|
||||
step_type = str(first.get("reasoning type") or "").strip()
|
||||
rationale = str(first.get("rationale") or "").strip()
|
||||
preview_parts.append(f"[path {path_count}] {step_type}: {rationale[:180]}")
|
||||
if not path_count:
|
||||
return {"fields": [], "preview": ""}
|
||||
return {
|
||||
"fields": ["reasoning_steps"],
|
||||
"preview": "\n".join(preview_parts)[:500],
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
split_dir: str = "",
|
||||
data_path: str = "",
|
||||
split_mode: str = "ratio",
|
||||
split_ratio: str = "2:1:7",
|
||||
split_seed: int = 42,
|
||||
split_output_dir: str = "",
|
||||
max_turns: int = 1,
|
||||
workers: int = 16,
|
||||
analyst_workers: int = 16,
|
||||
failure_only: bool = False,
|
||||
minibatch_size: int = 8,
|
||||
edit_budget: int = 4,
|
||||
seed: int = 42,
|
||||
limit: int = 0,
|
||||
image_detail: str = "auto",
|
||||
use_deep_reflect: bool = False,
|
||||
deep_reflect_failures: int = 4,
|
||||
deep_reflect_successes: int = 2,
|
||||
) -> None:
|
||||
self.max_turns = max_turns
|
||||
self.workers = workers
|
||||
self.analyst_workers = analyst_workers
|
||||
self.failure_only = failure_only
|
||||
self.minibatch_size = minibatch_size
|
||||
self.edit_budget = edit_budget
|
||||
self.image_detail = image_detail
|
||||
self.use_deep_reflect = use_deep_reflect
|
||||
self.deep_reflect_failures = deep_reflect_failures
|
||||
self.deep_reflect_successes = deep_reflect_successes
|
||||
self.dataloader = MMRBDataLoader(
|
||||
split_dir=split_dir,
|
||||
data_path=data_path,
|
||||
split_mode=split_mode,
|
||||
split_ratio=split_ratio,
|
||||
split_seed=split_seed,
|
||||
split_output_dir=split_output_dir,
|
||||
seed=seed,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
def setup(self, cfg: dict) -> None:
|
||||
super().setup(cfg)
|
||||
self.dataloader.setup(cfg)
|
||||
|
||||
def get_dataloader(self):
|
||||
return self.dataloader
|
||||
|
||||
def build_env_from_batch(self, batch: BatchSpec, **kwargs):
|
||||
return list(batch.payload or [])
|
||||
|
||||
def build_train_env(self, batch_size: int, seed: int, **kwargs):
|
||||
batch = self.dataloader.build_train_batch(batch_size=batch_size, seed=seed, **kwargs)
|
||||
return self.build_env_from_batch(batch, **kwargs)
|
||||
|
||||
def build_eval_env(self, env_num: int, split: str, seed: int, **kwargs):
|
||||
batch = self.dataloader.build_eval_batch(env_num=env_num, split=split, seed=seed, **kwargs)
|
||||
return self.build_env_from_batch(batch, **kwargs)
|
||||
|
||||
def rollout(
|
||||
self,
|
||||
env_manager,
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict]:
|
||||
items: list[dict] = env_manager
|
||||
return run_batch(
|
||||
items=items,
|
||||
out_root=out_dir,
|
||||
skill_content=skill_content,
|
||||
max_turns=self.max_turns,
|
||||
workers=self.workers,
|
||||
image_detail=self.image_detail,
|
||||
diagnostic_mode=kwargs.get("diagnostic_mode", False),
|
||||
diagnostic_instruction=kwargs.get("diagnostic_instruction", ""),
|
||||
diagnostic_trace_context_by_id=kwargs.get("diagnostic_trace_context_by_id"),
|
||||
)
|
||||
|
||||
def reflect(
|
||||
self,
|
||||
results: list[dict],
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict | None]:
|
||||
prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
|
||||
patches_dir = kwargs.get("patches_dir", os.path.join(out_dir, "patches"))
|
||||
random_seed = kwargs.get("random_seed")
|
||||
step_buffer_context = kwargs.get("step_buffer_context", "")
|
||||
meta_skill_context = kwargs.get("meta_skill_context", "")
|
||||
|
||||
return run_minibatch_reflect(
|
||||
results=results,
|
||||
skill_content=skill_content,
|
||||
prediction_dir=prediction_dir,
|
||||
patches_dir=patches_dir,
|
||||
workers=self.analyst_workers,
|
||||
failure_only=self.failure_only,
|
||||
minibatch_size=self.minibatch_size,
|
||||
edit_budget=self.edit_budget,
|
||||
random_seed=random_seed,
|
||||
error_system=self.get_error_minibatch_prompt(),
|
||||
success_system=self.get_success_minibatch_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
meta_skill_context=meta_skill_context,
|
||||
update_mode=getattr(self, "_cfg", {}).get("skill_update_mode", "patch"),
|
||||
)
|
||||
|
||||
def deep_reflect(
|
||||
self,
|
||||
results: list[dict],
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict | None]:
|
||||
if not self.use_deep_reflect:
|
||||
return []
|
||||
|
||||
env_manager = kwargs.get("env_manager")
|
||||
prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
|
||||
random_seed = kwargs.get("random_seed")
|
||||
step_buffer_context = kwargs.get("step_buffer_context", "")
|
||||
meta_skill_context = kwargs.get("meta_skill_context", "")
|
||||
codex_backend = get_student_backend() == "codex_exec"
|
||||
selected_items = self.select_representative_items(
|
||||
results,
|
||||
env_manager if isinstance(env_manager, list) else None,
|
||||
n_failures=self.deep_reflect_failures,
|
||||
n_successes=self.deep_reflect_successes,
|
||||
seed=random_seed,
|
||||
)
|
||||
if not selected_items:
|
||||
return []
|
||||
selected_ids = {str(item["id"]) for item in selected_items}
|
||||
selected_results = [row for row in results if str(row.get("id")) in selected_ids]
|
||||
selected_examples = self.attach_reference_context(selected_results, selected_items)
|
||||
if codex_backend:
|
||||
selected_examples = self.attach_codex_probe_context(selected_examples, prediction_dir)
|
||||
|
||||
reasoning_count = 0
|
||||
selected_metadata = []
|
||||
for item in selected_items:
|
||||
meta = self.get_reference_metadata(item)
|
||||
if meta["fields"]:
|
||||
reasoning_count += 1
|
||||
selected_metadata.append({
|
||||
"id": str(item["id"]),
|
||||
"task_type": str(item.get("subtask") or item.get("task_type") or "mmrb"),
|
||||
"reference_fields": meta["fields"],
|
||||
"reference_preview": meta["preview"],
|
||||
})
|
||||
|
||||
deep_dir = os.path.join(out_dir, "deep_reflect")
|
||||
rollout_dir = os.path.join(deep_dir, "rollout")
|
||||
patches_dir = os.path.join(deep_dir, "patches")
|
||||
os.makedirs(deep_dir, exist_ok=True)
|
||||
print(
|
||||
f" [2b/6 DEEP REFLECT setup] selected={len(selected_items)} "
|
||||
f"reference_fields=reasoning_steps({reasoning_count}/{len(selected_items)})"
|
||||
)
|
||||
probe = generate_deep_probe_instruction(
|
||||
skill_content=skill_content,
|
||||
items=selected_examples,
|
||||
prediction_dir=prediction_dir,
|
||||
system_prompt=self.get_codex_deep_probe_prompt() if codex_backend else self.get_deep_probe_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
meta_skill_context=meta_skill_context,
|
||||
)
|
||||
if not probe:
|
||||
return []
|
||||
diagnostic_trace_context_by_id = None
|
||||
if codex_backend:
|
||||
selected_items, diagnostic_trace_context_by_id, probe = self.resolve_codex_probe_target(
|
||||
selected_items=selected_items,
|
||||
selected_examples=selected_examples,
|
||||
prediction_dir=prediction_dir,
|
||||
probe=probe,
|
||||
)
|
||||
probe_record = {
|
||||
**probe,
|
||||
"reference_summary": {
|
||||
"selected_count": len(selected_items),
|
||||
"field_counts": {"reasoning_steps": reasoning_count},
|
||||
},
|
||||
"selected_examples": selected_metadata,
|
||||
}
|
||||
with open(os.path.join(deep_dir, "probe.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(probe_record, f, ensure_ascii=False, indent=2)
|
||||
deep_results = run_batch(
|
||||
items=selected_items,
|
||||
out_root=rollout_dir,
|
||||
skill_content=skill_content,
|
||||
max_turns=self.max_turns,
|
||||
workers=min(self.workers, max(len(selected_items), 1)),
|
||||
image_detail=self.image_detail,
|
||||
diagnostic_mode=True,
|
||||
diagnostic_instruction=probe["probe_instruction"],
|
||||
diagnostic_trace_context_by_id=diagnostic_trace_context_by_id,
|
||||
)
|
||||
deep_results = self.attach_reference_context(deep_results, selected_items)
|
||||
return run_minibatch_reflect(
|
||||
results=deep_results,
|
||||
skill_content=skill_content,
|
||||
prediction_dir=os.path.join(rollout_dir, "predictions"),
|
||||
patches_dir=patches_dir,
|
||||
workers=self.analyst_workers,
|
||||
failure_only=self.failure_only,
|
||||
minibatch_size=self.minibatch_size,
|
||||
edit_budget=self.edit_budget,
|
||||
random_seed=random_seed,
|
||||
error_system=self.get_error_minibatch_prompt(),
|
||||
success_system=self.get_success_minibatch_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
meta_skill_context=meta_skill_context,
|
||||
update_mode=getattr(self, "_cfg", {}).get("skill_update_mode", "patch"),
|
||||
)
|
||||
|
||||
def get_task_types(self) -> list[str]:
|
||||
return self.dataloader.get_task_types()
|
||||
@@ -0,0 +1,146 @@
|
||||
"""MMRB task dataloader."""
|
||||
from __future__ import annotations
|
||||
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from skillopt.datasets.base import SplitDataLoader
|
||||
|
||||
|
||||
# ── Raw data loading utilities (for preprocessing / standalone eval) ─────
|
||||
|
||||
def _load_json(path: str) -> Any:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def _iter_data_files(data_path: str) -> list[str]:
|
||||
if not data_path:
|
||||
return []
|
||||
if os.path.isfile(data_path):
|
||||
return [data_path]
|
||||
if os.path.isdir(data_path):
|
||||
nested = glob.glob(os.path.join(data_path, "**", "*_human.json"), recursive=True)
|
||||
flat = glob.glob(os.path.join(data_path, "*_human.json"))
|
||||
return sorted(set(nested + flat))
|
||||
return []
|
||||
|
||||
|
||||
def _normalize_space(text: str) -> str:
|
||||
return re.sub(r"\s+", " ", str(text or "").strip())
|
||||
|
||||
|
||||
def _normalize_item(item: dict, row_idx: int, source_path: str) -> dict | None:
|
||||
question = _normalize_space(item.get("question") or "")
|
||||
answer = _normalize_space(item.get("answer") or "")
|
||||
raw_image_paths = item.get("image_paths") or []
|
||||
if not question or not answer or not isinstance(raw_image_paths, list) or not raw_image_paths:
|
||||
return None
|
||||
|
||||
base_dir = os.path.dirname(source_path)
|
||||
image_paths: list[str] = []
|
||||
for raw_path in raw_image_paths:
|
||||
rel = str(raw_path or "").strip()
|
||||
if not rel:
|
||||
continue
|
||||
abs_path = rel if os.path.isabs(rel) else os.path.abspath(os.path.join(base_dir, rel))
|
||||
if os.path.exists(abs_path):
|
||||
image_paths.append(abs_path)
|
||||
if not image_paths:
|
||||
return None
|
||||
|
||||
options_raw = item.get("options") or []
|
||||
options = [_normalize_space(opt) for opt in options_raw if _normalize_space(opt)]
|
||||
source = _normalize_space(item.get("source") or "unknown")
|
||||
subtask = _normalize_space(item.get("subtask") or "unknown")
|
||||
item_index = item.get("index", row_idx)
|
||||
item_id = f"{source}:{subtask}:{item_index}"
|
||||
|
||||
return {
|
||||
"id": item_id,
|
||||
"source": source,
|
||||
"subtask": subtask,
|
||||
"task_type": subtask,
|
||||
"question": question,
|
||||
"answer": answer,
|
||||
"options": options,
|
||||
"is_choice": bool(options),
|
||||
"image_paths": image_paths,
|
||||
"reasoning_steps": item.get("reasoning_steps") or [],
|
||||
"annotation_time": item.get("annotation_time"),
|
||||
"source_path": os.path.abspath(source_path),
|
||||
}
|
||||
|
||||
|
||||
def load_items(data_path: str) -> list[dict]:
|
||||
"""Load and normalise MMRB items from JSON files."""
|
||||
files = _iter_data_files(data_path)
|
||||
if not files:
|
||||
raise ValueError(
|
||||
"MMRB requires data_path to be a *_human.json file or a directory "
|
||||
"containing extracted MMRB subtask folders."
|
||||
)
|
||||
|
||||
items: list[dict] = []
|
||||
for path in files:
|
||||
raw = _load_json(path)
|
||||
if not isinstance(raw, list):
|
||||
raise ValueError(f"Expected JSON array in {path}, got {type(raw).__name__}")
|
||||
for row_idx, item in enumerate(raw):
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
norm = _normalize_item(item, row_idx=row_idx, source_path=path)
|
||||
if norm is not None:
|
||||
items.append(norm)
|
||||
|
||||
if not items:
|
||||
raise ValueError(f"No valid MMRB items loaded from {data_path}")
|
||||
return items
|
||||
|
||||
|
||||
# ── Dataloader ───────────────────────────────────────────────────────────
|
||||
|
||||
class MMRBDataLoader(SplitDataLoader):
|
||||
"""MMRB dataloader."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
split_dir: str = "",
|
||||
data_path: str = "",
|
||||
split_mode: str = "ratio",
|
||||
split_ratio: str = "2:1:7",
|
||||
split_seed: int = 42,
|
||||
split_output_dir: str = "",
|
||||
seed: int = 42,
|
||||
limit: int = 0,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
split_dir=split_dir,
|
||||
data_path=data_path,
|
||||
split_mode=split_mode,
|
||||
split_ratio=split_ratio,
|
||||
split_seed=split_seed,
|
||||
split_output_dir=split_output_dir,
|
||||
seed=seed,
|
||||
limit=limit,
|
||||
)
|
||||
self._task_types: list[str] = []
|
||||
|
||||
def load_raw_items(self, data_path: str) -> list[dict]:
|
||||
return load_items(data_path)
|
||||
|
||||
def setup(self, cfg: dict) -> None:
|
||||
super().setup(cfg)
|
||||
all_items = self.train_items + self.val_items + self.test_items
|
||||
task_types = {
|
||||
item.get("subtask") or item.get("task_type") or "unknown"
|
||||
for item in all_items
|
||||
}
|
||||
self._task_types = sorted(task_types)
|
||||
|
||||
def get_task_types(self) -> list[str]:
|
||||
return list(self._task_types)
|
||||
@@ -0,0 +1,102 @@
|
||||
"""MMRB evaluation helpers."""
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import string
|
||||
|
||||
|
||||
_EVAL_MODE = "mmrb_exact_match_v1"
|
||||
|
||||
|
||||
def normalize_text(text: str) -> str:
|
||||
text = str(text or "").strip().lower()
|
||||
text = "".join(ch for ch in text if ch not in string.punctuation)
|
||||
return " ".join(text.split())
|
||||
|
||||
|
||||
def extract_answer(text: str | None) -> str:
|
||||
raw = str(text or "").strip()
|
||||
if not raw:
|
||||
return ""
|
||||
|
||||
answer_tags = re.findall(r"<answer>\s*(.*?)\s*</answer>", raw, re.IGNORECASE | re.DOTALL)
|
||||
if answer_tags:
|
||||
return answer_tags[-1].strip()
|
||||
|
||||
bracket = re.findall(r"Answer\s*\[\s*(.*?)\s*\]", raw, re.IGNORECASE | re.DOTALL)
|
||||
if bracket:
|
||||
return bracket[-1].strip()
|
||||
|
||||
boxed = re.findall(r"\\boxed\{(.*?)\}", raw, re.IGNORECASE | re.DOTALL)
|
||||
if boxed:
|
||||
return boxed[-1].strip()
|
||||
|
||||
single = raw.strip().rstrip(".):")
|
||||
if re.fullmatch(r"[A-Z]", single, re.IGNORECASE):
|
||||
return single.strip()
|
||||
|
||||
patterns = [
|
||||
r"final answer\s*(?:is)?\s*[::]?\s*(.+)",
|
||||
r"the answer is\s*[::]?\s*(.+)",
|
||||
r"answer\s*[::]?\s*(.+)$",
|
||||
]
|
||||
for pattern in patterns:
|
||||
match = re.search(pattern, raw, re.IGNORECASE)
|
||||
if match:
|
||||
return match.group(1).strip().strip("*")
|
||||
|
||||
return raw
|
||||
|
||||
|
||||
def evaluate_item(*, item: dict, prediction_text: str) -> dict:
|
||||
predicted_answer = extract_answer(prediction_text)
|
||||
gold_answer = str(item.get("answer") or "").strip()
|
||||
predicted_norm = normalize_text(predicted_answer)
|
||||
gold_norm = normalize_text(gold_answer)
|
||||
|
||||
hard = 0.0
|
||||
matched_gold = ""
|
||||
predicted_label = ""
|
||||
predicted_text = predicted_answer
|
||||
|
||||
if item.get("is_choice"):
|
||||
predicted_label = str(predicted_answer).strip().upper().rstrip(".):")
|
||||
if predicted_label == str(gold_answer).strip().upper():
|
||||
hard = 1.0
|
||||
matched_gold = gold_answer
|
||||
else:
|
||||
for option in item.get("options") or []:
|
||||
label_match = re.match(r"\(?([A-Z])\)", option)
|
||||
if not label_match:
|
||||
continue
|
||||
label = label_match.group(1).upper()
|
||||
option_text = option[label_match.end():].strip(" .:-")
|
||||
if predicted_norm and normalize_text(option_text) == predicted_norm:
|
||||
predicted_label = label
|
||||
predicted_text = option_text
|
||||
break
|
||||
if predicted_label == str(gold_answer).strip().upper():
|
||||
hard = 1.0
|
||||
matched_gold = gold_answer
|
||||
else:
|
||||
if predicted_norm and gold_norm and (
|
||||
predicted_norm == gold_norm or predicted_norm in gold_norm or gold_norm in predicted_norm
|
||||
):
|
||||
hard = 1.0
|
||||
matched_gold = gold_answer
|
||||
|
||||
return {
|
||||
"evaluation_mode": _EVAL_MODE,
|
||||
"predicted_answer": predicted_answer,
|
||||
"predicted_label": predicted_label,
|
||||
"predicted_text": predicted_text,
|
||||
"em": hard,
|
||||
"f1": hard,
|
||||
"sub_em": hard,
|
||||
"matched_gold": matched_gold,
|
||||
}
|
||||
|
||||
|
||||
def evaluation_mode() -> str:
|
||||
return _EVAL_MODE
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
You are an expert multi-image reasoning agent.
|
||||
|
||||
{skill_section}## Task Format
|
||||
You will receive a question grounded in multiple images.
|
||||
Use the image order exactly as presented in the prompt and compare evidence across images carefully.
|
||||
|
||||
## Answer Format
|
||||
- Put the final answer inside <answer>...</answer>.
|
||||
- For multiple-choice questions, output only the single option letter inside <answer>...</answer>.
|
||||
- For open questions, output only the short final answer inside <answer>...</answer>.
|
||||
@@ -0,0 +1,455 @@
|
||||
"""MMRB rollout."""
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
import re
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
from skillopt.envs.mmrb.evaluator import evaluate_item, evaluation_mode
|
||||
from skillopt.model import chat_student_messages, get_student_backend, is_student_exec_backend
|
||||
from skillopt.model.codex_harness import prepare_workspace, render_skill_md, run_student_exec
|
||||
from skillopt.prompts import load_prompt
|
||||
|
||||
_IMAGE_REF_RE = re.compile(r"\{image#(\d+)\}", re.IGNORECASE)
|
||||
|
||||
|
||||
def _build_system(skill_content: str) -> str:
|
||||
if skill_content.strip():
|
||||
skill_section = f"## Skill\n{skill_content.strip()}\n\n"
|
||||
else:
|
||||
skill_section = ""
|
||||
return load_prompt("rollout_system", env="mmrb").format(skill_section=skill_section)
|
||||
|
||||
|
||||
def _image_to_data_uri(path: str) -> str:
|
||||
mime = mimetypes.guess_type(path)[0] or "image/png"
|
||||
with open(path, "rb") as f:
|
||||
encoded = base64.b64encode(f.read()).decode("ascii")
|
||||
return f"data:{mime};base64,{encoded}"
|
||||
|
||||
|
||||
def _build_user_content(
|
||||
item: dict,
|
||||
image_detail: str,
|
||||
*,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context: str = "",
|
||||
) -> tuple[list[dict], str]:
|
||||
raw_question = str(item["question"])
|
||||
content: list[dict] = []
|
||||
text_parts: list[str] = []
|
||||
used_indices: set[int] = set()
|
||||
cursor = 0
|
||||
|
||||
if diagnostic_trace_context.strip():
|
||||
prefix = (
|
||||
"## Previous Codex Trace Snapshot\n"
|
||||
"This is a partial transcript from an earlier attempt. Use it as your current reasoning context.\n\n"
|
||||
f"{diagnostic_trace_context.strip()}\n\n"
|
||||
)
|
||||
content.append({"type": "text", "text": prefix})
|
||||
text_parts.append(prefix)
|
||||
|
||||
for match in _IMAGE_REF_RE.finditer(raw_question):
|
||||
if match.start() > cursor:
|
||||
chunk = raw_question[cursor:match.start()]
|
||||
if chunk:
|
||||
content.append({"type": "text", "text": chunk})
|
||||
text_parts.append(chunk)
|
||||
|
||||
image_idx = int(match.group(1)) - 1
|
||||
marker = f"[Image #{image_idx + 1}]"
|
||||
text_parts.append(marker)
|
||||
if 0 <= image_idx < len(item["image_paths"]):
|
||||
image_url = {"url": _image_to_data_uri(item["image_paths"][image_idx])}
|
||||
if image_detail and image_detail != "auto":
|
||||
image_url["detail"] = image_detail
|
||||
content.append({"type": "image_url", "image_url": image_url})
|
||||
used_indices.add(image_idx)
|
||||
else:
|
||||
content.append({"type": "text", "text": marker})
|
||||
cursor = match.end()
|
||||
|
||||
if cursor < len(raw_question):
|
||||
tail = raw_question[cursor:]
|
||||
if tail:
|
||||
content.append({"type": "text", "text": tail})
|
||||
text_parts.append(tail)
|
||||
|
||||
for idx, path in enumerate(item["image_paths"]):
|
||||
if idx in used_indices:
|
||||
continue
|
||||
marker = f"\n[Additional Image #{idx + 1}]"
|
||||
text_parts.append(marker)
|
||||
content.append({"type": "text", "text": marker})
|
||||
image_url = {"url": _image_to_data_uri(path)}
|
||||
if image_detail and image_detail != "auto":
|
||||
image_url["detail"] = image_detail
|
||||
content.append({"type": "image_url", "image_url": image_url})
|
||||
|
||||
answer_instruction = (
|
||||
"\n\nAnswer with the single correct option letter inside <answer>...</answer>."
|
||||
if item.get("is_choice")
|
||||
else "\n\nAnswer with the short final answer inside <answer>...</answer>."
|
||||
)
|
||||
content.append({"type": "text", "text": answer_instruction})
|
||||
text_parts.append(answer_instruction)
|
||||
|
||||
if diagnostic_mode and diagnostic_instruction.strip():
|
||||
diag_block = f"\n\n## Training Readout\n{diagnostic_instruction.strip()}"
|
||||
content.append({"type": "text", "text": diag_block})
|
||||
text_parts.append(diag_block)
|
||||
|
||||
return content, "".join(text_parts)
|
||||
|
||||
|
||||
def _build_messages(
|
||||
item: dict,
|
||||
skill_content: str,
|
||||
image_detail: str,
|
||||
*,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
) -> tuple[list[dict], str, str]:
|
||||
system = _build_system(skill_content)
|
||||
user_content, user_text = _build_user_content(
|
||||
item,
|
||||
image_detail,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
)
|
||||
messages = [
|
||||
{"role": "system", "content": system},
|
||||
{"role": "user", "content": user_content},
|
||||
]
|
||||
return messages, system, user_text
|
||||
|
||||
|
||||
def _build_codex_skill(skill_content: str) -> str:
|
||||
return render_skill_md(
|
||||
skill_content,
|
||||
description="Dynamic ReflACT skill for solving the current MMRB multi-image reasoning question.",
|
||||
preamble=(
|
||||
"Use this skill when solving the current multi-image reasoning task.\n"
|
||||
"Inspect all attached images carefully and return the final answer inside <answer>...</answer>."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _run_codex_once(
|
||||
*,
|
||||
pred_dir: str,
|
||||
item: dict,
|
||||
skill_content: str,
|
||||
model: str,
|
||||
timeout: int,
|
||||
image_detail: str,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context: str = "",
|
||||
previous_response: str = "",
|
||||
) -> tuple[str, str, str, str]:
|
||||
user_text = _build_user_content(
|
||||
item,
|
||||
image_detail,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
diagnostic_trace_context=diagnostic_trace_context,
|
||||
)[1]
|
||||
task_parts = [user_text]
|
||||
if previous_response:
|
||||
task_parts.append(
|
||||
"## Previous Attempt\n"
|
||||
f"{previous_response}\n\n"
|
||||
"Review the same images carefully and answer again."
|
||||
)
|
||||
task_text = "\n\n".join(task_parts)
|
||||
skill_md = _build_codex_skill(skill_content)
|
||||
work_dir = os.path.join(pred_dir, "codex_exec")
|
||||
prepare_workspace(
|
||||
work_dir=work_dir,
|
||||
skill_md=skill_md,
|
||||
task_text=task_text,
|
||||
images=item["image_paths"],
|
||||
)
|
||||
prompt = (
|
||||
"Use the `skillopt-student` skill available in this workspace.\n"
|
||||
"Read `task.md`, inspect all attached images, and answer the question.\n"
|
||||
"Keep the final answer inside <answer>...</answer>."
|
||||
)
|
||||
final_message, raw = run_student_exec(
|
||||
work_dir=work_dir,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
timeout=timeout,
|
||||
images=item["image_paths"],
|
||||
)
|
||||
return final_message or raw, raw, skill_md, task_text
|
||||
|
||||
|
||||
def process_one(
|
||||
item: dict,
|
||||
out_root: str,
|
||||
skill_content: str,
|
||||
*,
|
||||
max_turns: int = 1,
|
||||
image_detail: str = "auto",
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context: str = "",
|
||||
) -> dict:
|
||||
item_id = str(item["id"])
|
||||
result = {
|
||||
"id": item_id,
|
||||
"question": item["question"],
|
||||
"task_type": item.get("subtask") or item.get("task_type") or "mmrb",
|
||||
"task_description": item["question"],
|
||||
"hard": 0,
|
||||
"soft": 0.0,
|
||||
"predicted_answer": "",
|
||||
"predicted_label": "",
|
||||
"predicted_text": "",
|
||||
"response": "",
|
||||
"fail_reason": "",
|
||||
"agent_ok": False,
|
||||
"n_turns": 0,
|
||||
"image_paths": item["image_paths"],
|
||||
"gold_answer": item["answer"],
|
||||
"evaluation_mode": evaluation_mode(),
|
||||
}
|
||||
|
||||
try:
|
||||
pred_dir = os.path.join(out_root, "predictions", item_id)
|
||||
os.makedirs(pred_dir, exist_ok=True)
|
||||
|
||||
if is_student_exec_backend():
|
||||
from skillopt.model import azure_openai as _llm
|
||||
|
||||
response = ""
|
||||
conversation: list[dict] = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": item["question"] + "\n\n" + "\n".join(
|
||||
f"[image] {os.path.basename(path)}" for path in item["image_paths"]
|
||||
),
|
||||
}
|
||||
]
|
||||
system_prompt = ""
|
||||
user_text = ""
|
||||
for turn in range(max_turns):
|
||||
response, raw, system_prompt, user_text = _run_codex_once(
|
||||
pred_dir=pred_dir,
|
||||
item=item,
|
||||
skill_content=skill_content,
|
||||
model=_llm.STUDENT_DEPLOYMENT,
|
||||
timeout=120,
|
||||
image_detail=image_detail,
|
||||
diagnostic_mode=diagnostic_mode if turn == 0 else False,
|
||||
diagnostic_instruction=diagnostic_instruction if turn == 0 else "",
|
||||
diagnostic_trace_context=diagnostic_trace_context if turn == 0 else "",
|
||||
previous_response=response if turn > 0 else "",
|
||||
)
|
||||
conversation.append({"type": "message", "turn": turn + 1, "content": response})
|
||||
if "<answer>" in response.lower():
|
||||
break
|
||||
|
||||
result["response"] = response
|
||||
result["agent_ok"] = True
|
||||
result["n_turns"] = len(conversation) - 1
|
||||
with open(os.path.join(pred_dir, "student_system_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(system_prompt)
|
||||
with open(os.path.join(pred_dir, "student_user_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(user_text)
|
||||
|
||||
eval_result = evaluate_item(item=item, prediction_text=response)
|
||||
result["evaluation_mode"] = eval_result["evaluation_mode"]
|
||||
result["predicted_answer"] = eval_result["predicted_answer"]
|
||||
result["predicted_label"] = eval_result["predicted_label"]
|
||||
result["predicted_text"] = eval_result["predicted_text"]
|
||||
result["matched_gold"] = eval_result["matched_gold"]
|
||||
result["hard"] = int(eval_result["em"])
|
||||
result["soft"] = eval_result["f1"]
|
||||
if not result["hard"]:
|
||||
result["fail_reason"] = (
|
||||
f"predicted '{eval_result['predicted_answer']}' but expected '{item['answer']}'"
|
||||
)
|
||||
eval_detail = (
|
||||
"[EVALUATION RESULT]\n"
|
||||
f"Question: {item['question']}\n"
|
||||
f"Predicted answer: {eval_result['predicted_answer']!r}\n"
|
||||
f"Predicted label: {eval_result['predicted_label']!r}\n"
|
||||
f"Gold answer: {item['answer']!r}\n"
|
||||
f"Correct: {eval_result['em']}\n"
|
||||
)
|
||||
conversation.append({"role": "system", "content": eval_detail})
|
||||
with open(os.path.join(pred_dir, "conversation.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(conversation, f, ensure_ascii=False, indent=2)
|
||||
return result
|
||||
|
||||
messages, system_prompt, user_text = _build_messages(
|
||||
item,
|
||||
skill_content,
|
||||
image_detail,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
diagnostic_trace_context=diagnostic_trace_context,
|
||||
)
|
||||
response = ""
|
||||
conversation: list[dict] = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": user_text + "\n\n" + "\n".join(
|
||||
f"[image] {os.path.basename(path)}" for path in item["image_paths"]
|
||||
),
|
||||
}
|
||||
]
|
||||
|
||||
for turn in range(max_turns):
|
||||
if turn == 0:
|
||||
resp_text, _ = chat_student_messages(
|
||||
messages=messages,
|
||||
max_completion_tokens=768,
|
||||
retries=5,
|
||||
stage="rollout",
|
||||
)
|
||||
else:
|
||||
refinement_messages = [
|
||||
messages[0],
|
||||
messages[1],
|
||||
{"role": "assistant", "content": response},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Review the same images carefully and answer again. Keep the final answer inside <answer>...</answer>.",
|
||||
},
|
||||
]
|
||||
resp_text, _ = chat_student_messages(
|
||||
messages=refinement_messages,
|
||||
max_completion_tokens=512,
|
||||
retries=5,
|
||||
stage="rollout",
|
||||
)
|
||||
response = resp_text
|
||||
conversation.append({"type": "message", "turn": turn + 1, "content": resp_text})
|
||||
if "<answer>" in resp_text.lower():
|
||||
break
|
||||
|
||||
result["response"] = response
|
||||
result["agent_ok"] = True
|
||||
result["n_turns"] = len(conversation) - 1
|
||||
|
||||
with open(os.path.join(pred_dir, "student_system_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(system_prompt)
|
||||
with open(os.path.join(pred_dir, "student_user_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(user_text)
|
||||
|
||||
eval_result = evaluate_item(item=item, prediction_text=response)
|
||||
result["evaluation_mode"] = eval_result["evaluation_mode"]
|
||||
result["predicted_answer"] = eval_result["predicted_answer"]
|
||||
result["predicted_label"] = eval_result["predicted_label"]
|
||||
result["predicted_text"] = eval_result["predicted_text"]
|
||||
result["matched_gold"] = eval_result["matched_gold"]
|
||||
result["hard"] = int(eval_result["em"])
|
||||
result["soft"] = eval_result["f1"]
|
||||
if not result["hard"]:
|
||||
result["fail_reason"] = (
|
||||
f"predicted '{eval_result['predicted_answer']}' but expected '{item['answer']}'"
|
||||
)
|
||||
|
||||
eval_detail = (
|
||||
"[EVALUATION RESULT]\n"
|
||||
f"Question: {item['question']}\n"
|
||||
f"Predicted answer: {eval_result['predicted_answer']!r}\n"
|
||||
f"Predicted label: {eval_result['predicted_label']!r}\n"
|
||||
f"Gold answer: {item['answer']!r}\n"
|
||||
f"Correct: {eval_result['em']}\n"
|
||||
)
|
||||
conversation.append({"role": "system", "content": eval_detail})
|
||||
with open(os.path.join(pred_dir, "conversation.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(conversation, f, ensure_ascii=False, indent=2)
|
||||
except Exception as e: # noqa: BLE001
|
||||
result["fail_reason"] = f"error: {e}"
|
||||
return result
|
||||
|
||||
|
||||
def run_batch(
|
||||
items: list[dict],
|
||||
out_root: str,
|
||||
skill_content: str,
|
||||
*,
|
||||
max_turns: int = 1,
|
||||
workers: int = 16,
|
||||
image_detail: str = "auto",
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
diagnostic_trace_context_by_id: dict[str, str] | None = None,
|
||||
) -> list[dict]:
|
||||
results_path = os.path.join(out_root, "results.jsonl")
|
||||
os.makedirs(out_root, exist_ok=True)
|
||||
|
||||
expected_eval_mode = evaluation_mode()
|
||||
done_ids: set[str] = set()
|
||||
existing: list[dict] = []
|
||||
rewrite_results = False
|
||||
if os.path.exists(results_path):
|
||||
with open(results_path, encoding="utf-8") as f:
|
||||
for line in f:
|
||||
try:
|
||||
row = json.loads(line)
|
||||
if row.get("evaluation_mode") != expected_eval_mode:
|
||||
rewrite_results = True
|
||||
continue
|
||||
done_ids.add(str(row["id"]))
|
||||
existing.append(row)
|
||||
except Exception:
|
||||
rewrite_results = True
|
||||
|
||||
pending = [item for item in items if str(item["id"]) not in done_ids]
|
||||
if not pending and not rewrite_results:
|
||||
return existing
|
||||
|
||||
total = len(existing) + len(pending)
|
||||
completed = len(existing)
|
||||
correct_count = sum(1 for r in existing if r.get("hard", 0))
|
||||
if existing:
|
||||
print(f" [rollout] resuming: {completed}/{total} already done", flush=True)
|
||||
|
||||
results = list(existing)
|
||||
file_mode = "w" if rewrite_results else "a"
|
||||
with open(results_path, file_mode, encoding="utf-8") as outf, ThreadPoolExecutor(max_workers=workers) as ex:
|
||||
if rewrite_results:
|
||||
for row in existing:
|
||||
outf.write(json.dumps(row, ensure_ascii=False) + "\n")
|
||||
futs = {
|
||||
ex.submit(
|
||||
process_one,
|
||||
item,
|
||||
out_root,
|
||||
skill_content,
|
||||
max_turns=max_turns,
|
||||
image_detail=image_detail,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
diagnostic_trace_context=(diagnostic_trace_context_by_id or {}).get(str(item["id"]), ""),
|
||||
): item
|
||||
for item in pending
|
||||
}
|
||||
for fut in as_completed(futs):
|
||||
row = fut.result()
|
||||
results.append(row)
|
||||
completed += 1
|
||||
if row.get("hard", 0):
|
||||
correct_count += 1
|
||||
acc = correct_count / completed if completed else 0
|
||||
print(
|
||||
f" [rollout] {completed}/{total} "
|
||||
f"(acc={acc:.3f}) id={row.get('id', '?')} "
|
||||
f"hard={row.get('hard', '?')}",
|
||||
flush=True,
|
||||
)
|
||||
outf.write(json.dumps(row, ensure_ascii=False) + "\n")
|
||||
outf.flush()
|
||||
return results
|
||||
@@ -0,0 +1,17 @@
|
||||
# MMRB Multi-Image Reasoning Heuristics
|
||||
|
||||
## Cross-Image Alignment
|
||||
- Track the role of each image by its index and compare evidence across all referenced images before deciding.
|
||||
- When the question depends on sequence, correspondence, or retrieval, verify the relation between images instead of judging each image independently.
|
||||
|
||||
## Option Elimination
|
||||
- For multiple-choice tasks, compare all options and reject choices that match only part of the visual evidence.
|
||||
- If options differ by a small visual detail, use the most discriminative cue rather than a coarse scene impression.
|
||||
|
||||
## Open Answers
|
||||
- For open-ended tasks, give the shortest answer that is fully supported by the combined images.
|
||||
- Preserve exact entities, attributes, counts, and directions when the images support them directly.
|
||||
|
||||
## Final Answer
|
||||
- Output only the final answer inside <answer>...</answer>.
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""OfficeQA environment package for ReflACT."""
|
||||
@@ -0,0 +1,174 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
from skillopt.datasets.base import BatchSpec
|
||||
from skillopt.envs.base import EnvAdapter
|
||||
from skillopt.envs.deep_reflect import run_no_reference_deep_reflect
|
||||
from skillopt.envs.officeqa.dataloader import OfficeQADataLoader
|
||||
from skillopt.envs.officeqa.rollout import run_batch
|
||||
from skillopt.gradient.reflect import run_minibatch_reflect
|
||||
|
||||
|
||||
class OfficeQAAdapter(EnvAdapter):
|
||||
def __init__(
|
||||
self,
|
||||
split_dir: str = "",
|
||||
data_path: str = "",
|
||||
split_mode: str = "split_dir",
|
||||
split_ratio: str = "2:1:7",
|
||||
split_seed: int = 42,
|
||||
split_output_dir: str = "",
|
||||
workers: int = 8,
|
||||
analyst_workers: int = 8,
|
||||
failure_only: bool = False,
|
||||
minibatch_size: int = 8,
|
||||
edit_budget: int = 4,
|
||||
seed: int = 42,
|
||||
limit: int = 0,
|
||||
max_tool_turns: int = 12,
|
||||
max_completion_tokens: int = 64000,
|
||||
search_mode: str = "offline",
|
||||
max_queries_per_turn: int = 4,
|
||||
search_api_url: str = os.environ.get("OFFICEQA_SEARCH_API_URL", "http://localhost:8080/search_tool/search"),
|
||||
search_auth_env: str = "OFFICEQA_CUSTOM_SEARCH_AUTH",
|
||||
search_provider: str = "duckduckgo",
|
||||
search_max_num_results: int = 4,
|
||||
search_timeout_seconds: int = 20,
|
||||
use_local_tools: bool = True,
|
||||
data_dirs: list[str] | str | None = None,
|
||||
docs_dirs: list[str] | str | None = None,
|
||||
use_deep_reflect: bool = False,
|
||||
deep_reflect_failures: int = 4,
|
||||
deep_reflect_successes: int = 2,
|
||||
) -> None:
|
||||
self.workers = workers
|
||||
self.analyst_workers = analyst_workers
|
||||
self.failure_only = failure_only
|
||||
self.minibatch_size = minibatch_size
|
||||
self.edit_budget = edit_budget
|
||||
self.max_tool_turns = max_tool_turns
|
||||
self.max_completion_tokens = int(max_completion_tokens)
|
||||
self.search_mode = str(search_mode or "offline")
|
||||
self.max_queries_per_turn = int(max_queries_per_turn)
|
||||
self.search_api_url = str(search_api_url or "").strip()
|
||||
self.search_auth_env = str(search_auth_env or "OFFICEQA_CUSTOM_SEARCH_AUTH").strip()
|
||||
self.search_provider = str(search_provider or "duckduckgo").strip()
|
||||
self.search_max_num_results = int(search_max_num_results)
|
||||
self.search_timeout_seconds = int(search_timeout_seconds)
|
||||
self.use_local_tools = bool(use_local_tools)
|
||||
self.data_dirs = data_dirs if data_dirs is not None else docs_dirs
|
||||
self.use_deep_reflect = use_deep_reflect
|
||||
self.deep_reflect_failures = deep_reflect_failures
|
||||
self.deep_reflect_successes = deep_reflect_successes
|
||||
self.dataloader = OfficeQADataLoader(
|
||||
split_dir=split_dir,
|
||||
data_path=data_path,
|
||||
split_mode=split_mode,
|
||||
split_ratio=split_ratio,
|
||||
split_seed=split_seed,
|
||||
split_output_dir=split_output_dir,
|
||||
seed=seed,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
def setup(self, cfg: dict) -> None:
|
||||
super().setup(cfg)
|
||||
self.dataloader.setup(cfg)
|
||||
|
||||
def get_dataloader(self):
|
||||
return self.dataloader
|
||||
|
||||
def build_env_from_batch(self, batch: BatchSpec, **kwargs):
|
||||
return list(batch.payload or [])
|
||||
|
||||
def build_train_env(self, batch_size: int, seed: int, **kwargs):
|
||||
batch = self.dataloader.build_train_batch(batch_size=batch_size, seed=seed, **kwargs)
|
||||
return self.build_env_from_batch(batch, **kwargs)
|
||||
|
||||
def build_eval_env(self, env_num: int, split: str, seed: int, **kwargs):
|
||||
batch = self.dataloader.build_eval_batch(env_num=env_num, split=split, seed=seed, **kwargs)
|
||||
return self.build_env_from_batch(batch, **kwargs)
|
||||
|
||||
def rollout(self, env_manager, skill_content: str, out_dir: str, **kwargs) -> list[dict]:
|
||||
items: list[dict] = env_manager
|
||||
return run_batch(
|
||||
items=items,
|
||||
out_root=out_dir,
|
||||
skill_content=skill_content,
|
||||
workers=self.workers,
|
||||
max_tool_turns=self.max_tool_turns,
|
||||
max_completion_tokens=self.max_completion_tokens,
|
||||
search_mode=self.search_mode,
|
||||
max_queries_per_turn=self.max_queries_per_turn,
|
||||
search_api_url=self.search_api_url,
|
||||
search_auth_env=self.search_auth_env,
|
||||
search_provider=self.search_provider,
|
||||
search_max_num_results=self.search_max_num_results,
|
||||
search_timeout_seconds=self.search_timeout_seconds,
|
||||
use_local_tools=self.use_local_tools,
|
||||
data_dirs=self.data_dirs,
|
||||
diagnostic_mode=kwargs.get("diagnostic_mode", False),
|
||||
diagnostic_instruction=kwargs.get("diagnostic_instruction", ""),
|
||||
)
|
||||
|
||||
def reflect(self, results: list[dict], skill_content: str, out_dir: str, **kwargs) -> list[dict | None]:
|
||||
prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
|
||||
patches_dir = kwargs.get("patches_dir", os.path.join(out_dir, "patches"))
|
||||
random_seed = kwargs.get("random_seed")
|
||||
step_buffer_context = kwargs.get("step_buffer_context", "")
|
||||
return run_minibatch_reflect(
|
||||
results=results,
|
||||
skill_content=skill_content,
|
||||
prediction_dir=prediction_dir,
|
||||
patches_dir=patches_dir,
|
||||
workers=self.analyst_workers,
|
||||
failure_only=self.failure_only,
|
||||
minibatch_size=self.minibatch_size,
|
||||
edit_budget=self.edit_budget,
|
||||
random_seed=random_seed,
|
||||
error_system=self.get_error_minibatch_prompt(),
|
||||
success_system=self.get_success_minibatch_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
update_mode=getattr(self, "_cfg", {}).get("skill_update_mode", "patch"),
|
||||
)
|
||||
|
||||
def deep_reflect(
|
||||
self,
|
||||
results: list[dict],
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict | None]:
|
||||
return run_no_reference_deep_reflect(
|
||||
self,
|
||||
results,
|
||||
skill_content,
|
||||
out_dir,
|
||||
env_manager=kwargs.get("env_manager"),
|
||||
prediction_dir=kwargs.get("prediction_dir"),
|
||||
random_seed=kwargs.get("random_seed"),
|
||||
step_buffer_context=kwargs.get("step_buffer_context", ""),
|
||||
output_requirements=[
|
||||
"- There is no hidden reference block. Use only the question, candidate files, tool trace, student output, and evaluation result to infer what intermediate state is worth probing.",
|
||||
"- The instruction must explicitly request a short <analysis>...</analysis> block before the final <answer>...</answer>.",
|
||||
"- The readout should focus on selected document/file, evidence span or table, extracted value, units, and any date or fiscal-period normalization.",
|
||||
"- Do not ask for exhaustive copying of source text or a full chain-of-thought.",
|
||||
"- The instruction text should be ready to append directly to the student's prompt.",
|
||||
],
|
||||
metadata_builder=lambda item: {
|
||||
"id": str(item.get("id")),
|
||||
"task_type": str(item.get("task_type") or "officeqa"),
|
||||
"question_preview": str(item.get("question") or "")[:200],
|
||||
"source_files": item.get("source_files", []),
|
||||
"source_docs": item.get("source_docs", []),
|
||||
},
|
||||
)
|
||||
|
||||
def get_task_types(self) -> list[str]:
|
||||
seen: list[str] = []
|
||||
for item in self.dataloader.train_items + self.dataloader.val_items + self.dataloader.test_items:
|
||||
task_type = str(item.get("task_type") or "officeqa")
|
||||
if task_type not in seen:
|
||||
seen.append(task_type)
|
||||
return seen or ["officeqa"]
|
||||
@@ -0,0 +1,71 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from skillopt.datasets.base import SplitDataLoader
|
||||
|
||||
|
||||
def _parse_list_field(value: str | list[str] | None) -> list[str]:
|
||||
if value is None:
|
||||
return []
|
||||
if isinstance(value, list):
|
||||
return [str(item).strip() for item in value if str(item).strip()]
|
||||
text = str(value).strip()
|
||||
if not text:
|
||||
return []
|
||||
try:
|
||||
loaded = json.loads(text)
|
||||
except json.JSONDecodeError:
|
||||
loaded = None
|
||||
if isinstance(loaded, list):
|
||||
return [str(item).strip() for item in loaded if str(item).strip()]
|
||||
if "\n" in text:
|
||||
return [part.strip() for part in text.splitlines() if part.strip()]
|
||||
if "," in text and not text.lower().endswith(".txt"):
|
||||
return [part.strip() for part in text.split(",") if part.strip()]
|
||||
return [text]
|
||||
|
||||
|
||||
def _normalize_row(row: dict[str, str]) -> dict:
|
||||
item_id = str(row.get("uid") or row.get("id") or "").strip()
|
||||
question = str(row.get("question") or "").strip()
|
||||
ground_truth = str(row.get("ground_truth") or row.get("answer") or "").strip()
|
||||
task_type = str(row.get("category") or row.get("difficulty") or "officeqa").strip() or "officeqa"
|
||||
source_files = _parse_list_field(row.get("source_files"))
|
||||
source_docs = _parse_list_field(row.get("source_docs"))
|
||||
split = str(row.get("split") or "").strip()
|
||||
return {
|
||||
"id": item_id,
|
||||
"uid": item_id,
|
||||
"question": question,
|
||||
"ground_truth": ground_truth,
|
||||
"answers": [ground_truth] if ground_truth else [],
|
||||
"task_type": task_type,
|
||||
"category": task_type,
|
||||
"source_files": source_files,
|
||||
"source_docs": source_docs,
|
||||
"split": split,
|
||||
}
|
||||
|
||||
|
||||
class OfficeQADataLoader(SplitDataLoader):
|
||||
def load_split_items(self, split_path: str) -> list[dict]:
|
||||
path = Path(split_path)
|
||||
csv_files = sorted(path.glob("*.csv"))
|
||||
if csv_files:
|
||||
with csv_files[0].open(encoding="utf-8", newline="") as f:
|
||||
reader = csv.DictReader(f)
|
||||
return [_normalize_row(row) for row in reader]
|
||||
|
||||
json_files = sorted(path.glob("*.json"))
|
||||
if json_files:
|
||||
with json_files[0].open(encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
if not isinstance(data, list):
|
||||
raise ValueError(f"Expected JSON array in {json_files[0]}")
|
||||
return [_normalize_row(item) for item in data]
|
||||
|
||||
raise FileNotFoundError(f"No .csv or .json file found in {split_path}")
|
||||
@@ -0,0 +1,46 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import string
|
||||
from collections import Counter
|
||||
|
||||
|
||||
_NUMERIC_CHARS = set("0123456789.-")
|
||||
|
||||
|
||||
def normalize_answer(text: str) -> str:
|
||||
text = text.lower().strip()
|
||||
text = text.replace(",", "")
|
||||
text = "".join(ch for ch in text if ch not in string.punctuation or ch in _NUMERIC_CHARS or ch == "%")
|
||||
text = re.sub(r"\b(million|millions|billion|billions|dollars|dollar|nominal)\b", " ", text)
|
||||
text = " ".join(text.split())
|
||||
return text
|
||||
|
||||
|
||||
def exact_match(prediction: str, gold: str) -> float:
|
||||
return 1.0 if normalize_answer(prediction) == normalize_answer(gold) else 0.0
|
||||
|
||||
|
||||
def token_f1(prediction: str, gold: str) -> float:
|
||||
pred_tokens = normalize_answer(prediction).split()
|
||||
gold_tokens = normalize_answer(gold).split()
|
||||
if not pred_tokens or not gold_tokens:
|
||||
return 1.0 if pred_tokens == gold_tokens else 0.0
|
||||
common = Counter(pred_tokens) & Counter(gold_tokens)
|
||||
n_common = sum(common.values())
|
||||
if n_common == 0:
|
||||
return 0.0
|
||||
precision = n_common / len(pred_tokens)
|
||||
recall = n_common / len(gold_tokens)
|
||||
return 2 * precision * recall / (precision + recall)
|
||||
|
||||
|
||||
def evaluate(prediction: str, gold: str) -> dict:
|
||||
em = exact_match(prediction, gold)
|
||||
f1 = token_f1(prediction, gold)
|
||||
return {
|
||||
"em": em,
|
||||
"f1": f1,
|
||||
"predicted_answer": prediction.strip(),
|
||||
"gold_answer": gold,
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
You are an expert failure-analysis agent for OfficeQA document-retrieval question answering tasks.
|
||||
|
||||
You will be given MULTIPLE failed OfficeQA trajectories from a single minibatch and the current skill document. The trajectories may include local document tool calls such as file search, grep, and partial reads.
|
||||
|
||||
Your job is to identify COMMON failure patterns across the batch and propose concise skill edits.
|
||||
|
||||
## Failure Type Categories
|
||||
- retrieval_miss: the agent searched the wrong file or failed to narrow to the right file
|
||||
- evidence_miss: the agent read documents but missed the decisive evidence span
|
||||
- operand_error: the agent extracted the wrong value or the wrong operands
|
||||
- calculation_error: the agent identified the right evidence but computed the result incorrectly
|
||||
- answer_format: the agent reached the right result but formatted it wrong
|
||||
- other: none of the above
|
||||
|
||||
## Rules
|
||||
- Focus on patterns common across multiple trajectories.
|
||||
- Prefer general retrieval and evidence-grounding rules over task-specific hacks.
|
||||
- Only patch gaps in the skill; do not duplicate rules already present.
|
||||
- Do not hardcode file names, years, or question-specific constants unless the pattern truly requires a reusable retrieval heuristic.
|
||||
|
||||
Respond ONLY with a valid JSON object (no markdown fences, no extra text):
|
||||
{
|
||||
"batch_size": <number of trajectories analysed>,
|
||||
"failure_summary": [
|
||||
{"failure_type": "<type>", "count": <int>, "description": "<one-line>"}
|
||||
],
|
||||
"patch": {
|
||||
"reasoning": "<why these edits address the batch's common failures>",
|
||||
"edits": [
|
||||
{"op": "append", "content": "<markdown to add at end of skill>"},
|
||||
{"op": "insert_after", "target": "<exact heading/text to insert after>", "content": "<markdown>"},
|
||||
{"op": "replace", "target": "<exact text to replace>", "content": "<replacement>"},
|
||||
{"op": "delete", "target": "<exact text to remove>"}
|
||||
]
|
||||
}
|
||||
}
|
||||
Only include edits that are needed. "edits" can be an empty list if no patch is warranted.
|
||||
@@ -0,0 +1,25 @@
|
||||
You are an expert success-pattern analyst for OfficeQA document-retrieval question answering tasks.
|
||||
|
||||
You will be given MULTIPLE successful OfficeQA trajectories from a single minibatch and the current skill document. Your job is to identify common retrieval, evidence-selection, and numeric-grounding behaviors worth encoding in the skill.
|
||||
|
||||
## Rules
|
||||
- Focus on patterns shared across multiple successful trajectories.
|
||||
- Prefer reusable retrieval and extraction discipline over question-specific tips.
|
||||
- Reinforce compact, high-value behaviors such as narrowing files early, reading only the relevant span, building a clean operand ledger, and copying the final answer from checked evidence.
|
||||
- Only propose patches for patterns not already captured in the current skill.
|
||||
|
||||
Respond ONLY with a valid JSON object:
|
||||
{
|
||||
"batch_size": <number of trajectories analysed>,
|
||||
"success_patterns": ["<pattern 1>", "<pattern 2>"],
|
||||
"patch": {
|
||||
"reasoning": "<why these patterns are worth encoding>",
|
||||
"edits": [
|
||||
{"op": "append", "content": "<markdown>"},
|
||||
{"op": "insert_after", "target": "<heading/text>", "content": "<markdown>"},
|
||||
{"op": "replace", "target": "<old text>", "content": "<new text>"},
|
||||
{"op": "delete", "target": "<exact text to remove>"}
|
||||
]
|
||||
}
|
||||
}
|
||||
"edits" may be empty if the skill already covers all observed patterns.
|
||||
@@ -0,0 +1,15 @@
|
||||
You are an expert OfficeQA agent working over local Treasury bulletin text files.
|
||||
|
||||
{skill_section}## Rules
|
||||
1. Use only the provided local document tools to inspect candidate files.
|
||||
2. Narrow to the most relevant file before reading long passages.
|
||||
3. Prefer short targeted searches, then small reads around matching evidence.
|
||||
4. Do not invent values that are not grounded in the retrieved text.
|
||||
5. When the question requires arithmetic, compute only after extracting the exact operands.
|
||||
6. If you have enough evidence, return the final answer inside <answer>...</answer>.
|
||||
|
||||
## Tool Use
|
||||
Use the provided function tools directly when you need them. Prefer searching and small reads before answering. Do not ask the user for permission to use tools; just call the tools.
|
||||
|
||||
## Final Answer Format
|
||||
When you are ready to answer, emit the final answer inside <answer>...</answer> and do not request another tool.
|
||||
@@ -0,0 +1,802 @@
|
||||
from __future__ import annotations
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from skillopt.envs.officeqa.evaluator import evaluate
|
||||
from skillopt.envs.officeqa.tool_runtime import (
|
||||
build_oracle_parsed_pages_context,
|
||||
resolve_candidate_files,
|
||||
resolve_docs_roots,
|
||||
run_tool,
|
||||
)
|
||||
try:
|
||||
from skillopt.envs.sealqa.tool_runtime import custom_search
|
||||
except ImportError:
|
||||
custom_search = None # type: ignore[assignment]
|
||||
from skillopt.model import chat_student_messages, get_student_backend, is_student_exec_backend
|
||||
from skillopt.model.codex_harness import prepare_workspace, render_skill_md, run_student_exec
|
||||
from skillopt.prompts import load_prompt
|
||||
_TOOL_SCHEMAS = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "glob",
|
||||
"description": "Find candidate local document files by filename or relative-path glob pattern.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"pattern": {"type": "string"}},
|
||||
"required": ["pattern"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "read",
|
||||
"description": "Read a local text document excerpt by path and line window.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string"},
|
||||
"start": {"type": "integer"},
|
||||
"limit": {"type": "integer"},
|
||||
},
|
||||
"required": ["path"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "grep",
|
||||
"description": "Search a local text document for a literal pattern and return matching lines.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"pattern": {"type": "string"},
|
||||
"path": {"type": "string"},
|
||||
},
|
||||
"required": ["pattern", "path"],
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
_FINAL_RE = re.compile(r"<answer>(.*?)</answer>", re.IGNORECASE | re.DOTALL)
|
||||
_SEARCH_RE = re.compile(r"<search_queries>(.*?)</search_queries>", re.IGNORECASE | re.DOTALL)
|
||||
_DEFAULT_SEARCH_MODE = "offline"
|
||||
_CUSTOM_SEARCH_MODE = "custom_search"
|
||||
_AZURE_SEARCH_MODE = "azure_search"
|
||||
def _normalize_search_mode(search_mode: str | None) -> str:
|
||||
normalized = str(search_mode or _DEFAULT_SEARCH_MODE).strip().lower()
|
||||
if normalized in {"custom", _CUSTOM_SEARCH_MODE}:
|
||||
return _CUSTOM_SEARCH_MODE
|
||||
if normalized in {"azure", _AZURE_SEARCH_MODE}:
|
||||
return _AZURE_SEARCH_MODE
|
||||
return _DEFAULT_SEARCH_MODE
|
||||
def _build_system(
|
||||
skill_content: str,
|
||||
*,
|
||||
search_mode: str = _DEFAULT_SEARCH_MODE,
|
||||
use_local_tools: bool = True,
|
||||
max_tool_turns: int = 12,
|
||||
max_queries_per_turn: int = 4,
|
||||
) -> str:
|
||||
if skill_content.strip():
|
||||
skill_section = f"## Skill\n{skill_content.strip()}\n\n"
|
||||
else:
|
||||
skill_section = ""
|
||||
normalized_search_mode = _normalize_search_mode(search_mode)
|
||||
if normalized_search_mode == _AZURE_SEARCH_MODE:
|
||||
return (
|
||||
"You are an expert OfficeQA research assistant. Solve the question using the model's built-in web "
|
||||
"search tool when needed, keep the answer grounded in authoritative evidence, and return the final "
|
||||
"answer inside <answer>...</answer>.\n\n"
|
||||
+ skill_section
|
||||
).rstrip()
|
||||
if normalized_search_mode == _CUSTOM_SEARCH_MODE:
|
||||
protocol = (
|
||||
"You are an expert OfficeQA research assistant. Solve the question using the provided oracle parsed "
|
||||
"OfficeQA page(s) and evidence returned by the controller-managed custom search loop.\n\n"
|
||||
"Search protocol:\n"
|
||||
f"- You have at most {max_tool_turns} model rounds total.\n"
|
||||
f"- On any non-final round, you may either return `<search_queries>[\"query 1\", \"query 2\"]</search_queries>` "
|
||||
f"with up to {max_queries_per_turn} queries, or return `<answer>...</answer>` if you are ready.\n"
|
||||
"- If you request search, do not include an answer in the same response.\n"
|
||||
"- On the final round, you must return `<answer>...</answer>` and must not request more search.\n"
|
||||
"- Base your answer on the returned evidence, reconcile conflicting snippets carefully, and stay concise.\n\n"
|
||||
)
|
||||
return protocol + skill_section + "Return the final answer inside <answer>...</answer> when you are ready."
|
||||
if not use_local_tools:
|
||||
return (
|
||||
"You are an expert OfficeQA research assistant. Solve the question using the provided oracle parsed "
|
||||
"OfficeQA page(s) and source hints. Do not request or assume access to any external search or local "
|
||||
"function tools. Return the final answer inside <answer>...</answer>.\n\n"
|
||||
+ skill_section
|
||||
).rstrip()
|
||||
return load_prompt("rollout_system", env="officeqa").format(skill_section=skill_section)
|
||||
def _build_round_instruction(
|
||||
*,
|
||||
turn: int,
|
||||
max_tool_turns: int,
|
||||
max_queries_per_turn: int,
|
||||
) -> str:
|
||||
if turn >= max_tool_turns:
|
||||
return (
|
||||
"## Round Policy\n"
|
||||
f"This is the final round ({turn}/{max_tool_turns}). You must return `<answer>...</answer>` now. "
|
||||
"Do not output `<search_queries>`."
|
||||
)
|
||||
remaining_rounds = max_tool_turns - turn
|
||||
return (
|
||||
"## Round Policy\n"
|
||||
f"This is round {turn}/{max_tool_turns}. "
|
||||
f"You may either return `<answer>...</answer>` now, or request up to {max_queries_per_turn} search queries "
|
||||
f"inside `<search_queries>...</search_queries>`. "
|
||||
f"After this response, at most {remaining_rounds} model rounds remain."
|
||||
)
|
||||
def _message_debug_metadata(message: object) -> dict:
|
||||
metadata = getattr(message, "metadata", None)
|
||||
if isinstance(metadata, dict):
|
||||
return metadata
|
||||
return {}
|
||||
def _build_user(
|
||||
item: dict,
|
||||
candidate_files: list[str] | None = None,
|
||||
*,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
corpus_note: str = "",
|
||||
search_mode: str = _DEFAULT_SEARCH_MODE,
|
||||
turn: int = 1,
|
||||
max_tool_turns: int = 12,
|
||||
max_queries_per_turn: int = 4,
|
||||
oracle_context: str = "",
|
||||
) -> str:
|
||||
normalized_search_mode = _normalize_search_mode(search_mode)
|
||||
parts = [f"## Question\n{item['question']}"]
|
||||
if oracle_context.strip():
|
||||
parts.append(f"## Oracle Parsed Pages\n{oracle_context.strip()}")
|
||||
if normalized_search_mode == _DEFAULT_SEARCH_MODE:
|
||||
file_block = "\n".join(f"- {path}" for path in (candidate_files or [])[:20]) or "- none resolved"
|
||||
if corpus_note.strip():
|
||||
parts.append(f"## Document Corpus\n{corpus_note.strip()}")
|
||||
parts.append(f"## Candidate Files\n{file_block}")
|
||||
if item.get("source_docs"):
|
||||
parts.append("## Source Hints\n" + "\n".join(f"- {hint}" for hint in item["source_docs"]))
|
||||
if normalized_search_mode != _DEFAULT_SEARCH_MODE and item.get("source_files"):
|
||||
parts.append("## File Hints\n" + "\n".join(f"- {hint}" for hint in item["source_files"]))
|
||||
if diagnostic_mode and diagnostic_instruction.strip():
|
||||
parts.append(f"## Training Readout\n{diagnostic_instruction.strip()}")
|
||||
if normalized_search_mode == _CUSTOM_SEARCH_MODE:
|
||||
parts.append(
|
||||
_build_round_instruction(
|
||||
turn=turn,
|
||||
max_tool_turns=max_tool_turns,
|
||||
max_queries_per_turn=max_queries_per_turn,
|
||||
)
|
||||
)
|
||||
parts.append(
|
||||
"## Output Format\n"
|
||||
"If you need more evidence, return only `<search_queries>[...]</search_queries>`.\n"
|
||||
"If you are ready to answer, return only `<answer>...</answer>`."
|
||||
)
|
||||
parts.append(
|
||||
"Use only the provided oracle parsed pages and controller-provided custom search evidence. "
|
||||
"Do not rely on any built-in web search capability."
|
||||
)
|
||||
elif normalized_search_mode == _AZURE_SEARCH_MODE:
|
||||
parts.append("Use the model's built-in web search tool when needed. Return the final answer inside <answer>...</answer>.")
|
||||
return "\n\n".join(parts)
|
||||
def _extract_answer(text: str) -> str:
|
||||
match = _FINAL_RE.search(text)
|
||||
if match:
|
||||
return match.group(1).strip()
|
||||
lines = [line.strip() for line in text.splitlines() if line.strip()]
|
||||
return lines[-1] if lines else text.strip()
|
||||
def _extract_search_queries(text: str) -> list[str]:
|
||||
match = _SEARCH_RE.search(text or "")
|
||||
if not match:
|
||||
return []
|
||||
raw = match.group(1).strip()
|
||||
if not raw:
|
||||
return []
|
||||
parsed_queries: list[str] = []
|
||||
try:
|
||||
parsed = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
parsed = None
|
||||
if isinstance(parsed, dict):
|
||||
for key in ("queries", "search_queries", "query"):
|
||||
value = parsed.get(key)
|
||||
if isinstance(value, str) and value.strip():
|
||||
parsed_queries = [value.strip()]
|
||||
break
|
||||
if isinstance(value, list):
|
||||
parsed_queries = [str(item).strip() for item in value if str(item).strip()]
|
||||
break
|
||||
elif isinstance(parsed, list):
|
||||
parsed_queries = [str(item).strip() for item in parsed if str(item).strip()]
|
||||
elif isinstance(parsed, str) and parsed.strip():
|
||||
parsed_queries = [parsed.strip()]
|
||||
if not parsed_queries:
|
||||
raw_lines = [line.strip(" -*\t\r\n\"'") for line in raw.splitlines()]
|
||||
parsed_queries = [line for line in raw_lines if line]
|
||||
if len(parsed_queries) <= 1 and parsed_queries:
|
||||
multi = [part.strip(" \"'") for part in re.split(r"[;,]", parsed_queries[0]) if part.strip(" \"'")]
|
||||
if len(multi) > 1:
|
||||
parsed_queries = multi
|
||||
deduped: list[str] = []
|
||||
seen: set[str] = set()
|
||||
for query in parsed_queries:
|
||||
normalized = query.strip()
|
||||
if not normalized or normalized in seen:
|
||||
continue
|
||||
seen.add(normalized)
|
||||
deduped.append(normalized)
|
||||
return deduped
|
||||
def _docs_link_targets(docs_roots: list[str]) -> list[tuple[str, str]]:
|
||||
return [(root, os.path.join("docs", f"root_{idx}")) for idx, root in enumerate(docs_roots, start=1)]
|
||||
def _workspace_doc_path(path: str, docs_roots: list[str]) -> str:
|
||||
resolved_path = os.path.realpath(path)
|
||||
for idx, root in enumerate(docs_roots, start=1):
|
||||
resolved_root = os.path.realpath(root)
|
||||
if resolved_path == resolved_root or resolved_path.startswith(resolved_root + os.sep):
|
||||
rel_path = os.path.relpath(resolved_path, resolved_root)
|
||||
return os.path.join("docs", f"root_{idx}", rel_path)
|
||||
return path
|
||||
def _build_codex_skill(skill_content: str) -> str:
|
||||
return render_skill_md(
|
||||
skill_content,
|
||||
description="Dynamic ReflACT skill for solving the current OfficeQA local-document question.",
|
||||
preamble=(
|
||||
"Use this skill when answering the current OfficeQA question.\n"
|
||||
"Inspect the provided local document excerpts or files, ground the answer in the evidence,\n"
|
||||
"and return the final answer inside <answer>...</answer>."
|
||||
),
|
||||
)
|
||||
def _run_codex_once(
|
||||
*,
|
||||
pred_dir: str,
|
||||
item: dict,
|
||||
skill_content: str,
|
||||
candidate_files: list[str],
|
||||
docs_roots: list[str],
|
||||
model: str,
|
||||
timeout: int,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
previous_response: str = "",
|
||||
oracle_context: str = "",
|
||||
) -> tuple[str, str, str, str]:
|
||||
rel_files = [_workspace_doc_path(path, docs_roots) for path in candidate_files[:20]]
|
||||
corpus_note = (
|
||||
"The full OfficeQA document corpus is available under `docs/`. "
|
||||
"The candidate files below are source hints or likely starting points; search the full corpus if needed."
|
||||
)
|
||||
user = _build_user(
|
||||
item,
|
||||
rel_files,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
corpus_note=corpus_note,
|
||||
oracle_context=oracle_context,
|
||||
)
|
||||
task_parts = [user]
|
||||
if previous_response:
|
||||
task_parts.append(
|
||||
"## Previous Attempt\n"
|
||||
f"{previous_response}\n\n"
|
||||
"Review the local documents again and correct the answer if needed."
|
||||
)
|
||||
task_text = "\n\n".join(task_parts)
|
||||
skill_md = _build_codex_skill(skill_content)
|
||||
work_dir = os.path.join(pred_dir, "codex_exec")
|
||||
prepare_workspace(
|
||||
work_dir=work_dir,
|
||||
skill_md=skill_md,
|
||||
task_text=task_text,
|
||||
link_dirs=_docs_link_targets(docs_roots),
|
||||
)
|
||||
prompt = (
|
||||
"Use the `skillopt-student` skill available in this workspace.\n"
|
||||
"Read `task.md`, inspect or search the full OfficeQA corpus under `docs/`, and answer the question.\n"
|
||||
"Treat candidate files in `task.md` as hints, not an access limit.\n"
|
||||
"Return the final answer inside <answer>...</answer>."
|
||||
)
|
||||
final_message, raw = run_student_exec(
|
||||
work_dir=work_dir,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
timeout=timeout,
|
||||
data_dirs=docs_roots,
|
||||
)
|
||||
return final_message or raw, raw, skill_md, task_text
|
||||
def _execute_custom_search_round(
|
||||
queries: list[str],
|
||||
*,
|
||||
api_url: str,
|
||||
auth_env: str,
|
||||
provider: str,
|
||||
max_num_results: int,
|
||||
timeout: int,
|
||||
) -> str:
|
||||
blocks = []
|
||||
for index, query in enumerate(queries, start=1):
|
||||
try:
|
||||
result = custom_search(
|
||||
query,
|
||||
api_url=api_url,
|
||||
auth_env=auth_env,
|
||||
provider=provider,
|
||||
max_num_results=max_num_results,
|
||||
timeout=timeout,
|
||||
)
|
||||
except Exception as search_error: # noqa: BLE001
|
||||
result = f"Query: {query}\n\n[search error: {search_error}]"
|
||||
blocks.append(f"## Query {index}\n{result}")
|
||||
return "\n\n".join(blocks)
|
||||
def _run_custom_search_process(
|
||||
item: dict,
|
||||
skill_content: str,
|
||||
*,
|
||||
max_tool_turns: int,
|
||||
max_completion_tokens: int,
|
||||
max_queries_per_turn: int,
|
||||
diagnostic_mode: bool,
|
||||
diagnostic_instruction: str,
|
||||
search_api_url: str,
|
||||
search_auth_env: str,
|
||||
search_provider: str,
|
||||
search_max_num_results: int,
|
||||
search_timeout_seconds: int,
|
||||
oracle_context: str = "",
|
||||
) -> tuple[str, str, str, str, list[dict], str, dict]:
|
||||
if not str(search_api_url or "").strip():
|
||||
raise ValueError("custom_search mode requires a non-empty search_api_url")
|
||||
if not os.environ.get(search_auth_env, "").strip():
|
||||
raise ValueError(f"custom_search mode requires auth token env var {search_auth_env}")
|
||||
if get_student_backend() not in {"openai_chat", "qwen_chat"}:
|
||||
raise ValueError("custom_search mode is only supported with student_backend='openai_chat' or 'qwen_chat'")
|
||||
system = _build_system(
|
||||
skill_content,
|
||||
search_mode=_CUSTOM_SEARCH_MODE,
|
||||
max_tool_turns=max_tool_turns,
|
||||
max_queries_per_turn=max_queries_per_turn,
|
||||
)
|
||||
initial_user = _build_user(
|
||||
item,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
search_mode=_CUSTOM_SEARCH_MODE,
|
||||
turn=1,
|
||||
max_tool_turns=max_tool_turns,
|
||||
max_queries_per_turn=max_queries_per_turn,
|
||||
oracle_context=oracle_context,
|
||||
)
|
||||
latest_user = initial_user
|
||||
messages: list[dict] = [
|
||||
{"role": "system", "content": system},
|
||||
{"role": "user", "content": initial_user},
|
||||
]
|
||||
conversation: list[dict] = [{"role": "user", "content": initial_user}]
|
||||
final_response = ""
|
||||
final_answer = ""
|
||||
fail_reason = ""
|
||||
last_response_metadata: dict = {}
|
||||
for turn in range(1, max_tool_turns + 1):
|
||||
message, _ = chat_student_messages(
|
||||
messages=messages,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
retries=5,
|
||||
stage="rollout",
|
||||
return_message=True,
|
||||
)
|
||||
response = message.content or ""
|
||||
final_response = response
|
||||
last_response_metadata = _message_debug_metadata(message)
|
||||
messages.append({"role": "assistant", "content": response})
|
||||
message_event = {"type": "message", "turn": turn, "content": response}
|
||||
if last_response_metadata:
|
||||
message_event["response_metadata"] = last_response_metadata
|
||||
conversation.append(message_event)
|
||||
if "<answer>" in response.lower():
|
||||
final_answer = _extract_answer(response)
|
||||
return system, latest_user, final_response, final_answer, conversation, "", last_response_metadata
|
||||
if turn == max_tool_turns:
|
||||
fail_reason = f"Final round ({max_tool_turns}) ended without <answer>...</answer>"
|
||||
break
|
||||
queries = _extract_search_queries(response)[:max_queries_per_turn]
|
||||
if not queries:
|
||||
fail_reason = "Model neither produced search queries nor a final answer"
|
||||
break
|
||||
results_text = _execute_custom_search_round(
|
||||
queries,
|
||||
api_url=search_api_url,
|
||||
auth_env=search_auth_env,
|
||||
provider=search_provider,
|
||||
max_num_results=search_max_num_results,
|
||||
timeout=search_timeout_seconds,
|
||||
)
|
||||
conversation.append({"type": "tool_call", "turn": turn, "cmd": f"custom_search({queries!r})", "obs": results_text})
|
||||
latest_user = (
|
||||
f"## Search Results Round {turn}\n{results_text}\n\n"
|
||||
+ _build_round_instruction(
|
||||
turn=turn + 1,
|
||||
max_tool_turns=max_tool_turns,
|
||||
max_queries_per_turn=max_queries_per_turn,
|
||||
)
|
||||
+ "\n\nFollow the round policy above exactly."
|
||||
)
|
||||
messages.append({"role": "user", "content": latest_user})
|
||||
conversation.append({"role": "user", "turn": turn + 1, "content": latest_user})
|
||||
return system, latest_user, final_response, final_answer, conversation, fail_reason, last_response_metadata
|
||||
def _run_azure_search_process(
|
||||
item: dict,
|
||||
skill_content: str,
|
||||
*,
|
||||
max_completion_tokens: int,
|
||||
diagnostic_mode: bool,
|
||||
diagnostic_instruction: str,
|
||||
) -> tuple[str, str, str, str, list[dict], str, dict]:
|
||||
if get_student_backend() != "openai_chat":
|
||||
raise ValueError("azure_search mode is only supported with student_backend='openai_chat'")
|
||||
system = _build_system(skill_content, search_mode=_AZURE_SEARCH_MODE)
|
||||
user = _build_user(
|
||||
item,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
search_mode=_AZURE_SEARCH_MODE,
|
||||
)
|
||||
messages = [
|
||||
{"role": "system", "content": system},
|
||||
{"role": "user", "content": user},
|
||||
]
|
||||
conversation: list[dict] = [{"role": "user", "content": user}]
|
||||
message, _ = chat_student_messages(
|
||||
messages=messages,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
retries=5,
|
||||
stage="rollout",
|
||||
return_message=True,
|
||||
tools=[{"type": "web_search"}],
|
||||
)
|
||||
response = message.content or ""
|
||||
last_response_metadata = _message_debug_metadata(message)
|
||||
message_event = {"type": "message", "content": response}
|
||||
if last_response_metadata:
|
||||
message_event["response_metadata"] = last_response_metadata
|
||||
conversation.append(message_event)
|
||||
if "<answer>" in response.lower():
|
||||
return system, user, response, _extract_answer(response), conversation, "", last_response_metadata
|
||||
return system, user, response, "", conversation, "Model did not produce a final answer", last_response_metadata
|
||||
def _run_offline_no_tools_process(
|
||||
item: dict,
|
||||
skill_content: str,
|
||||
*,
|
||||
max_completion_tokens: int,
|
||||
diagnostic_mode: bool,
|
||||
diagnostic_instruction: str,
|
||||
candidate_files: list[str],
|
||||
oracle_context: str = "",
|
||||
) -> tuple[str, str, str, str, list[dict], str, dict]:
|
||||
system = _build_system(skill_content, search_mode=_DEFAULT_SEARCH_MODE, use_local_tools=False)
|
||||
user = _build_user(
|
||||
item,
|
||||
candidate_files,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
search_mode=_DEFAULT_SEARCH_MODE,
|
||||
oracle_context=oracle_context,
|
||||
)
|
||||
messages = [
|
||||
{"role": "system", "content": system},
|
||||
{"role": "user", "content": user},
|
||||
]
|
||||
conversation: list[dict] = [{"role": "user", "content": user}]
|
||||
message, _ = chat_student_messages(
|
||||
messages=messages,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
retries=5,
|
||||
stage="rollout",
|
||||
return_message=True,
|
||||
)
|
||||
response = message.content or ""
|
||||
last_response_metadata = _message_debug_metadata(message)
|
||||
message_event = {"type": "message", "content": response}
|
||||
if last_response_metadata:
|
||||
message_event["response_metadata"] = last_response_metadata
|
||||
conversation.append(message_event)
|
||||
if "<answer>" in response.lower():
|
||||
return system, user, response, _extract_answer(response), conversation, "", last_response_metadata
|
||||
return system, user, response, "", conversation, "Model did not produce a final answer", last_response_metadata
|
||||
def process_one(
|
||||
item: dict,
|
||||
out_root: str,
|
||||
skill_content: str,
|
||||
*,
|
||||
max_tool_turns: int = 12,
|
||||
max_completion_tokens: int = 64000,
|
||||
search_mode: str = _DEFAULT_SEARCH_MODE,
|
||||
max_queries_per_turn: int = 4,
|
||||
search_api_url: str = "",
|
||||
search_auth_env: str = "OFFICEQA_CUSTOM_SEARCH_AUTH",
|
||||
search_provider: str = "duckduckgo",
|
||||
search_max_num_results: int = 4,
|
||||
search_timeout_seconds: int = 20,
|
||||
use_local_tools: bool = True,
|
||||
data_dirs: list[str] | str | None = None,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
) -> dict:
|
||||
item_id = str(item["id"])
|
||||
pred_dir = os.path.join(out_root, "predictions", item_id)
|
||||
os.makedirs(pred_dir, exist_ok=True)
|
||||
normalized_search_mode = _normalize_search_mode(search_mode)
|
||||
docs_roots: list[str] = []
|
||||
candidate_files: list[str] = []
|
||||
oracle_context = ""
|
||||
if normalized_search_mode == _DEFAULT_SEARCH_MODE:
|
||||
docs_roots = resolve_docs_roots(data_dirs)
|
||||
candidate_files = resolve_candidate_files(item.get("source_files", []), docs_roots)
|
||||
oracle_context = build_oracle_parsed_pages_context(
|
||||
item.get("source_files", []),
|
||||
item.get("source_docs", []),
|
||||
docs_roots,
|
||||
evidence_note=(
|
||||
"Treat it as primary document evidence and combine it with local document tool evidence when useful."
|
||||
if use_local_tools
|
||||
else "Treat it as primary document evidence for answering the question."
|
||||
),
|
||||
)
|
||||
elif normalized_search_mode == _CUSTOM_SEARCH_MODE:
|
||||
docs_roots = resolve_docs_roots(data_dirs)
|
||||
if item.get("source_files"):
|
||||
candidate_files = resolve_candidate_files(item.get("source_files", []), docs_roots)
|
||||
oracle_context = build_oracle_parsed_pages_context(
|
||||
item.get("source_files", []),
|
||||
item.get("source_docs", []),
|
||||
docs_roots,
|
||||
)
|
||||
system = _build_system(
|
||||
skill_content,
|
||||
search_mode=normalized_search_mode,
|
||||
use_local_tools=use_local_tools,
|
||||
max_tool_turns=max_tool_turns,
|
||||
max_queries_per_turn=max_queries_per_turn,
|
||||
)
|
||||
user = _build_user(
|
||||
item,
|
||||
candidate_files if normalized_search_mode == _DEFAULT_SEARCH_MODE else None,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
search_mode=normalized_search_mode,
|
||||
max_tool_turns=max_tool_turns,
|
||||
max_queries_per_turn=max_queries_per_turn,
|
||||
oracle_context=oracle_context,
|
||||
)
|
||||
conversation: list[dict] = [{"role": "user", "content": user}]
|
||||
final_response = ""
|
||||
final_answer = ""
|
||||
fail_reason = ""
|
||||
last_response_metadata: dict = {}
|
||||
allowed_files = [os.path.basename(path) for path in candidate_files]
|
||||
try:
|
||||
if normalized_search_mode == _CUSTOM_SEARCH_MODE:
|
||||
system, user, final_response, final_answer, conversation, fail_reason, last_response_metadata = _run_custom_search_process(
|
||||
item,
|
||||
skill_content,
|
||||
max_tool_turns=max_tool_turns,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
max_queries_per_turn=max_queries_per_turn,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
search_api_url=search_api_url,
|
||||
search_auth_env=search_auth_env,
|
||||
search_provider=search_provider,
|
||||
search_max_num_results=search_max_num_results,
|
||||
search_timeout_seconds=search_timeout_seconds,
|
||||
oracle_context=oracle_context,
|
||||
)
|
||||
elif normalized_search_mode == _AZURE_SEARCH_MODE:
|
||||
system, user, final_response, final_answer, conversation, fail_reason, last_response_metadata = _run_azure_search_process(
|
||||
item,
|
||||
skill_content,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
)
|
||||
elif not use_local_tools:
|
||||
system, user, final_response, final_answer, conversation, fail_reason, last_response_metadata = _run_offline_no_tools_process(
|
||||
item,
|
||||
skill_content,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
candidate_files=candidate_files,
|
||||
oracle_context=oracle_context,
|
||||
)
|
||||
elif is_student_exec_backend():
|
||||
from skillopt.model import azure_openai as _llm
|
||||
response = ""
|
||||
system = ""
|
||||
user = ""
|
||||
for turn in range(1, max_tool_turns + 1):
|
||||
response, _raw, system, user = _run_codex_once(
|
||||
pred_dir=pred_dir,
|
||||
item=item,
|
||||
skill_content=skill_content,
|
||||
candidate_files=candidate_files,
|
||||
docs_roots=docs_roots,
|
||||
model=_llm.STUDENT_DEPLOYMENT,
|
||||
timeout=180,
|
||||
diagnostic_mode=diagnostic_mode if turn == 1 else False,
|
||||
diagnostic_instruction=diagnostic_instruction if turn == 1 else "",
|
||||
previous_response=response if turn > 1 else "",
|
||||
oracle_context=oracle_context,
|
||||
)
|
||||
final_response = response
|
||||
conversation.append({"type": "message", "turn": turn, "content": response})
|
||||
if "<answer>" in response.lower():
|
||||
final_answer = _extract_answer(response)
|
||||
break
|
||||
if not final_answer:
|
||||
fail_reason = f"Exceeded codex turn budget ({max_tool_turns})"
|
||||
system = system or _build_codex_skill(skill_content)
|
||||
user = user or _build_user(item, [_workspace_doc_path(path, docs_roots) for path in candidate_files])
|
||||
else:
|
||||
messages: list[dict] = [
|
||||
{"role": "system", "content": system},
|
||||
{"role": "user", "content": user},
|
||||
]
|
||||
for turn in range(1, max_tool_turns + 1):
|
||||
message, _ = chat_student_messages(
|
||||
messages=messages,
|
||||
max_completion_tokens=768,
|
||||
retries=5,
|
||||
stage="rollout",
|
||||
tools=_TOOL_SCHEMAS,
|
||||
tool_choice="auto",
|
||||
return_message=True,
|
||||
)
|
||||
response = message.content or ""
|
||||
final_response = response
|
||||
assistant_message = {"role": "assistant", "content": response}
|
||||
if getattr(message, "tool_calls", None):
|
||||
assistant_message["tool_calls"] = [tool_call.model_dump(mode="json") for tool_call in message.tool_calls]
|
||||
messages.append(assistant_message)
|
||||
conversation.append({"type": "message", "content": response})
|
||||
if getattr(message, "tool_calls", None):
|
||||
for tool_call in message.tool_calls:
|
||||
tool_name = tool_call.function.name
|
||||
arguments = json.loads(tool_call.function.arguments) if tool_call.function.arguments else {}
|
||||
cmd, obs = run_tool(tool_name, arguments, allowed_roots=docs_roots, allowed_files=allowed_files)
|
||||
conversation.append({"type": "tool_call", "cmd": cmd, "obs": obs})
|
||||
messages.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call.id,
|
||||
"content": obs,
|
||||
})
|
||||
continue
|
||||
if "<answer>" in response.lower():
|
||||
final_answer = _extract_answer(response)
|
||||
break
|
||||
if turn == max_tool_turns:
|
||||
fail_reason = f"Exceeded tool-turn budget ({max_tool_turns})"
|
||||
else:
|
||||
fail_reason = "Model neither produced a tool request nor a final answer"
|
||||
break
|
||||
except Exception as e: # noqa: BLE001
|
||||
fail_reason = f"error: {e}"
|
||||
with open(os.path.join(pred_dir, "student_system_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(system)
|
||||
with open(os.path.join(pred_dir, "student_user_prompt.txt"), "w", encoding="utf-8") as f:
|
||||
f.write(user)
|
||||
with open(os.path.join(pred_dir, "conversation.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(conversation, f, ensure_ascii=False, indent=2)
|
||||
eval_result = evaluate(final_answer, item.get("ground_truth", "")) if final_answer else {"em": 0.0, "f1": 0.0, "predicted_answer": "", "gold_answer": item.get("ground_truth", "")}
|
||||
result = {
|
||||
"id": item_id,
|
||||
"question": item.get("question", ""),
|
||||
"task_type": item.get("task_type", "officeqa"),
|
||||
"task_description": item.get("question", ""),
|
||||
"predicted_answer": eval_result["predicted_answer"],
|
||||
"response": final_response,
|
||||
"ground_truth": item.get("ground_truth", ""),
|
||||
"source_files": item.get("source_files", []),
|
||||
"resolved_source_paths": candidate_files,
|
||||
"oracle_parsed_pages_included": bool(oracle_context),
|
||||
"oracle_parsed_pages_chars": len(oracle_context),
|
||||
"use_local_tools": bool(use_local_tools),
|
||||
"hard": int(eval_result["em"]),
|
||||
"soft": eval_result["f1"],
|
||||
"fail_reason": fail_reason or ("" if eval_result["em"] else f"predicted '{eval_result['predicted_answer']}' but expected '{item.get('ground_truth', '')}'"),
|
||||
"agent_ok": not fail_reason,
|
||||
"n_turns": len(conversation),
|
||||
"last_finish_reason": last_response_metadata.get("finish_reason", ""),
|
||||
"student_system_prompt": system,
|
||||
"student_user_prompt": user,
|
||||
}
|
||||
return result
|
||||
def run_batch(
|
||||
items: list[dict],
|
||||
out_root: str,
|
||||
skill_content: str,
|
||||
*,
|
||||
workers: int = 8,
|
||||
max_tool_turns: int = 12,
|
||||
max_completion_tokens: int = 64000,
|
||||
search_mode: str = _DEFAULT_SEARCH_MODE,
|
||||
max_queries_per_turn: int = 4,
|
||||
search_api_url: str = "",
|
||||
search_auth_env: str = "OFFICEQA_CUSTOM_SEARCH_AUTH",
|
||||
search_provider: str = "duckduckgo",
|
||||
search_max_num_results: int = 4,
|
||||
search_timeout_seconds: int = 20,
|
||||
use_local_tools: bool = True,
|
||||
data_dirs: list[str] | str | None = None,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = "",
|
||||
) -> list[dict]:
|
||||
results_path = os.path.join(out_root, "results.jsonl")
|
||||
os.makedirs(out_root, exist_ok=True)
|
||||
done_ids: set[str] = set()
|
||||
existing: list[dict] = []
|
||||
if os.path.exists(results_path):
|
||||
with open(results_path, encoding="utf-8") as f:
|
||||
for line in f:
|
||||
try:
|
||||
row = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
done_ids.add(str(row.get("id")))
|
||||
existing.append(row)
|
||||
pending = [item for item in items if str(item["id"]) not in done_ids]
|
||||
if not pending:
|
||||
return existing
|
||||
total = len(existing) + len(pending)
|
||||
completed = len(existing)
|
||||
correct_count = sum(1 for r in existing if r.get("hard", 0))
|
||||
if existing:
|
||||
print(f" [rollout] resuming: {completed}/{total} already done", flush=True)
|
||||
|
||||
results = list(existing)
|
||||
with open(results_path, "a", encoding="utf-8") as outf, ThreadPoolExecutor(max_workers=workers) as ex:
|
||||
futs = {
|
||||
ex.submit(
|
||||
process_one,
|
||||
item,
|
||||
out_root,
|
||||
skill_content,
|
||||
max_tool_turns=max_tool_turns,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
search_mode=search_mode,
|
||||
max_queries_per_turn=max_queries_per_turn,
|
||||
search_api_url=search_api_url,
|
||||
search_auth_env=search_auth_env,
|
||||
search_provider=search_provider,
|
||||
search_max_num_results=search_max_num_results,
|
||||
search_timeout_seconds=search_timeout_seconds,
|
||||
use_local_tools=use_local_tools,
|
||||
data_dirs=data_dirs,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
): item
|
||||
for item in pending
|
||||
}
|
||||
for fut in as_completed(futs):
|
||||
res = fut.result()
|
||||
results.append(res)
|
||||
completed += 1
|
||||
if res.get("hard", 0):
|
||||
correct_count += 1
|
||||
acc = correct_count / completed if completed else 0
|
||||
print(
|
||||
f" [rollout] {completed}/{total} "
|
||||
f"(acc={acc:.3f}) id={res.get('id', '?')} "
|
||||
f"hard={res.get('hard', '?')}",
|
||||
flush=True,
|
||||
)
|
||||
outf.write(json.dumps(res, ensure_ascii=False) + "\n")
|
||||
outf.flush()
|
||||
return results
|
||||
@@ -0,0 +1,15 @@
|
||||
# OfficeQA Skill
|
||||
|
||||
## Retrieval Discipline
|
||||
- Start by narrowing to the most likely candidate file before reading long passages.
|
||||
- Prefer targeted search terms that name the exact entity, period, measure, or table concept from the question.
|
||||
- After a promising match, read only a small surrounding span and verify it matches the requested year, basis, and unit.
|
||||
|
||||
## Evidence Discipline
|
||||
- Extract the exact value from the retrieved text before doing any arithmetic.
|
||||
- Keep track of each operand's period, unit, and semantic role so nearby proxy values are not mixed in.
|
||||
- If the question asks for a transformed or derived quantity, compute only after confirming every operand.
|
||||
|
||||
## Final Answer Discipline
|
||||
- Return the final answer only after one last consistency check against the retrieved evidence.
|
||||
- Copy the final answer from a checked value, not from an unverified intermediate guess.
|
||||
@@ -0,0 +1,402 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import fnmatch
|
||||
import html
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from functools import lru_cache
|
||||
from html.parser import HTMLParser
|
||||
from pathlib import Path
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
_MAX_READ_CHARS = 4000
|
||||
_MAX_GREP_MATCHES = 20
|
||||
_MAX_GLOB_MATCHES = 50
|
||||
_MAX_ORACLE_PAGE_CHARS = 24000
|
||||
_MAX_ORACLE_CONTEXT_CHARS = 80000
|
||||
|
||||
|
||||
def _normalize_data_dirs(data_dirs: list[str] | tuple[str, ...] | str | None, project_root: Path) -> list[str]:
|
||||
if data_dirs is None:
|
||||
return []
|
||||
if isinstance(data_dirs, str):
|
||||
items = [part.strip() for chunk in data_dirs.split(os.pathsep) for part in chunk.split(",")]
|
||||
else:
|
||||
items = [str(item).strip() for item in data_dirs]
|
||||
resolved: list[str] = []
|
||||
for item in items:
|
||||
if not item:
|
||||
continue
|
||||
path = Path(item).expanduser()
|
||||
if not path.is_absolute():
|
||||
path = project_root / path
|
||||
resolved.append(str(path))
|
||||
return resolved
|
||||
|
||||
|
||||
def resolve_docs_roots(data_dirs: list[str] | tuple[str, ...] | str | None = None) -> list[str]:
|
||||
project_root = Path(__file__).resolve().parents[3]
|
||||
env_value = os.environ.get("OFFICEQA_DOCS_DIR", "").strip()
|
||||
candidates = _normalize_data_dirs(data_dirs, project_root)
|
||||
candidates.extend(_normalize_data_dirs(env_value, project_root))
|
||||
candidates.extend([
|
||||
str(project_root / "data" / "officeqa_docs_official"),
|
||||
str(project_root / "data" / "officeqa_smoke_docs"),
|
||||
os.path.expanduser("~/officeqa-sparse/treasury_bulletins_parsed"),
|
||||
os.path.expanduser("~/officeqa/treasury_bulletins_parsed"),
|
||||
])
|
||||
roots: list[str] = []
|
||||
seen: set[str] = set()
|
||||
for candidate in candidates:
|
||||
path = Path(candidate).expanduser()
|
||||
if not path.is_dir():
|
||||
continue
|
||||
transformed = path / "transformed"
|
||||
resolved = str((transformed if transformed.is_dir() else path).resolve())
|
||||
if resolved in seen:
|
||||
continue
|
||||
seen.add(resolved)
|
||||
roots.append(resolved)
|
||||
if not roots:
|
||||
raise FileNotFoundError("OfficeQA docs directory not found. Set OFFICEQA_DOCS_DIR or env.data_dirs.")
|
||||
return roots
|
||||
|
||||
|
||||
def _is_allowed(path: str, allowed_roots: list[str], allowed_files: list[str]) -> bool:
|
||||
try:
|
||||
resolved = str(Path(path).resolve())
|
||||
except FileNotFoundError:
|
||||
return False
|
||||
if not any(resolved.startswith(root + os.sep) or resolved == root for root in allowed_roots):
|
||||
return False
|
||||
if not allowed_files:
|
||||
return True
|
||||
base = os.path.basename(resolved)
|
||||
return base in allowed_files
|
||||
|
||||
|
||||
def resolve_candidate_files(source_files: list[str], allowed_roots: list[str]) -> list[str]:
|
||||
resolved: list[str] = []
|
||||
seen: set[str] = set()
|
||||
for root in allowed_roots:
|
||||
for dirpath, _, filenames in os.walk(root):
|
||||
for filename in filenames:
|
||||
if source_files and filename not in source_files:
|
||||
continue
|
||||
full = str(Path(dirpath, filename).resolve())
|
||||
if full in seen:
|
||||
continue
|
||||
seen.add(full)
|
||||
resolved.append(full)
|
||||
return resolved
|
||||
|
||||
|
||||
def _as_list(value: object) -> list[str]:
|
||||
if value is None:
|
||||
return []
|
||||
if isinstance(value, list):
|
||||
return [str(item).strip() for item in value if str(item).strip()]
|
||||
text = str(value).strip()
|
||||
if not text:
|
||||
return []
|
||||
try:
|
||||
loaded = json.loads(text)
|
||||
except json.JSONDecodeError:
|
||||
loaded = None
|
||||
if isinstance(loaded, list):
|
||||
return [str(item).strip() for item in loaded if str(item).strip()]
|
||||
if "\n" in text:
|
||||
return [part.strip() for part in text.splitlines() if part.strip()]
|
||||
return [text]
|
||||
|
||||
|
||||
def _extract_page_number(source_doc: str) -> int | None:
|
||||
text = str(source_doc or "").strip()
|
||||
if not text:
|
||||
return None
|
||||
parsed = urlparse(text)
|
||||
query = parse_qs(parsed.query)
|
||||
for key in ("page", "pagenum", "page_id"):
|
||||
for raw_value in query.get(key, []):
|
||||
try:
|
||||
return int(str(raw_value).strip())
|
||||
except ValueError:
|
||||
continue
|
||||
match = re.search(r"(?:[?&]|^)page=(\d+)", text)
|
||||
if match:
|
||||
return int(match.group(1))
|
||||
return None
|
||||
|
||||
|
||||
def _iter_oracle_refs(source_files: object, source_docs: object) -> list[tuple[str, int, str]]:
|
||||
files = _as_list(source_files)
|
||||
docs = _as_list(source_docs)
|
||||
refs: list[tuple[str, int, str]] = []
|
||||
seen: set[tuple[str, int, str]] = set()
|
||||
if not files or not docs:
|
||||
return refs
|
||||
for index, source_doc in enumerate(docs):
|
||||
page_number = _extract_page_number(source_doc)
|
||||
if page_number is None:
|
||||
continue
|
||||
if index < len(files):
|
||||
source_file = files[index]
|
||||
elif len(files) == 1:
|
||||
source_file = files[0]
|
||||
else:
|
||||
continue
|
||||
key = (source_file, page_number, source_doc)
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
refs.append(key)
|
||||
return refs
|
||||
|
||||
|
||||
def _parsed_root_candidates(docs_roots: list[str]) -> list[Path]:
|
||||
candidates: list[Path] = []
|
||||
seen: set[str] = set()
|
||||
for root in docs_roots:
|
||||
path = Path(root).expanduser()
|
||||
for candidate in (
|
||||
path,
|
||||
path.parent,
|
||||
path / "treasury_bulletins_parsed",
|
||||
path.parent / "treasury_bulletins_parsed",
|
||||
):
|
||||
resolved = str(candidate.resolve()) if candidate.exists() else str(candidate)
|
||||
if resolved in seen:
|
||||
continue
|
||||
seen.add(resolved)
|
||||
candidates.append(candidate)
|
||||
return candidates
|
||||
|
||||
|
||||
def _locate_parsed_json(source_file: str, docs_roots: list[str]) -> Path | None:
|
||||
source_path = Path(str(source_file).strip())
|
||||
stem = source_path.stem if source_path.suffix else source_path.name
|
||||
if not stem:
|
||||
return None
|
||||
candidate_names = [stem + ".json"]
|
||||
if source_path.suffix == ".json":
|
||||
candidate_names.insert(0, source_path.name)
|
||||
for root in _parsed_root_candidates(docs_roots):
|
||||
for name in candidate_names:
|
||||
path = root / "jsons" / name
|
||||
if path.is_file():
|
||||
return path
|
||||
return None
|
||||
|
||||
|
||||
class _TableMarkdownParser(HTMLParser):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(convert_charrefs=True)
|
||||
self.rows: list[list[str]] = []
|
||||
self._row: list[str] | None = None
|
||||
self._cell: list[str] | None = None
|
||||
|
||||
def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
|
||||
if tag.lower() == "tr":
|
||||
self._row = []
|
||||
elif tag.lower() in {"td", "th"} and self._row is not None:
|
||||
self._cell = []
|
||||
|
||||
def handle_data(self, data: str) -> None:
|
||||
if self._cell is not None:
|
||||
self._cell.append(data)
|
||||
|
||||
def handle_endtag(self, tag: str) -> None:
|
||||
normalized_tag = tag.lower()
|
||||
if normalized_tag in {"td", "th"} and self._cell is not None and self._row is not None:
|
||||
cell = re.sub(r"\s+", " ", "".join(self._cell)).strip()
|
||||
self._row.append(cell)
|
||||
self._cell = None
|
||||
elif normalized_tag == "tr" and self._row is not None:
|
||||
if any(cell for cell in self._row):
|
||||
self.rows.append(self._row)
|
||||
self._row = None
|
||||
self._cell = None
|
||||
|
||||
|
||||
def _escape_markdown_cell(value: str) -> str:
|
||||
return str(value).replace("\n", " ").replace("|", "\\|").strip()
|
||||
|
||||
|
||||
def _html_table_to_markdown(raw_html: str) -> str:
|
||||
parser = _TableMarkdownParser()
|
||||
try:
|
||||
parser.feed(raw_html)
|
||||
except Exception: # noqa: BLE001
|
||||
parser.rows = []
|
||||
rows = parser.rows
|
||||
if not rows:
|
||||
text = re.sub(r"(?is)<[^>]+>", " ", raw_html)
|
||||
return re.sub(r"\s+", " ", html.unescape(text)).strip()
|
||||
width = max(len(row) for row in rows)
|
||||
normalized_rows = [row + [""] * (width - len(row)) for row in rows]
|
||||
header = normalized_rows[0]
|
||||
body = normalized_rows[1:]
|
||||
lines = [
|
||||
"| " + " | ".join(_escape_markdown_cell(cell) for cell in header) + " |",
|
||||
"| " + " | ".join(["---"] * width) + " |",
|
||||
]
|
||||
lines.extend("| " + " | ".join(_escape_markdown_cell(cell) for cell in row) + " |" for row in body)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _render_parsed_content(content: str) -> str:
|
||||
text = content.strip()
|
||||
if not text:
|
||||
return ""
|
||||
if "<table" in text.lower():
|
||||
return _html_table_to_markdown(text)
|
||||
text = html.unescape(text)
|
||||
text = re.sub(r"\r\n?", "\n", text)
|
||||
return re.sub(r"\n{3,}", "\n\n", text).strip()
|
||||
|
||||
|
||||
def _element_page_ids(element: dict) -> set[int]:
|
||||
page_ids: set[int] = set()
|
||||
bbox = element.get("bbox")
|
||||
if not isinstance(bbox, list):
|
||||
return page_ids
|
||||
for box in bbox:
|
||||
if not isinstance(box, dict):
|
||||
continue
|
||||
raw_page_id = box.get("page_id")
|
||||
try:
|
||||
page_ids.add(int(raw_page_id))
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
return page_ids
|
||||
|
||||
|
||||
@lru_cache(maxsize=256)
|
||||
def _load_parsed_elements(json_path: str) -> tuple[dict, ...]:
|
||||
with open(json_path, encoding="utf-8") as f:
|
||||
payload = json.load(f)
|
||||
document = payload.get("document") if isinstance(payload, dict) else {}
|
||||
elements = document.get("elements") if isinstance(document, dict) else []
|
||||
if not isinstance(elements, list):
|
||||
return ()
|
||||
return tuple(element for element in elements if isinstance(element, dict))
|
||||
|
||||
|
||||
@lru_cache(maxsize=2048)
|
||||
def _render_parsed_page(json_path: str, page_number: int) -> str:
|
||||
rendered: list[str] = []
|
||||
for element in _load_parsed_elements(json_path):
|
||||
if page_number not in _element_page_ids(element):
|
||||
continue
|
||||
content = element.get("content")
|
||||
if not isinstance(content, str) or not content.strip():
|
||||
continue
|
||||
section = _render_parsed_content(content)
|
||||
if section:
|
||||
rendered.append(section)
|
||||
return "\n\n".join(rendered).strip()
|
||||
|
||||
|
||||
def build_oracle_parsed_pages_context(
|
||||
source_files: object,
|
||||
source_docs: object,
|
||||
docs_roots: list[str],
|
||||
*,
|
||||
max_page_chars: int = _MAX_ORACLE_PAGE_CHARS,
|
||||
max_total_chars: int = _MAX_ORACLE_CONTEXT_CHARS,
|
||||
evidence_note: str = "Treat it as primary document evidence and combine it with custom web search results when useful.",
|
||||
) -> str:
|
||||
"""Render oracle parsed OfficeQA pages referenced by source_docs/source_files."""
|
||||
refs = _iter_oracle_refs(source_files, source_docs)
|
||||
if not refs:
|
||||
return ""
|
||||
|
||||
blocks: list[str] = []
|
||||
total_chars = 0
|
||||
seen_pages: set[tuple[str, int]] = set()
|
||||
for source_file, page_number, source_doc in refs:
|
||||
json_path = _locate_parsed_json(source_file, docs_roots)
|
||||
if json_path is None:
|
||||
continue
|
||||
page_key = (str(json_path), page_number)
|
||||
if page_key in seen_pages:
|
||||
continue
|
||||
seen_pages.add(page_key)
|
||||
page_text = _render_parsed_page(str(json_path), page_number)
|
||||
if not page_text:
|
||||
continue
|
||||
if len(page_text) > max_page_chars:
|
||||
omitted = len(page_text) - max_page_chars
|
||||
page_text = page_text[:max_page_chars].rstrip() + f"\n\n[... {omitted} characters omitted from this parsed page ...]"
|
||||
block = (
|
||||
f"### {source_file} page {page_number}\n"
|
||||
f"Source URL: {source_doc}\n\n"
|
||||
f"{page_text}"
|
||||
)
|
||||
if total_chars + len(block) > max_total_chars:
|
||||
remaining = max_total_chars - total_chars
|
||||
if remaining <= 0:
|
||||
break
|
||||
block = block[:remaining].rstrip() + "\n\n[... oracle parsed page context truncated ...]"
|
||||
blocks.append(block)
|
||||
break
|
||||
blocks.append(block)
|
||||
total_chars += len(block)
|
||||
if not blocks:
|
||||
return ""
|
||||
return (
|
||||
"The following content is pre-parsed from the oracle OfficeQA source page(s). "
|
||||
f"{evidence_note.strip()}\n\n"
|
||||
+ "\n\n".join(blocks)
|
||||
)
|
||||
|
||||
|
||||
def run_tool(name: str, arguments: dict, *, allowed_roots: list[str], allowed_files: list[str]) -> tuple[str, str]:
|
||||
if name == "glob":
|
||||
pattern = str(arguments.get("pattern") or "*")
|
||||
matches: list[str] = []
|
||||
for root in allowed_roots:
|
||||
for dirpath, _, filenames in os.walk(root):
|
||||
for filename in filenames:
|
||||
if allowed_files and filename not in allowed_files:
|
||||
continue
|
||||
rel = os.path.relpath(os.path.join(dirpath, filename), root)
|
||||
if fnmatch.fnmatch(rel, pattern) or fnmatch.fnmatch(filename, pattern):
|
||||
matches.append(os.path.join(dirpath, filename))
|
||||
if len(matches) >= _MAX_GLOB_MATCHES:
|
||||
break
|
||||
if len(matches) >= _MAX_GLOB_MATCHES:
|
||||
break
|
||||
return f"glob(pattern={pattern!r})", "\n".join(matches) if matches else "[no matches]"
|
||||
|
||||
if name == "read":
|
||||
path = str(arguments.get("path") or "")
|
||||
if not path:
|
||||
return "read(path='')", "[read error: missing path]"
|
||||
if not _is_allowed(path, allowed_roots, allowed_files):
|
||||
return f"read(path={path!r})", "[read error: path not allowed]"
|
||||
start = max(int(arguments.get("start") or 1), 1)
|
||||
limit = max(int(arguments.get("limit") or 80), 1)
|
||||
with open(path, encoding="utf-8") as f:
|
||||
lines = f.readlines()
|
||||
excerpt = "".join(lines[start - 1:start - 1 + limit])
|
||||
return f"read(path={path!r}, start={start}, limit={limit})", excerpt[:_MAX_READ_CHARS] or "[empty file]"
|
||||
|
||||
if name == "grep":
|
||||
pattern = str(arguments.get("pattern") or "").lower()
|
||||
path = str(arguments.get("path") or "")
|
||||
if not pattern or not path:
|
||||
return f"grep(pattern={pattern!r}, path={path!r})", "[grep error: missing pattern or path]"
|
||||
if not _is_allowed(path, allowed_roots, allowed_files):
|
||||
return f"grep(pattern={pattern!r}, path={path!r})", "[grep error: path not allowed]"
|
||||
matches: list[str] = []
|
||||
with open(path, encoding="utf-8") as f:
|
||||
for idx, line in enumerate(f, start=1):
|
||||
if pattern in line.lower():
|
||||
matches.append(f"{idx}: {line.rstrip()}")
|
||||
if len(matches) >= _MAX_GREP_MATCHES:
|
||||
break
|
||||
return f"grep(pattern={pattern!r}, path={path!r})", "\n".join(matches) if matches else "[no matches]"
|
||||
|
||||
return name, f"[tool error: unknown tool {name}]"
|
||||
@@ -0,0 +1 @@
|
||||
"""SealQA environment package for ReflACT."""
|
||||
@@ -0,0 +1,130 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
from skillopt.datasets.base import BatchSpec
|
||||
from skillopt.envs.base import EnvAdapter
|
||||
from skillopt.envs.deep_reflect import run_no_reference_deep_reflect
|
||||
from skillopt.envs.sealqa.dataloader import SealQADataLoader
|
||||
from skillopt.envs.sealqa.rollout import run_batch
|
||||
from skillopt.gradient.reflect import run_minibatch_reflect
|
||||
|
||||
|
||||
class SealQAAdapter(EnvAdapter):
|
||||
def __init__(
|
||||
self,
|
||||
split_dir: str = '',
|
||||
workers: int = 4,
|
||||
analyst_workers: int = 8,
|
||||
failure_only: bool = False,
|
||||
minibatch_size: int = 8,
|
||||
edit_budget: int = 4,
|
||||
seed: int = 42,
|
||||
limit: int = 0,
|
||||
max_tool_turns: int = 12,
|
||||
use_deep_reflect: bool = False,
|
||||
deep_reflect_failures: int = 4,
|
||||
deep_reflect_successes: int = 2,
|
||||
) -> None:
|
||||
self.workers = workers
|
||||
self.analyst_workers = analyst_workers
|
||||
self.failure_only = failure_only
|
||||
self.minibatch_size = minibatch_size
|
||||
self.edit_budget = edit_budget
|
||||
self.max_tool_turns = max_tool_turns
|
||||
self.use_deep_reflect = use_deep_reflect
|
||||
self.deep_reflect_failures = deep_reflect_failures
|
||||
self.deep_reflect_successes = deep_reflect_successes
|
||||
self.dataloader = SealQADataLoader(split_dir=split_dir, seed=seed, limit=limit)
|
||||
|
||||
def setup(self, cfg: dict) -> None:
|
||||
super().setup(cfg)
|
||||
self.dataloader.setup(cfg)
|
||||
|
||||
def get_dataloader(self):
|
||||
return self.dataloader
|
||||
|
||||
def build_env_from_batch(self, batch: BatchSpec, **kwargs):
|
||||
return list(batch.payload or [])
|
||||
|
||||
def build_train_env(self, batch_size: int, seed: int, **kwargs):
|
||||
batch = self.dataloader.build_train_batch(batch_size=batch_size, seed=seed, **kwargs)
|
||||
return self.build_env_from_batch(batch, **kwargs)
|
||||
|
||||
def build_eval_env(self, env_num: int, split: str, seed: int, **kwargs):
|
||||
batch = self.dataloader.build_eval_batch(env_num=env_num, split=split, seed=seed, **kwargs)
|
||||
return self.build_env_from_batch(batch, **kwargs)
|
||||
|
||||
def rollout(self, env_manager, skill_content: str, out_dir: str, **kwargs) -> list[dict]:
|
||||
items: list[dict] = env_manager
|
||||
return run_batch(
|
||||
items=items,
|
||||
out_root=out_dir,
|
||||
skill_content=skill_content,
|
||||
workers=self.workers,
|
||||
max_tool_turns=self.max_tool_turns,
|
||||
diagnostic_mode=kwargs.get('diagnostic_mode', False),
|
||||
diagnostic_instruction=kwargs.get('diagnostic_instruction', ''),
|
||||
)
|
||||
|
||||
def reflect(self, results: list[dict], skill_content: str, out_dir: str, **kwargs) -> list[dict | None]:
|
||||
prediction_dir = kwargs.get('prediction_dir', os.path.join(out_dir, 'predictions'))
|
||||
patches_dir = kwargs.get('patches_dir', os.path.join(out_dir, 'patches'))
|
||||
random_seed = kwargs.get('random_seed')
|
||||
step_buffer_context = kwargs.get('step_buffer_context', '')
|
||||
return run_minibatch_reflect(
|
||||
results=results,
|
||||
skill_content=skill_content,
|
||||
prediction_dir=prediction_dir,
|
||||
patches_dir=patches_dir,
|
||||
workers=self.analyst_workers,
|
||||
failure_only=self.failure_only,
|
||||
minibatch_size=self.minibatch_size,
|
||||
edit_budget=self.edit_budget,
|
||||
random_seed=random_seed,
|
||||
error_system=self.get_error_minibatch_prompt(),
|
||||
success_system=self.get_success_minibatch_prompt(),
|
||||
step_buffer_context=step_buffer_context,
|
||||
update_mode=getattr(self, "_cfg", {}).get("skill_update_mode", "patch"),
|
||||
)
|
||||
|
||||
def deep_reflect(
|
||||
self,
|
||||
results: list[dict],
|
||||
skill_content: str,
|
||||
out_dir: str,
|
||||
**kwargs,
|
||||
) -> list[dict | None]:
|
||||
return run_no_reference_deep_reflect(
|
||||
self,
|
||||
results,
|
||||
skill_content,
|
||||
out_dir,
|
||||
env_manager=kwargs.get('env_manager'),
|
||||
prediction_dir=kwargs.get('prediction_dir'),
|
||||
random_seed=kwargs.get('random_seed'),
|
||||
step_buffer_context=kwargs.get('step_buffer_context', ''),
|
||||
output_requirements=[
|
||||
"- There is no hidden reference block. Use only the question, provided evidence, URL/fetch trace, student output, and evaluation result to infer what intermediate state is worth probing.",
|
||||
"- The instruction must explicitly request a short <analysis>...</analysis> block before the final <answer>...</answer>.",
|
||||
"- The readout should focus on effective time frame, conflicting evidence, decisive source, candidate answer, and answer-finalization rule.",
|
||||
"- Do not ask for exhaustive web summaries or a full chain-of-thought.",
|
||||
"- The instruction text should be ready to append directly to the student's prompt.",
|
||||
],
|
||||
metadata_builder=lambda item: {
|
||||
"id": str(item.get('id')),
|
||||
"task_type": str(item.get('task_type') or item.get('topic') or 'sealqa'),
|
||||
"question_preview": str(item.get('question') or '')[:200],
|
||||
"freshness": item.get('freshness', ''),
|
||||
"question_types": item.get('question_types', ''),
|
||||
"topic": item.get('topic', ''),
|
||||
},
|
||||
)
|
||||
|
||||
def get_task_types(self) -> list[str]:
|
||||
seen: list[str] = []
|
||||
for item in self.dataloader.train_items + self.dataloader.val_items + self.dataloader.test_items:
|
||||
task_type = str(item.get('task_type') or 'sealqa')
|
||||
if task_type not in seen:
|
||||
seen.append(task_type)
|
||||
return seen or ['sealqa']
|
||||
@@ -0,0 +1,37 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
from pathlib import Path
|
||||
|
||||
from skillopt.datasets.base import SplitDataLoader
|
||||
|
||||
|
||||
def _normalize_row(row: dict[str, str], index: int) -> dict:
|
||||
canary = str(row.get('canary') or '').strip()
|
||||
base_id = str(row.get('question_id') or row.get('id') or '').strip()
|
||||
if not base_id:
|
||||
base_id = f"{canary or 'sealqa'}:{index:04d}"
|
||||
return {
|
||||
'id': base_id,
|
||||
'question': str(row.get('question') or '').strip(),
|
||||
'ground_truth': str(row.get('answer') or row.get('ground_truth') or '').strip(),
|
||||
'answers': [str(row.get('answer') or row.get('ground_truth') or '').strip()],
|
||||
'task_type': str(row.get('topic') or 'sealqa').strip() or 'sealqa',
|
||||
'topic': str(row.get('topic') or 'sealqa').strip() or 'sealqa',
|
||||
'urls': str(row.get('urls') or '').strip(),
|
||||
'search_results': str(row.get('search_results') or '').strip(),
|
||||
'freshness': str(row.get('freshness') or '').strip(),
|
||||
'question_types': str(row.get('question_types') or '').strip(),
|
||||
'canary': canary,
|
||||
}
|
||||
|
||||
|
||||
class SealQADataLoader(SplitDataLoader):
|
||||
def load_split_items(self, split_path: str) -> list[dict]:
|
||||
path = Path(split_path)
|
||||
csv_files = sorted(path.glob('*.csv'))
|
||||
if not csv_files:
|
||||
raise FileNotFoundError(f'No .csv file found in {split_path}')
|
||||
with csv_files[0].open(encoding='utf-8', newline='') as f:
|
||||
reader = csv.DictReader(f)
|
||||
return [_normalize_row(row, idx) for idx, row in enumerate(reader, start=1)]
|
||||
@@ -0,0 +1,126 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
from openai import AzureOpenAI, OpenAI
|
||||
|
||||
GRADER_TEMPLATE = """
|
||||
Your job is to look at a question, a gold target, and a predicted answer, and then assign a grade of either [\"CORRECT\", \"INCORRECT\", \"NOT_ATTEMPTED\"].
|
||||
First, I will give examples of each grade, and then you will grade a new example.
|
||||
|
||||
The following are examples of CORRECT predicted answers.
|
||||
```
|
||||
Question: What are the names of Barack Obama's children?
|
||||
Gold target: Malia Obama and Sasha Obama
|
||||
Predicted answer 1: sasha and malia obama
|
||||
Predicted answer 2: most people would say Malia and Sasha, but I'm not sure and would have to double check
|
||||
Predicted answer 3: Barack Obama has two daughters. Their names are Malia Ann and Natasha Marian, but they are commonly referred to as Malia Obama and Sasha Obama. Malia was born on July 4, 1998, and Sasha was born on June 10, 2001.
|
||||
```
|
||||
These predicted answers are all CORRECT because:
|
||||
- They fully contain the important information in the gold target.
|
||||
- They do not contain any information that contradicts the gold target.
|
||||
- Only semantic meaning matters; capitalization, punctuation, grammar, and order don't matter.
|
||||
- Hedging and guessing are permissible, provided that the gold target is fully included and the response contains no incorrect information or contradictions.
|
||||
|
||||
The following are examples of INCORRECT predicted answers.
|
||||
```
|
||||
Question: What are the names of Barack Obama's children?
|
||||
Gold target: Malia and Sasha
|
||||
Predicted answer 1: Malia.
|
||||
Predicted answer 2: Malia, Sasha, and Susan.
|
||||
Predicted answer 3: Barack Obama does not have any children.
|
||||
```
|
||||
These predicted answers are all INCORRECT because a factual statement in the answer contradicts the gold target.
|
||||
|
||||
The following are examples of NOT_ATTEMPTED predicted answers.
|
||||
```
|
||||
Question: What are the names of Barack Obama's children?
|
||||
Gold target: Malia and Sasha
|
||||
Predicted answer 1: I don't know.
|
||||
Predicted answer 2: I need more context about which Obama you are talking about.
|
||||
```
|
||||
These predicted answers are all NOT_ATTEMPTED because the important information in the gold target is not included and there is no contradiction.
|
||||
|
||||
Here is a new example. Simply reply with either CORRECT, INCORRECT, NOT ATTEMPTED. Just return the letters \"A\", \"B\", or \"C\", with no text around it.
|
||||
```
|
||||
Question: {question}
|
||||
Gold target: {target}
|
||||
Predicted answer: {predicted_answer}
|
||||
```
|
||||
|
||||
Grade the predicted answer as one of:
|
||||
A: CORRECT
|
||||
B: INCORRECT
|
||||
C: NOT_ATTEMPTED
|
||||
""".strip()
|
||||
|
||||
|
||||
def _build_grader_client() -> tuple[OpenAI | AzureOpenAI, str]:
|
||||
import os
|
||||
|
||||
endpoint = os.environ.get('AZURE_OPENAI_ENDPOINT', '').strip()
|
||||
api_version = os.environ.get('AZURE_OPENAI_API_VERSION', '').strip() or '2025-04-01-preview'
|
||||
azure_key = os.environ.get('AZURE_OPENAI_API_KEY', '').strip()
|
||||
openai_key = os.environ.get('OPENAI_API_KEY', '').strip()
|
||||
api_key = azure_key or openai_key
|
||||
if endpoint and api_version and api_key:
|
||||
model = os.environ.get('SEALQA_GRADER_AZURE_MODEL', '').strip() or os.environ.get('SEALQA_GRADER_MODEL', '').strip() or os.environ.get('AZURE_MODEL_NAME', '').strip() or os.environ.get('TEACHER_DEPLOYMENT', '').strip() or 'gpt-5.4'
|
||||
client = AzureOpenAI(api_key=api_key, api_version=api_version, azure_endpoint=endpoint.rstrip('/'))
|
||||
return client, model
|
||||
|
||||
if openai_key:
|
||||
model = os.environ.get('SEALQA_GRADER_OPENAI_MODEL', '').strip() or os.environ.get('SEALQA_GRADER_MODEL', '').strip() or 'gpt-4.1-mini'
|
||||
return OpenAI(api_key=openai_key), model
|
||||
|
||||
raise ValueError('Missing grader credentials for SealQA scoring.')
|
||||
|
||||
|
||||
def _extract_text_content(content) -> str:
|
||||
if content is None:
|
||||
return ''
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
parts = []
|
||||
for part in content:
|
||||
if isinstance(part, dict) and part.get('type') == 'text':
|
||||
parts.append(str(part.get('text', '')))
|
||||
else:
|
||||
text = getattr(part, 'text', None)
|
||||
if text:
|
||||
parts.append(str(text))
|
||||
return '\n'.join(parts).strip()
|
||||
return str(content).strip()
|
||||
|
||||
|
||||
def _normalize_text(text: str) -> str:
|
||||
lowered = text.strip().lower()
|
||||
lowered = re.sub(r'\s+', ' ', lowered)
|
||||
lowered = re.sub(r'[^\w\s%.-]', '', lowered)
|
||||
return lowered.strip()
|
||||
|
||||
|
||||
def _fallback_score(ground_truth: str, predicted: str) -> float:
|
||||
gold = _normalize_text(ground_truth)
|
||||
pred = _normalize_text(predicted)
|
||||
if not gold or not pred:
|
||||
return 0.0
|
||||
if gold == pred:
|
||||
return 1.0
|
||||
if gold in pred or pred in gold:
|
||||
return 1.0
|
||||
return 0.0
|
||||
|
||||
|
||||
def score_sealqa(question: str, ground_truth: str, predicted: str) -> float:
|
||||
try:
|
||||
client, model = _build_grader_client()
|
||||
except ValueError:
|
||||
return _fallback_score(ground_truth, predicted)
|
||||
|
||||
prompt = GRADER_TEMPLATE.format(question=question, target=ground_truth, predicted_answer=predicted)
|
||||
completion = client.chat.completions.create(model=model, messages=[{'role': 'user', 'content': prompt}])
|
||||
content = _extract_text_content(completion.choices[0].message.content).strip().upper()
|
||||
if content.startswith('A'):
|
||||
return 1.0
|
||||
return 0.0
|
||||
@@ -0,0 +1,30 @@
|
||||
You are an expert failure-analysis agent for evidence-seeking factual question answering tasks.
|
||||
|
||||
You will be given MULTIPLE failed SealQA trajectories from a single minibatch and the current skill document. The trajectories may include tool calls such as search, fetch, local reads, or evidence gathering steps.
|
||||
|
||||
Your job is to identify COMMON failure patterns across the batch and propose concise skill edits.
|
||||
|
||||
## Failure Type Categories
|
||||
- retrieval_miss: the agent failed to gather the right evidence
|
||||
- evidence_conflict: the agent saw conflicting evidence but resolved it badly
|
||||
- answer_selection: the agent found evidence but chose the wrong final answer
|
||||
- not_attempted: the agent never reached a grounded answer
|
||||
- other: none of the above
|
||||
|
||||
Respond ONLY with a valid JSON object (no markdown fences, no extra text):
|
||||
{
|
||||
"batch_size": <number of trajectories analysed>,
|
||||
"failure_summary": [
|
||||
{"failure_type": "<type>", "count": <int>, "description": "<one-line>"}
|
||||
],
|
||||
"patch": {
|
||||
"reasoning": "<why these edits address the batch's common failures>",
|
||||
"edits": [
|
||||
{"op": "append", "content": "<markdown to add at end of skill>"},
|
||||
{"op": "insert_after", "target": "<exact heading/text to insert after>", "content": "<markdown>"},
|
||||
{"op": "replace", "target": "<exact text to replace>", "content": "<replacement>"},
|
||||
{"op": "delete", "target": "<exact text to remove>"}
|
||||
]
|
||||
}
|
||||
}
|
||||
Only include edits that are needed. "edits" can be an empty list if no patch is warranted.
|
||||
@@ -0,0 +1,19 @@
|
||||
You are an expert success-pattern analyst for evidence-seeking factual question answering tasks.
|
||||
|
||||
You will be given MULTIPLE successful SealQA trajectories from a single minibatch and the current skill document. Your job is to identify common evidence-gathering and answer-selection behaviors worth encoding in the skill.
|
||||
|
||||
Respond ONLY with a valid JSON object:
|
||||
{
|
||||
"batch_size": <number of trajectories analysed>,
|
||||
"success_patterns": ["<pattern 1>", "<pattern 2>"],
|
||||
"patch": {
|
||||
"reasoning": "<why these patterns are worth encoding>",
|
||||
"edits": [
|
||||
{"op": "append", "content": "<markdown>"},
|
||||
{"op": "insert_after", "target": "<heading/text>", "content": "<markdown>"},
|
||||
{"op": "replace", "target": "<old text>", "content": "<new text>"},
|
||||
{"op": "delete", "target": "<exact text to remove>"}
|
||||
]
|
||||
}
|
||||
}
|
||||
"edits" may be empty if the skill already covers all observed patterns.
|
||||
@@ -0,0 +1,3 @@
|
||||
You are an expert research assistant. Use the provided search evidence first, and only if that is insufficient, inspect the provided URL content fetched for you. Reconcile conflicting information when necessary and return a concise final answer grounded in the evidence you found.
|
||||
|
||||
{skill_section}Return the final answer inside <answer>...</answer> when you are ready.
|
||||
@@ -0,0 +1,284 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
from skillopt.envs.sealqa.evaluator import score_sealqa
|
||||
from skillopt.envs.sealqa.tool_runtime import web_fetch
|
||||
from skillopt.model import chat_student, get_student_backend, is_student_exec_backend
|
||||
from skillopt.model.codex_harness import prepare_workspace, render_skill_md, run_student_exec
|
||||
from skillopt.prompts import load_prompt
|
||||
|
||||
_FINAL_RE = re.compile(r"<answer>(.*?)</answer>", re.IGNORECASE | re.DOTALL)
|
||||
|
||||
|
||||
def _build_system(skill_content: str) -> str:
|
||||
if skill_content.strip():
|
||||
skill_section = f"## Skill\n{skill_content.strip()}\n\n"
|
||||
else:
|
||||
skill_section = ""
|
||||
return load_prompt("rollout_system", env="sealqa").format(skill_section=skill_section)
|
||||
|
||||
|
||||
def _build_user(item: dict, *, diagnostic_mode: bool = False, diagnostic_instruction: str = '') -> str:
|
||||
parts = [f"## Question\n{item['question']}"]
|
||||
if item.get('search_results'):
|
||||
parts.append(f"## Search Results\n{item['search_results']}")
|
||||
if item.get('urls'):
|
||||
parts.append(f"## URL Hints\n{item['urls']}")
|
||||
if item.get('freshness'):
|
||||
parts.append(f"## Freshness\n{item['freshness']}")
|
||||
if item.get('question_types'):
|
||||
parts.append(f"## Question Types\n{item['question_types']}")
|
||||
if diagnostic_mode and diagnostic_instruction.strip():
|
||||
parts.append(f"## Training Readout\n{diagnostic_instruction.strip()}")
|
||||
parts.append('Use the provided search evidence as your primary context. Do not rely on external tool use.')
|
||||
return "\n\n".join(parts)
|
||||
|
||||
|
||||
def _extract_answer(text: str) -> str:
|
||||
match = _FINAL_RE.search(text)
|
||||
if match:
|
||||
return match.group(1).strip()
|
||||
lines = [line.strip() for line in text.splitlines() if line.strip()]
|
||||
return lines[-1] if lines else text.strip()
|
||||
|
||||
|
||||
def _build_codex_skill(skill_content: str) -> str:
|
||||
return render_skill_md(
|
||||
skill_content,
|
||||
description="Dynamic ReflACT skill for solving the current SealQA evidence-grounded question.",
|
||||
preamble=(
|
||||
"Use this skill when answering the current SealQA question.\n"
|
||||
"Use the provided search evidence first, reconcile conflicts carefully,\n"
|
||||
"and return the final answer inside <answer>...</answer>."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _run_codex_once(
|
||||
*,
|
||||
pred_dir: str,
|
||||
skill_content: str,
|
||||
task_text: str,
|
||||
model: str,
|
||||
timeout: int,
|
||||
previous_response: str = '',
|
||||
) -> tuple[str, str, str, str]:
|
||||
task_parts = [task_text]
|
||||
if previous_response:
|
||||
task_parts.append(
|
||||
"## Previous Attempt\n"
|
||||
f"{previous_response}\n\n"
|
||||
"Review the evidence again and correct the final answer if needed."
|
||||
)
|
||||
final_task_text = "\n\n".join(task_parts)
|
||||
skill_md = _build_codex_skill(skill_content)
|
||||
work_dir = os.path.join(pred_dir, 'codex_exec')
|
||||
prepare_workspace(
|
||||
work_dir=work_dir,
|
||||
skill_md=skill_md,
|
||||
task_text=final_task_text,
|
||||
)
|
||||
prompt = (
|
||||
"Use the `skillopt-student` skill available in this workspace.\n"
|
||||
"Read `task.md`, answer the SealQA question using the provided evidence,\n"
|
||||
"and return the final answer inside <answer>...</answer>."
|
||||
)
|
||||
final_message, raw = run_student_exec(
|
||||
work_dir=work_dir,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
timeout=timeout,
|
||||
)
|
||||
return final_message or raw, raw, skill_md, final_task_text
|
||||
|
||||
|
||||
def process_one(
|
||||
item: dict,
|
||||
out_root: str,
|
||||
skill_content: str,
|
||||
*,
|
||||
max_tool_turns: int = 12,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = '',
|
||||
) -> dict:
|
||||
item_id = str(item['id'])
|
||||
pred_dir = os.path.join(out_root, 'predictions', item_id)
|
||||
os.makedirs(pred_dir, exist_ok=True)
|
||||
|
||||
system = _build_system(skill_content)
|
||||
user = _build_user(
|
||||
item,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
)
|
||||
conversation: list[dict] = [{'role': 'user', 'content': user}]
|
||||
final_response = ''
|
||||
final_answer = ''
|
||||
fail_reason = ''
|
||||
|
||||
try:
|
||||
if is_student_exec_backend():
|
||||
from skillopt.model import azure_openai as _llm
|
||||
|
||||
response, _raw, system, user_for_save = _run_codex_once(
|
||||
pred_dir=pred_dir,
|
||||
skill_content=skill_content,
|
||||
task_text=user,
|
||||
model=_llm.STUDENT_DEPLOYMENT,
|
||||
timeout=120,
|
||||
)
|
||||
final_response = response
|
||||
conversation.append({'type': 'message', 'content': response})
|
||||
if '<answer>' in response.lower():
|
||||
final_answer = _extract_answer(response)
|
||||
else:
|
||||
user = user_for_save
|
||||
else:
|
||||
response, _ = chat_student(
|
||||
system=system,
|
||||
user=user,
|
||||
max_completion_tokens=768,
|
||||
retries=5,
|
||||
stage='rollout',
|
||||
)
|
||||
final_response = response
|
||||
conversation.append({'type': 'message', 'content': response})
|
||||
if '<answer>' in response.lower():
|
||||
final_answer = _extract_answer(response)
|
||||
|
||||
if not final_answer:
|
||||
urls_text = str(item.get('urls') or '').strip()
|
||||
fetched_blocks = []
|
||||
for raw_url in re.findall(r'https?://[^\s\]\[\'\",]+', urls_text)[:2]:
|
||||
try:
|
||||
fetched = web_fetch(raw_url)
|
||||
except Exception as fetch_error: # noqa: BLE001
|
||||
fetched = f'URL: {raw_url}\n\n[fetch error: {fetch_error}]'
|
||||
fetched_blocks.append(fetched)
|
||||
conversation.append({'type': 'tool_call', 'cmd': f'web_fetch({raw_url!r})', 'obs': fetched})
|
||||
if fetched_blocks:
|
||||
retry_user = user + '\n\n## Fetched URL Content\n' + '\n\n'.join(fetched_blocks)
|
||||
if is_student_exec_backend():
|
||||
retry_response, _raw, system, retry_user = _run_codex_once(
|
||||
pred_dir=pred_dir,
|
||||
skill_content=skill_content,
|
||||
task_text=retry_user,
|
||||
model=_llm.STUDENT_DEPLOYMENT,
|
||||
timeout=120,
|
||||
previous_response=final_response,
|
||||
)
|
||||
else:
|
||||
retry_response, _ = chat_student(
|
||||
system=system,
|
||||
user=retry_user,
|
||||
max_completion_tokens=768,
|
||||
retries=5,
|
||||
stage='rollout',
|
||||
)
|
||||
final_response = retry_response
|
||||
conversation.append({'type': 'message', 'content': retry_response})
|
||||
if '<answer>' in retry_response.lower():
|
||||
final_answer = _extract_answer(retry_response)
|
||||
else:
|
||||
fail_reason = 'Model did not produce a final answer'
|
||||
else:
|
||||
fail_reason = 'Model did not produce a final answer'
|
||||
except Exception as e: # noqa: BLE001
|
||||
fail_reason = f'error: {e}'
|
||||
|
||||
with open(os.path.join(pred_dir, 'student_system_prompt.txt'), 'w', encoding='utf-8') as f:
|
||||
f.write(system)
|
||||
with open(os.path.join(pred_dir, 'student_user_prompt.txt'), 'w', encoding='utf-8') as f:
|
||||
f.write(user)
|
||||
with open(os.path.join(pred_dir, 'conversation.json'), 'w', encoding='utf-8') as f:
|
||||
json.dump(conversation, f, ensure_ascii=False, indent=2)
|
||||
|
||||
score = score_sealqa(item.get('question', ''), item.get('ground_truth', ''), final_answer) if final_answer else 0.0
|
||||
result = {
|
||||
'id': item_id,
|
||||
'question': item.get('question', ''),
|
||||
'task_type': item.get('task_type', 'sealqa'),
|
||||
'task_description': item.get('question', ''),
|
||||
'predicted_answer': final_answer,
|
||||
'response': final_response,
|
||||
'ground_truth': item.get('ground_truth', ''),
|
||||
'hard': int(score >= 1.0),
|
||||
'soft': float(score),
|
||||
'fail_reason': fail_reason or ('' if score >= 1.0 else f"predicted '{final_answer}' but expected '{item.get('ground_truth', '')}'"),
|
||||
'agent_ok': not fail_reason,
|
||||
'n_turns': len(conversation),
|
||||
'student_system_prompt': system,
|
||||
'student_user_prompt': user,
|
||||
}
|
||||
return result
|
||||
|
||||
|
||||
def run_batch(
|
||||
items: list[dict],
|
||||
out_root: str,
|
||||
skill_content: str,
|
||||
*,
|
||||
workers: int = 4,
|
||||
max_tool_turns: int = 12,
|
||||
diagnostic_mode: bool = False,
|
||||
diagnostic_instruction: str = '',
|
||||
) -> list[dict]:
|
||||
results_path = os.path.join(out_root, 'results.jsonl')
|
||||
os.makedirs(out_root, exist_ok=True)
|
||||
|
||||
done_ids: set[str] = set()
|
||||
existing: list[dict] = []
|
||||
if os.path.exists(results_path):
|
||||
with open(results_path, encoding='utf-8') as f:
|
||||
for line in f:
|
||||
try:
|
||||
row = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
done_ids.add(str(row.get('id')))
|
||||
existing.append(row)
|
||||
|
||||
pending = [item for item in items if str(item['id']) not in done_ids]
|
||||
if not pending:
|
||||
return existing
|
||||
|
||||
total = len(existing) + len(pending)
|
||||
completed = len(existing)
|
||||
correct_count = sum(1 for r in existing if r.get("hard", 0))
|
||||
if existing:
|
||||
print(f" [rollout] resuming: {completed}/{total} already done", flush=True)
|
||||
|
||||
results = list(existing)
|
||||
with open(results_path, 'a', encoding='utf-8') as outf, ThreadPoolExecutor(max_workers=workers) as ex:
|
||||
futs = {
|
||||
ex.submit(
|
||||
process_one,
|
||||
item,
|
||||
out_root,
|
||||
skill_content,
|
||||
max_tool_turns=max_tool_turns,
|
||||
diagnostic_mode=diagnostic_mode,
|
||||
diagnostic_instruction=diagnostic_instruction,
|
||||
): item
|
||||
for item in pending
|
||||
}
|
||||
for fut in as_completed(futs):
|
||||
res = fut.result()
|
||||
results.append(res)
|
||||
completed += 1
|
||||
if res.get("hard", 0):
|
||||
correct_count += 1
|
||||
acc = correct_count / completed if completed else 0
|
||||
print(
|
||||
f" [rollout] {completed}/{total} "
|
||||
f"(acc={acc:.3f}) id={res.get('id', '?')} "
|
||||
f"hard={res.get('hard', '?')}",
|
||||
flush=True,
|
||||
)
|
||||
outf.write(json.dumps(res, ensure_ascii=False) + '\n')
|
||||
outf.flush()
|
||||
return results
|
||||
@@ -0,0 +1,11 @@
|
||||
# SealQA Skill
|
||||
|
||||
## Evidence Gathering
|
||||
- Search for the most directly relevant evidence before answering.
|
||||
- If multiple sources conflict, prefer the source that best matches the question's entity, date, and scope.
|
||||
- Keep notes on which evidence directly answers the question versus which evidence is only contextual.
|
||||
|
||||
## Final Answer Discipline
|
||||
- Do not answer until the supporting evidence is specific enough.
|
||||
- Choose the final answer that is best grounded in the gathered evidence.
|
||||
- Keep the final answer concise.
|
||||
@@ -0,0 +1,30 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import html
|
||||
import re
|
||||
from urllib.request import Request, urlopen
|
||||
|
||||
DEFAULT_USER_AGENT = (
|
||||
'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 '
|
||||
'(KHTML, like Gecko) Chrome/135.0 Safari/537.36'
|
||||
)
|
||||
_MAX_FETCH_CHARS = 6000
|
||||
|
||||
|
||||
def _strip_html(raw_html: str) -> str:
|
||||
cleaned = re.sub(r'(?is)<script.*?>.*?</script>', ' ', raw_html)
|
||||
cleaned = re.sub(r'(?is)<style.*?>.*?</style>', ' ', cleaned)
|
||||
cleaned = re.sub(r'(?is)<[^>]+>', ' ', cleaned)
|
||||
cleaned = html.unescape(cleaned)
|
||||
return re.sub(r'\s+', ' ', cleaned).strip()
|
||||
|
||||
|
||||
def web_fetch(url: str, max_chars: int = _MAX_FETCH_CHARS) -> str:
|
||||
req = Request(url, headers={'User-Agent': DEFAULT_USER_AGENT})
|
||||
with urlopen(req, timeout=20) as response:
|
||||
body = response.read().decode('utf-8', errors='ignore')
|
||||
text = _strip_html(body)
|
||||
if len(text) > max_chars:
|
||||
omitted = len(text) - max_chars
|
||||
text = text[:max_chars] + f"\n\n[... {omitted} characters omitted ...]"
|
||||
return f"URL: {url}\n\n{text}"
|
||||
@@ -0,0 +1 @@
|
||||
"""SearchQA environment package for ReflACT."""
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user