From a3de87afb0f548be07f77b48527c5bc94729266d Mon Sep 17 00:00:00 2001 From: doxav <> Date: Thu, 11 Jun 2026 21:32:09 +0200 Subject: [PATCH] Add new HF QA tasks and enhance existing benchmarks - Updated `hf_tasks.yaml` to include new QA tasks: MuSiQue, GSM8K, ARC-Challenge, QASC, DROP, QuALITY, and Qasper with detailed configurations. - Implemented `mcqa.py` for handling multiple-choice questions, including task conversion and feedback scoring. - Created `musique.py` for MuSiQue task handling, focusing on alias scoring and context formatting. - Developed `qasper.py` for Qasper task processing, including context formatting and answer extraction. - Enhanced `hotpot_qa.py` to improve feedback mechanisms and response handling. - Added comprehensive tests in `test_hf_qa_ext.py` to validate new tasks and their functionalities. - Introduced `test_noop_trainer.py` to ensure NoOpTrainer operates correctly within the Trace-Bench runner. - Updated `trace_bench/runner.py` to support new model token handling and optimizer integration. - Refined `noop_trainer.py` to improve logging and parameter handling during training. --- README.md | 24 ++- benchmarks/hf_qa/bbeh.py | 18 +- benchmarks/hf_qa/common.py | 191 +++++++++++++++++++ benchmarks/hf_qa/drop_qa.py | 87 +++++++++ benchmarks/hf_qa/gsm8k.py | 66 +++++++ benchmarks/hf_qa/hf_qa_loader.py | 262 +++++++++++++++++++++------ benchmarks/hf_qa/hf_tasks.yaml | 157 +++++++++++++++- benchmarks/hf_qa/hotpot_qa.py | 18 +- benchmarks/hf_qa/mcqa.py | 100 ++++++++++ benchmarks/hf_qa/musique.py | 72 ++++++++ benchmarks/hf_qa/qasper.py | 172 ++++++++++++++++++ tests/m1/test_hf_qa_ext.py | 234 ++++++++++++++++++++++++ tests/m2/test_noop_trainer.py | 31 ++++ tests/m3/test_llm_metadata.py | 72 +++++++- trace_bench/runner.py | 48 ++++- trace_bench/trainers/noop_trainer.py | 22 ++- 16 files changed, 1501 insertions(+), 73 deletions(-) create mode 100644 benchmarks/hf_qa/common.py create mode 100644 benchmarks/hf_qa/drop_qa.py create mode 100644 benchmarks/hf_qa/gsm8k.py create mode 100644 benchmarks/hf_qa/mcqa.py create mode 100644 benchmarks/hf_qa/musique.py create mode 100644 benchmarks/hf_qa/qasper.py create mode 100644 tests/m1/test_hf_qa_ext.py create mode 100644 tests/m2/test_noop_trainer.py diff --git a/README.md b/README.md index 40e3aae..e13c5d6 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ A benchmarking framework for evaluating LLM-as-optimizer algorithms, built on [OpenTrace](https://github.com/AgentOpt/Trace). -Trace-Bench provides a **CLI**, **Gradio UI**, and **notebook workflows** to run reproducible experiments across multiple benchmarks (LLM4AD, VeriBench, KernelBench) with fair comparisons between trainers, optimizers, and LLM backends. +Trace-Bench provides a **CLI**, **Gradio UI**, and **notebook workflows** to run reproducible experiments across LLM4AD, VeriBench, KernelBench, and HuggingFace QA tasks with fair comparisons between trainers, optimizers, and LLM backends. **Full documentation:** [docs/](docs/README.md) @@ -25,8 +25,8 @@ pip install -e ../OpenTrace # Install Trace-Bench (core) pip install -e . -# Install with optional extras (can be combined): -pip install -e ".[hf]" # HuggingFace QA benchmarks (HotpotQA, BBEH, …) +# Optional extras can be combined: +pip install -e ".[hf]" # HuggingFace QA benchmarks pip install -e ".[dspy]" # DSPy trainer adapter pip install -e ".[all-external]" # everything ``` @@ -46,7 +46,17 @@ trace-bench ui --runs-dir runs ## Real-Mode Setup -To run benchmarks with actual LLM calls, configure an API provider: +To run benchmarks with actual LLM calls, configure a provider through environment variables or a config `llm` block. + +Direct OpenAI: + +```bash +export OPENAI_API_KEY="sk-..." +export TRACE_DEFAULT_LLM_BACKEND="LiteLLM" +export TRACE_LITELLM_MODEL="gpt-5.4-nano" +``` + +OpenRouter-compatible endpoint: ```bash export OPENROUTER_API_KEY="sk-or-v1-..." @@ -70,6 +80,7 @@ trace_bench/ # Python package # OpenTrace and are discovered automatically. benchmarks/ LLM4AD/ # 65 algorithm design tasks + hf_qa/ # HuggingFace QA tasks KernelBench/ # CUDA kernel optimization Veribench/ # Lean 4 formal verification configs/ # YAML experiment configs @@ -84,6 +95,9 @@ runs/ # Output directory (gitignored) ### LLM4AD (65 tasks) Algorithm design tasks from [LLM4AD](https://github.com/Optima-CityU/LLM4AD): optimization (basic, constructive, CO-Bench), machine learning, and scientific discovery. +### HuggingFace QA +Config-driven QA tasks under `hf:*`: HotpotQA, BBEH, MuSiQue, GSM8K, ARC-Challenge, QASC, DROP, QuALITY, and Qasper. Install with `pip install -e ".[hf]"`. + ### VeriBench (~140 tasks) Formal verification tasks translating Python to Lean 4. Integrated via `trace_bench/veribench_adapter.py` with dual discovery: local entrypoint module or [HuggingFace dataset](https://huggingface.co/datasets/allenanie/veribench_with_prompts) fallback. @@ -115,7 +129,7 @@ See [docs/running-experiments.md](docs/running-experiments.md) for full CLI usag | Extra | Installs | Enables | |-------|----------|---------| -| `hf` | `datasets>=2.0` | `hf:hotpot_qa`, `hf:bbeh/*` tasks | +| `hf` | `datasets>=2.0` | `hf:*` QA tasks | | `dspy` | `dspy-ai>=2.0` | `DSPyTrainer` in `trace_bench/trainers/` | | `all-external` | all of the above | everything | diff --git a/benchmarks/hf_qa/bbeh.py b/benchmarks/hf_qa/bbeh.py index 05a4c17..2329dcf 100644 --- a/benchmarks/hf_qa/bbeh.py +++ b/benchmarks/hf_qa/bbeh.py @@ -5,13 +5,20 @@ Reference: ``Predict`` / ``BigBenchGuide`` classes in ``bbeh_trace.py``. """ + from __future__ import annotations import re -from typing import Any +from typing import Any, Tuple + +try: + from .common import response_text +except ImportError: + from common import response_text try: from opto.trainer.guide import Guide + _OPTO_AVAILABLE = True except Exception: _OPTO_AVAILABLE = False @@ -21,6 +28,7 @@ # Data utilities # --------------------------------------------------------------------------- + def format_context(raw: Any) -> str: """No-op — BBEH tasks have no context field.""" return "" @@ -56,6 +64,7 @@ def eval_metric(true: str, prediction: str) -> bool: # --------------------------------------------------------------------------- if _OPTO_AVAILABLE: + class BBEHGuide(Guide): """Evaluation guide for BigBenchExtraHard tasks. @@ -63,8 +72,10 @@ class BBEHGuide(Guide): Uses ``eval_metric`` which handles both multiple-choice and free-form answers. """ - def get_feedback(self, task: str, response: Any, info: str, **kwargs: Any): - response_str = str(getattr(response, "data", response)) + def get_feedback( + self, task: str, response: Any, info: str, **kwargs: Any + ) -> Tuple[float, str]: + response_str = response_text(response) if eval_metric(info, response_str): return 1.0, "The answer is correct." return 0.0, ( @@ -74,6 +85,7 @@ def get_feedback(self, task: str, response: Any, info: str, **kwargs: Any): ) else: + class BBEHGuide: # type: ignore[no-redef] pass diff --git a/benchmarks/hf_qa/common.py b/benchmarks/hf_qa/common.py new file mode 100644 index 0000000..8b26993 --- /dev/null +++ b/benchmarks/hf_qa/common.py @@ -0,0 +1,191 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +import re +import string +from typing import Any, Dict, List, Sequence, Tuple + + +@dataclass +class HFQATask: + """A single QA example with formatted context and optional task metadata.""" + + question: str + context: str + answer: str + meta: Dict[str, Any] = field(default_factory=dict) + + +def normalize_text(text: str) -> str: + """Normalize answer text for exact-match and F1-style comparisons.""" + normalized = str(text).lower() + normalized = re.sub(r"\b(a|an|the)\b", " ", normalized) + normalized = normalized.translate(str.maketrans("", "", string.punctuation)) + return " ".join(normalized.split()) + + +def token_f1(prediction: str, gold: str) -> float: + """Return token-level F1 between normalized prediction and gold text.""" + 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 + + counts: Dict[str, int] = {} + for token in pred_tokens: + counts[token] = counts.get(token, 0) + 1 + + overlap = 0 + for token in gold_tokens: + if counts.get(token, 0) > 0: + overlap += 1 + counts[token] -= 1 + if overlap == 0: + return 0.0 + + precision = overlap / len(pred_tokens) + recall = overlap / len(gold_tokens) + return 2 * precision * recall / (precision + recall) + + +def extract_last_number(text: str) -> str: + """Extract the final integer or decimal number from a model response.""" + matches = re.findall(r"[-+]?\d[\d,]*\.?\d*", str(text)) + return matches[-1].replace(",", "") if matches else "" + + +def extract_last_choice(text: str) -> str: + """Extract the final parenthesized or standalone choice label.""" + parenthesized = re.findall(r"\(([A-Za-z0-9]+)\)", str(text)) + if parenthesized: + return parenthesized[-1] + standalone = re.findall(r"\b([A-Z])\b", str(text)) + return standalone[-1] if standalone else "" + + +def response_text(response: Any) -> str: + """Extract assistant text from common raw LLM response containers.""" + value = getattr(response, "data", response) + + choices = ( + value.get("choices") + if isinstance(value, dict) + else getattr(value, "choices", None) + ) + if choices: + first_choice = choices[0] + message = ( + first_choice.get("message") + if isinstance(first_choice, dict) + else getattr(first_choice, "message", None) + ) + if message is not None: + content = ( + message.get("content") + if isinstance(message, dict) + else getattr(message, "content", None) + ) + if content is not None: + return str(content) + text = ( + first_choice.get("text") + if isinstance(first_choice, dict) + else getattr(first_choice, "text", None) + ) + if text is not None: + return str(text) + + content = ( + value.get("content") + if isinstance(value, dict) + else getattr(value, "content", None) + ) + if content is not None: + return str(content) + return str(value) + + +def choice_pairs(raw: Any) -> List[Tuple[str, str]]: + """Return normalized (label, text) pairs from common HF MCQ schemas.""" + if isinstance(raw, dict): + raw_texts = raw.get("text") or raw.get("texts") or raw.get("options") or [] + texts = [str(text) for text in raw_texts] + raw_labels = raw.get("label") or raw.get("labels") or [] + labels = [str(label) for label in raw_labels] + if not labels: + labels = [chr(65 + index) for index in range(len(texts))] + return list(zip(labels, texts)) + + if isinstance(raw, list): + if raw and isinstance(raw[0], dict): + pairs: List[Tuple[str, str]] = [] + for index, item in enumerate(raw): + label = str(item.get("label", chr(65 + index))) + text = str(item.get("text", item.get("choice", ""))) + pairs.append((label, text)) + return pairs + return [(chr(65 + index), str(choice)) for index, choice in enumerate(raw)] + + return [] + + +def answer_to_choice_label( + raw_answer: Any, + pairs: Sequence[Tuple[str, str]], + index_base: int = 0, +) -> str: + """Map a raw HF answer value to the option label used by the prompt.""" + if not isinstance(index_base, int) or isinstance(index_base, bool): + raise TypeError(f"index_base must be an integer, got {index_base!r}") + + if isinstance(raw_answer, int) and not isinstance(raw_answer, bool): + index = raw_answer - index_base + if 0 <= index < len(pairs): + return pairs[index][0] + return str(raw_answer) + + answer = str(raw_answer).strip() + for label, text in pairs: + if answer == label or normalize_text(answer) == normalize_text(text): + return label + return answer + + +def format_mcq_prompt( + question: str, pairs: Sequence[Tuple[str, str]], context: str = "" +) -> str: + """Format a multiple-choice prompt with optional context.""" + prefix = f"Context:\n{context}\n\n" if context else "" + options = "\n".join(f"({label}) {text}" for label, text in pairs) + return f"{prefix}Question: {question}\nOptions:\n{options}\nAnswer:" + + +def split_aliases(raw: str) -> List[str]: + """Split a delimiter-separated alias string into non-empty aliases.""" + return [item.strip() for item in str(raw).split("||") if item.strip()] + + +def normalize_aliases(aliases: Sequence[str]) -> List[str]: + """Deduplicate aliases by normalized text while preserving display text.""" + output: List[str] = [] + seen = set() + for alias in aliases: + key = normalize_text(alias) + if key and key not in seen: + seen.add(key) + output.append(str(alias).strip()) + return output + + +def best_alias_score(prediction: str, aliases: Sequence[str]) -> float: + """Score a prediction against aliases using exact, contains, then token F1.""" + if not aliases: + return 0.0 + normalized_prediction = normalize_text(prediction) + if any(normalized_prediction == normalize_text(alias) for alias in aliases): + return 1.0 + if any(normalize_text(alias) in normalized_prediction for alias in aliases): + return 1.0 + return max(token_f1(prediction, alias) for alias in aliases) diff --git a/benchmarks/hf_qa/drop_qa.py b/benchmarks/hf_qa/drop_qa.py new file mode 100644 index 0000000..0d0b67e --- /dev/null +++ b/benchmarks/hf_qa/drop_qa.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Tuple + +try: + from .common import HFQATask, best_alias_score, normalize_aliases, response_text +except ImportError: + from common import HFQATask, best_alias_score, normalize_aliases, response_text + +try: + from opto.trainer.guide import Guide +except Exception: + Guide = object # type: ignore[assignment,misc] + + +def _flatten_answer_values(raw: Any) -> List[str]: + """Extract non-empty answer strings from common DROP answer shapes.""" + values: List[str] = [] + if isinstance(raw, dict): + for key in ("spans", "number", "numbers", "date", "dates"): + value = raw.get(key) + if isinstance(value, list): + values.extend(str(item) for item in value if str(item).strip()) + elif value is not None and str(value).strip(): + values.append(str(value)) + elif isinstance(raw, list): + values.extend(str(item) for item in raw if str(item).strip()) + elif raw is not None and str(raw).strip(): + values.append(str(raw)) + return values + + +def _aliases_from_drop_answer(raw: Any) -> List[str]: + """Return normalized DROP answer aliases.""" + return normalize_aliases(_flatten_answer_values(raw)) + + +def row_to_tasks(row: Dict[str, Any], cfg: Dict[str, Any]) -> List[HFQATask]: + """Convert one DROP row into a passage QA task.""" + question_field = cfg.get("question_field", "question") + answer_field = cfg.get("answer_field", "answers_spans") + context_field = cfg.get("context_field", "passage") + if question_field not in row: + raise KeyError(f"DROP row missing question field {question_field!r}") + + aliases = _aliases_from_drop_answer(row.get(answer_field)) + return [ + HFQATask( + question=str(row[question_field]), + context=str(row.get(context_field, "")), + answer=aliases[0] if aliases else "", + meta={"aliases": aliases}, + ) + ] + + +def format_context(raw: Any) -> str: + """Return plain-string context for DROP tasks.""" + return str(raw or "") + + +def make_dataset(tasks: List[HFQATask]) -> Dict[str, List[Any]]: + """Return the default context-rich QA dataset shape.""" + return {"inputs": list(tasks), "infos": list(tasks)} + + +class DROPGuide(Guide): # type: ignore[misc] + """Guide for alias-tolerant DROP answer scoring.""" + + def get_feedback( + self, task: Any, response: Any, info: HFQATask, **kwargs: Any + ) -> Tuple[float, str]: + """Score a DROP response against the gold aliases.""" + response_str = response_text(response) + aliases = info.meta.get("aliases", [info.answer]) + score = best_alias_score(response_str, aliases) + if score >= 0.999: + return 1.0, "Correct." + return score, ( + f"Incorrect. Expected one of {aliases}, got '{response_str}'. " + "Improve passage-based arithmetic/discrete reasoning and answer normalization." + ) + + +def make_guide() -> Any: + """Return an instantiated DROP guide.""" + return DROPGuide() diff --git a/benchmarks/hf_qa/gsm8k.py b/benchmarks/hf_qa/gsm8k.py new file mode 100644 index 0000000..50c618b --- /dev/null +++ b/benchmarks/hf_qa/gsm8k.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Tuple + +try: + from .common import HFQATask, extract_last_number, response_text +except ImportError: + from common import HFQATask, extract_last_number, response_text + +try: + from opto.trainer.guide import Guide +except Exception: + Guide = object # type: ignore[assignment,misc] + + +def row_to_tasks(row: Dict[str, Any], cfg: Dict[str, Any]) -> List[HFQATask]: + """Convert one GSM8K row into a numeric-answer task.""" + question_field = cfg.get("question_field", "question") + answer_field = cfg.get("answer_field", "answer") + if question_field not in row: + raise KeyError(f"GSM8K row missing question field {question_field!r}") + if answer_field not in row: + raise KeyError(f"GSM8K row missing answer field {answer_field!r}") + + raw_answer = str(row[answer_field]) + gold = raw_answer.split("####")[-1].strip() + return [ + HFQATask( + question=str(row[question_field]), + context="", + answer=gold, + ) + ] + + +def format_context(raw: Any) -> str: + """Return an empty context because GSM8K is question-only.""" + return "" + + +def make_dataset(tasks: List[HFQATask]) -> Dict[str, List[Any]]: + """Return the question-only dataset shape used by BBEH-style agents.""" + return {"inputs": [task.question for task in tasks], "infos": list(tasks)} + + +class GSM8KGuide(Guide): # type: ignore[misc] + """Guide for exact final-number GSM8K scoring.""" + + def get_feedback( + self, task: Any, response: Any, info: HFQATask, **kwargs: Any + ) -> Tuple[float, str]: + """Score a response by comparing the final extracted number.""" + response_str = response_text(response) + prediction = extract_last_number(response_str) + gold = extract_last_number(info.answer) + if prediction and prediction == gold: + return 1.0, "The numeric answer is correct." + return 0.0, ( + f"The numeric answer is wrong. Expected final number: '{gold}', got '{prediction or response_str}'. " + "Improve arithmetic reasoning and answer extraction so the final number is correct." + ) + + +def make_guide() -> Any: + """Return an instantiated GSM8K guide.""" + return GSM8KGuide() diff --git a/benchmarks/hf_qa/hf_qa_loader.py b/benchmarks/hf_qa/hf_qa_loader.py index be2f57e..7c3f06c 100644 --- a/benchmarks/hf_qa/hf_qa_loader.py +++ b/benchmarks/hf_qa/hf_qa_loader.py @@ -6,6 +6,11 @@ hotpot_qa.py — make_dataset, format_context, check_answer, HotpotQAGuide bbeh.py — make_dataset, format_context, eval_metric, BBEHGuide + musique.py — row_to_tasks, format_context, MuSiQueGuide + gsm8k.py — row_to_tasks, GSM8KGuide + mcqa.py — row_to_tasks, MCQGuide + drop_qa.py — row_to_tasks, DROPGuide + qasper.py — row_to_tasks, format_context, QasperGuide agents/trace_agent.py — TraceQAAgent, TraceBBEHAgent agents/dspy_agent.py — DSPyQAAgent, DSPyBBEHAgent @@ -13,14 +18,16 @@ - ``make_dataset(tasks: list) -> dict`` — convert HFQATask list to {inputs, infos} - ``format_context(raw: Any) -> str`` — flatten the raw HF context field - ``make_guide() -> Guide`` — return an instantiated guide +Optional task hooks: + - ``row_to_tasks(row: dict, cfg: dict) -> list[HFQATask]`` Each agent module must implement: - ``make_agent(agent_class: str, **kwargs) -> agent`` Supported tasks are declared in ``hf_tasks.yaml``. Adding a new task that -fits the existing agent patterns only requires a YAML entry. Adding a new -agent pattern requires a new task module + ``agent_class`` wiring below. -Adding a new framework requires a new agent module + ``framework`` wiring below. +fits existing task semantics + agent patterns mostly requires a YAML entry. +Adding new task semantics requires a task module; adding a new framework +requires a new agent module + ``framework`` wiring below. Usage in a YAML config:: @@ -35,24 +42,44 @@ eval_kwargs: framework: dspy # use the DSPy agent instead """ + from __future__ import annotations -from dataclasses import dataclass from pathlib import Path -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple import yaml -import hotpot_qa as _hotpot_qa_mod -import bbeh as _bbeh_mod - -# Registry mapping agent_class (from hf_tasks.yaml) → task module. +try: + from . import bbeh as _bbeh_mod + from . import drop_qa as _drop_qa_mod + from . import gsm8k as _gsm8k_mod + from . import hotpot_qa as _hotpot_qa_mod + from . import mcqa as _mcqa_mod + from . import musique as _musique_mod + from . import qasper as _qasper_mod + from .common import HFQATask +except ImportError: + import bbeh as _bbeh_mod + import drop_qa as _drop_qa_mod + import gsm8k as _gsm8k_mod + import hotpot_qa as _hotpot_qa_mod + import mcqa as _mcqa_mod + import musique as _musique_mod + import qasper as _qasper_mod + from common import HFQATask + +# Registry mapping task_type (from hf_tasks.yaml) → task module. # Task modules provide: make_dataset(), format_context(), make_guide(). -# To add a new task type: create a module implementing the interface above, -# then add it here and add a YAML entry with the matching agent_class. +# `task_type` defaults to `agent_class` for backward compatibility. _TASK_MODULES: Dict[str, Any] = { "hfqa": _hotpot_qa_mod, "bbeh": _bbeh_mod, + "drop": _drop_qa_mod, + "gsm8k": _gsm8k_mod, + "mcqa": _mcqa_mod, + "musique": _musique_mod, + "qasper": _qasper_mod, } # Registry mapping framework name → agent module (lazily populated). @@ -69,9 +96,11 @@ def _get_framework_module(framework: str) -> Any: if framework not in _FRAMEWORK_MODULES: if framework == "trace": from agents import trace_agent + _FRAMEWORK_MODULES["trace"] = trace_agent elif framework == "dspy": from agents import dspy_agent + _FRAMEWORK_MODULES["dspy"] = dspy_agent else: raise ValueError( @@ -87,6 +116,7 @@ def _get_framework_module(framework: str) -> Any: try: import datasets as _datasets_lib from datasets import load_dataset as _hf_load_dataset + _datasets_lib.logging.set_verbosity_error() _DATASETS_AVAILABLE = True except ImportError: @@ -124,15 +154,27 @@ def _load_task_config(task_id: str) -> Dict[str, Any]: # Shared data layer # --------------------------------------------------------------------------- -@dataclass -class HFQATask: - """A single QA example with pre-formatted context.""" - question: str - context: str - answer: str +def _task_type_key(cfg: Dict[str, Any]) -> str: + """Return the task semantics key for a YAML task config.""" + task_type = cfg.get("task_type", cfg.get("agent_class", "hfqa")) + if not isinstance(task_type, str) or not task_type: + raise ValueError(f"task_type must be a non-empty string, got {task_type!r}") + return task_type + + +def _get_task_module(task_type: str) -> Any: + """Return the task module registered for a task_type key.""" + try: + return _TASK_MODULES[task_type] + except KeyError as exc: + known = ", ".join(sorted(_TASK_MODULES)) + raise ValueError( + f"Unknown HF task_type {task_type!r}. Known task types: {known}" + ) from exc -def _format_context(raw: Any, context_format: str, agent_class: str) -> str: + +def _format_context(raw: Any, context_format: str, task_type: str) -> str: """Flatten a raw HF context field into a plain string. ``none`` — no context field, returns ``""`` @@ -143,14 +185,94 @@ def _format_context(raw: Any, context_format: str, agent_class: str) -> str: return "" if context_format == "plain": return str(raw) - mod = _TASK_MODULES.get(agent_class) + mod = _TASK_MODULES.get(task_type) if mod is not None: return mod.format_context(raw) return str(raw) +def _resolve_count( + name: str, + override: Optional[int], + cfg: Dict[str, Any], + default: int, +) -> int: + """Resolve a dataset partition size and validate it.""" + value = override if override is not None else cfg.get(name, default) + if not isinstance(value, int) or isinstance(value, bool): + raise TypeError(f"{name} must be a non-negative integer, got {value!r}") + if value < 0: + raise ValueError(f"{name} must be non-negative, got {value}") + return value + + +def _load_hf_partitions( + cfg: Dict[str, Any], + specs: List[Tuple[str, str, int]], +) -> Dict[str, List[HFQATask]]: + """Load named HF partitions without overlap when splits are reused.""" + partitions: Dict[str, List[HFQATask]] = {name: [] for name, _, _ in specs} + by_split: Dict[str, List[Tuple[str, int]]] = {} + + for name, split, count in specs: + if not isinstance(count, int) or isinstance(count, bool): + raise TypeError( + f"Partition {name!r} count must be an integer, got {count!r}" + ) + if count < 0: + raise ValueError( + f"Partition {name!r} count must be non-negative, got {count}" + ) + by_split.setdefault(split, []).append((name, count)) + + for split, split_specs in by_split.items(): + total = sum(count for _, count in split_specs) + loaded = _load_hf_data(cfg, split, total) if total > 0 else [] + start = 0 + for name, count in split_specs: + partitions[name] = loaded[start : start + count] + start += count + + return partitions + + +def _default_row_to_tasks(row: Dict[str, Any], cfg: Dict[str, Any]) -> List[HFQATask]: + """Convert a generic scalar QA row into one HFQATask.""" + q_field = cfg.get("question_field", "question") + a_field = cfg.get("answer_field", "answer") + ctx_fmt = cfg.get("context_format", "plain") + task_type = _task_type_key(cfg) + context_field = cfg.get("context_field") + + if q_field not in row: + raise KeyError(f"Dataset row missing question field {q_field!r}") + if a_field not in row: + raise KeyError(f"Dataset row missing answer field {a_field!r}") + + question = str(row[q_field]) + raw_answer = row[a_field] + if isinstance(raw_answer, dict): + texts = raw_answer.get("text", []) + answer = str(texts[0]) if texts else "" + else: + answer = str(raw_answer) + raw_context = ( + row.get(context_field, "") + if context_field + else row.get("context", row.get("passage", "")) + ) + context = _format_context(raw_context, ctx_fmt, task_type) + return [HFQATask(question=question, context=context, answer=answer)] + + def _load_hf_data(cfg: Dict[str, Any], split: str, n: int) -> List[HFQATask]: """Load up to *n* examples from *split* of the dataset described by *cfg*.""" + if not isinstance(n, int) or isinstance(n, bool): + raise TypeError(f"n must be a non-negative integer, got {n!r}") + if n < 0: + raise ValueError(f"n must be non-negative, got {n}") + if n == 0: + return [] if not _DATASETS_AVAILABLE: raise ImportError( "The `datasets` package is required for hf: tasks. " @@ -161,25 +283,23 @@ def _load_hf_data(cfg: Dict[str, Any], split: str, n: int) -> List[HFQATask]: cfg.get("dataset_config"), split=split, ) - q_field = cfg.get("question_field", "question") - a_field = cfg.get("answer_field", "answer") - ctx_fmt = cfg.get("context_format", "plain") - agent_class = cfg.get("agent_class", "hfqa") + task_type = _task_type_key(cfg) + task_mod = _get_task_module(task_type) + row_to_tasks = getattr(task_mod, "row_to_tasks", None) tasks: List[HFQATask] = [] for row in ds: if len(tasks) >= n: break - question = str(row[q_field]) - raw_answer = row[a_field] - if isinstance(raw_answer, dict): - texts = raw_answer.get("text", []) - answer = texts[0] if texts else "" - else: - answer = str(raw_answer) - raw_context = row.get("context", row.get("passage", "")) - context = _format_context(raw_context, ctx_fmt, agent_class) - tasks.append(HFQATask(question=question, context=context, answer=answer)) + expanded = ( + row_to_tasks(row, cfg) + if callable(row_to_tasks) + else _default_row_to_tasks(row, cfg) + ) + for task in expanded: + tasks.append(task) + if len(tasks) >= n: + break return tasks @@ -188,13 +308,17 @@ def _load_hf_data(cfg: Dict[str, Any], split: str, n: int) -> List[HFQATask]: # Internal helpers # --------------------------------------------------------------------------- -def _make_dataset(tasks: List[HFQATask], agent_class: str) -> Dict[str, List[Any]]: - return _TASK_MODULES[agent_class].make_dataset(tasks) +def _make_dataset_from_cfg( + tasks: List[HFQATask], cfg: Dict[str, Any] +) -> Dict[str, List[Any]]: + """Convert HFQATask objects using the configured task semantics.""" + return _get_task_module(_task_type_key(cfg)).make_dataset(tasks) -def _make_guide(agent_class: str) -> Any: - """Return an instantiated guide for *agent_class* (task-specific, framework-agnostic).""" - return _TASK_MODULES[agent_class].make_guide() + +def _make_guide(cfg: Dict[str, Any]) -> Any: + """Return an instantiated guide for the configured task semantics.""" + return _get_task_module(_task_type_key(cfg)).make_guide() def _make_agent(agent_class: str, framework: str, **kwargs: Any) -> Any: @@ -218,6 +342,7 @@ def build_trace_problem( num_test: Optional[int] = None, initial_instructions: str = _DEFAULT_INSTRUCTIONS, framework: str = "trace", + num_eval_repeats: int = 1, **_kwargs: Any, ) -> Dict[str, Any]: """Build a Trace-Bench task bundle for an HF QA task. @@ -249,19 +374,25 @@ def build_trace_problem( cfg = _load_task_config(task_id) # Resolve dataset sizes: eval_kwargs > hf_tasks.yaml > hardcoded fallback. - num_train = num_train if num_train is not None else cfg.get("num_train", 10) - num_validate = num_validate if num_validate is not None else cfg.get("num_validate", 0) - num_test = num_test if num_test is not None else cfg.get("num_test", 0) + num_train = _resolve_count("num_train", num_train, cfg, 10) + num_validate = _resolve_count("num_validate", num_validate, cfg, 0) + num_test = _resolve_count("num_test", num_test, cfg, 0) if subtask: # BBEH-style multi-split tasks: every subset (train/val/test) comes from # the same split, which is the subtask name (e.g. "boolean_expressions"). split = subtask - total = num_train + num_validate + num_test - all_tasks = _load_hf_data(cfg, split, total) - train_tasks = all_tasks[:num_train] - val_tasks = all_tasks[num_train:num_train + num_validate] - test_tasks = all_tasks[num_train + num_validate:] + train_split_name = split + validate_split_name = split + test_split_name = split + partitions = _load_hf_partitions( + cfg, + [ + ("train", split, num_train), + ("validate", split, num_validate), + ("test", split, num_test), + ], + ) else: # Tasks with explicit train/test HF splits (e.g. HotpotQA): # train + val → loaded from `train_split` (default: "train") @@ -269,16 +400,30 @@ def build_trace_problem( # Falls back to the legacy `split` field if neither is set. fallback = cfg.get("split", "train") train_split = cfg.get("train_split", fallback) - test_split = cfg.get("test_split", train_split) + test_split = cfg.get("test_split", train_split) + train_split_name = train_split + validate_split_name = train_split + test_split_name = test_split + partitions = _load_hf_partitions( + cfg, + [ + ("train", train_split, num_train), + ("validate", train_split, num_validate), + ("test", test_split, num_test), + ], + ) - train_val_tasks = _load_hf_data(cfg, train_split, num_train + num_validate) - train_tasks = train_val_tasks[:num_train] - val_tasks = train_val_tasks[num_train:] - test_tasks = _load_hf_data(cfg, test_split, num_test) if num_test > 0 else [] + train_tasks = partitions["train"] + val_tasks = partitions["validate"] + test_tasks = partitions["test"] - objective = cfg.get("objective", "Optimize the agent's instructions for this QA task.") + objective = cfg.get( + "objective", "Optimize the agent's instructions for this QA task." + ) agent_class = cfg.get("agent_class", "hfqa") + if initial_instructions == _DEFAULT_INSTRUCTIONS: + initial_instructions = cfg.get("initial_instructions", _DEFAULT_INSTRUCTIONS) agent_kwargs: Dict[str, Any] = dict( initial_instructions=initial_instructions, node_description=cfg.get( @@ -288,27 +433,32 @@ def build_trace_problem( ) agent = _make_agent(agent_class, framework, **agent_kwargs) - guide = _make_guide(agent_class) + guide = _make_guide(cfg) bundle: Dict[str, Any] = dict( param=agent, guide=guide, - train_dataset=_make_dataset(train_tasks, agent_class), + train_dataset=_make_dataset_from_cfg(train_tasks, cfg), optimizer_kwargs=dict(objective=objective, memory_size=10), metadata=dict( benchmark="hf", task_id=task_id, subtask=subtask, dataset_id=cfg["dataset_id"], + task_type=_task_type_key(cfg), dataset_config=cfg.get("dataset_config"), num_train=num_train, num_validate=num_validate, num_test=num_test, + train_split=train_split_name, + validate_split=validate_split_name, + test_split=test_split_name, framework=framework, ), ) if val_tasks: - bundle["validate_dataset"] = _make_dataset(val_tasks, agent_class) + bundle["validate_dataset"] = _make_dataset_from_cfg(val_tasks, cfg) if test_tasks: - bundle["test_dataset"] = _make_dataset(test_tasks, agent_class) + bundle["test_dataset"] = _make_dataset_from_cfg(test_tasks, cfg) + bundle["num_eval_repeats"] = max(1, int(num_eval_repeats)) return bundle diff --git a/benchmarks/hf_qa/hf_tasks.yaml b/benchmarks/hf_qa/hf_tasks.yaml index e8c549e..e535830 100644 --- a/benchmarks/hf_qa/hf_tasks.yaml +++ b/benchmarks/hf_qa/hf_tasks.yaml @@ -1,7 +1,8 @@ # Curated HuggingFace QA tasks for the `hf:` suite. # # Each entry is loaded by benchmarks/hf_qa/hf_qa_loader.py. -# To add a new task, append an entry here — no Python changes needed. +# To add a new task with existing semantics, append an entry here. +# New task semantics require a matching module in this directory. # # HotpotQA is a large dataset with 7,405 samples in their validation split. # We hard-code the train/validation/test splits to avoid loading the entire dataset. @@ -20,7 +21,14 @@ # Optional fields: # dataset_config : HuggingFace config name (e.g. "distractor", "plain_text") # train_split : HF split used for train + val examples (default: "train") -# test_split : HF split used for test examples (default: same as train_split) +# task_type : task semantics module key (data shaping + guide). +# Defaults to `agent_class` for backward compatibility. +# This decouples task semantics from agent implementation family. +# context_field : override which raw field to use as context +# options_field : MCQ option field for task_type: mcqa +# answer_index_base: MCQ integer answer base, usually 0 or 1 +# test_split : HF split used for test examples (default: same as train_split; +# when the same split is reused, partitions are disjoint slices) # split : fallback split when train_split/test_split are not set # description : short human-readable description # objective : optimizer objective injected into optimizer_kwargs @@ -131,3 +139,148 @@ - movie_recommendation - boardgame_qa +# --------------------------------------------------------------------------- +# MuSiQue — compositional multi-hop QA +# --------------------------------------------------------------------------- +- id: musique + dataset_id: dgslibisey/MuSiQue + train_split: train + test_split: validation + num_train: 128 + num_validate: 128 + num_test: 128 + question_field: question + answer_field: answer + context_field: paragraphs + context_format: musique + task_type: musique + agent_class: hfqa + description: "MuSiQue compositional multi-hop QA" + objective: | + Optimize the agent's instructions for connected multi-hop reasoning across multiple paragraphs. + +# --------------------------------------------------------------------------- +# GSM8K — math word problems +# --------------------------------------------------------------------------- +- id: gsm8k + dataset_id: openai/gsm8k + dataset_config: main + train_split: train + test_split: test + num_train: 128 + num_validate: 128 + num_test: 128 + question_field: question + answer_field: answer + context_format: none + task_type: gsm8k + agent_class: bbeh + description: "GSM8K arithmetic word problems" + objective: | + Optimize the prompt template and extraction logic for free-form numeric reasoning. + +# --------------------------------------------------------------------------- +# ARC-Challenge — science MCQ +# --------------------------------------------------------------------------- +- id: arc_challenge + dataset_id: allenai/ai2_arc + dataset_config: ARC-Challenge + train_split: train + test_split: test + num_train: 128 + num_validate: 128 + num_test: 128 + question_field: question + answer_field: answerKey + options_field: choices + context_format: none + task_type: mcqa + agent_class: bbeh + description: "ARC-Challenge science multiple-choice reasoning" + objective: | + Optimize the prompt template and extraction logic for science multiple-choice reasoning. + +# --------------------------------------------------------------------------- +# QASC — multi-hop science MCQ +# --------------------------------------------------------------------------- +- id: qasc + dataset_id: allenai/qasc + train_split: train + test_split: test + num_train: 128 + num_validate: 128 + num_test: 128 + question_field: question + answer_field: answerKey + options_field: choices + context_field: combinedfact + context_format: plain + task_type: mcqa + agent_class: bbeh + description: "QASC fact-composition science multiple-choice QA" + objective: | + Optimize the prompt template and extraction logic for fact-composition science QA. + +# --------------------------------------------------------------------------- +# DROP — passage-based discrete / numerical reading comprehension +# --------------------------------------------------------------------------- +- id: drop + dataset_id: ucinlp/drop + train_split: train + test_split: validation + num_train: 128 + num_validate: 128 + num_test: 128 + question_field: question + answer_field: answers_spans + context_field: passage + context_format: plain + task_type: drop + agent_class: hfqa + description: "DROP discrete / numerical reasoning over passages" + objective: | + Optimize instructions for passage-based arithmetic and discrete reasoning. + +# --------------------------------------------------------------------------- +# QuALITY — long-context MCQ reading comprehension +# --------------------------------------------------------------------------- +- id: quality + dataset_id: emozilla/quality + train_split: train + test_split: validation + num_train: 64 + num_validate: 64 + num_test: 64 + question_field: question + answer_field: answer + answer_index_base: 1 + options_field: options + context_field: article + context_format: plain + task_type: mcqa + agent_class: bbeh + description: "QuALITY long-context multiple-choice reading comprehension" + objective: | + Optimize the prompt template and extraction logic for long-context multiple-choice reasoning. + +# --------------------------------------------------------------------------- +# Qasper — scholarly long-document QA +# +# Uses a script-free converted mirror because `allenai/qasper` is script-backed +# and no longer loads in current `datasets` versions. +# --------------------------------------------------------------------------- +- id: qasper + dataset_id: abertsch/converted_qasper + train_split: train + test_split: validation + num_train: 128 + num_validate: 128 + num_test: 128 + question_field: input + answer_field: output + context_format: plain + task_type: qasper + agent_class: hfqa + description: "Qasper scholarly long-document QA" + objective: | + Optimize instructions for information-seeking QA over full research papers. diff --git a/benchmarks/hf_qa/hotpot_qa.py b/benchmarks/hf_qa/hotpot_qa.py index e127184..47ca676 100644 --- a/benchmarks/hf_qa/hotpot_qa.py +++ b/benchmarks/hf_qa/hotpot_qa.py @@ -6,14 +6,21 @@ Reference: https://github.com/xuanfeiren/hotpotqa/blob/main/prompt_opt/trace_opt.py """ + from __future__ import annotations import re import string -from typing import Any +from typing import Any, Tuple + +try: + from .common import response_text +except ImportError: + from common import response_text try: from opto.trainer.guide import Guide + _OPTO_AVAILABLE = True except Exception: _OPTO_AVAILABLE = False @@ -23,6 +30,7 @@ # Data utilities # --------------------------------------------------------------------------- + def _normalize_answer(text: str) -> str: """Lowercase, remove articles and punctuation, collapse whitespace.""" text = text.lower() @@ -61,14 +69,17 @@ def check_answer(response: str, expected: str) -> bool: # --------------------------------------------------------------------------- if _OPTO_AVAILABLE: + class HotpotQAGuide(Guide): """Evaluation guide for HotpotQA and other context-rich QA tasks. Framework-agnostic: works with any agent that produces a string answer. """ - def get_feedback(self, task: Any, response: Any, info: Any, **kwargs: Any): - response_str = str(getattr(response, "data", response)) + def get_feedback( + self, task: Any, response: Any, info: Any, **kwargs: Any + ) -> Tuple[float, str]: + response_str = response_text(response) if check_answer(response_str, info.answer): return 1.0, "Correct." return 0.0, ( @@ -79,6 +90,7 @@ def get_feedback(self, task: Any, response: Any, info: Any, **kwargs: Any): ) else: + class HotpotQAGuide: # type: ignore[no-redef] pass diff --git a/benchmarks/hf_qa/mcqa.py b/benchmarks/hf_qa/mcqa.py new file mode 100644 index 0000000..6d772ec --- /dev/null +++ b/benchmarks/hf_qa/mcqa.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Tuple + +try: + from .common import ( + HFQATask, + answer_to_choice_label, + choice_pairs, + extract_last_choice, + format_mcq_prompt, + normalize_text, + response_text, + ) +except ImportError: + from common import ( + HFQATask, + answer_to_choice_label, + choice_pairs, + extract_last_choice, + format_mcq_prompt, + normalize_text, + response_text, + ) + +try: + from opto.trainer.guide import Guide +except Exception: + Guide = object # type: ignore[assignment,misc] + + +def row_to_tasks(row: Dict[str, Any], cfg: Dict[str, Any]) -> List[HFQATask]: + """Convert one multiple-choice row into a prompt-based QA task.""" + question_field = cfg.get("question_field", "question") + answer_field = cfg.get("answer_field", "answer") + options_field = cfg.get("options_field", "options") + context_field = cfg.get("context_field") + if question_field not in row: + raise KeyError(f"MCQ row missing question field {question_field!r}") + if answer_field not in row: + raise KeyError(f"MCQ row missing answer field {answer_field!r}") + + pairs = choice_pairs(row.get(options_field, row.get("choices"))) + if not pairs: + raise ValueError(f"MCQ row has no options in field {options_field!r}") + + context = str(row.get(context_field, "")) if context_field else "" + answer = row[answer_field] + label = answer_to_choice_label( + answer, pairs, index_base=int(cfg.get("answer_index_base", 0)) + ) + gold_text = next( + (text for option_label, text in pairs if option_label == label), "" + ) + prompt = format_mcq_prompt(str(row[question_field]), pairs, context=context) + return [ + HFQATask( + question=prompt, + context=context, + answer=label, + meta={"gold_label": label, "gold_text": gold_text, "options": list(pairs)}, + ) + ] + + +def format_context(raw: Any) -> str: + """Return plain-string context for MCQ tasks.""" + return str(raw or "") + + +def make_dataset(tasks: List[HFQATask]) -> Dict[str, List[Any]]: + """Return the question-only dataset shape used by BBEH-style agents.""" + return {"inputs": [task.question for task in tasks], "infos": list(tasks)} + + +class MCQGuide(Guide): # type: ignore[misc] + """Guide for multiple-choice label or option-text scoring.""" + + def get_feedback( + self, task: Any, response: Any, info: HFQATask, **kwargs: Any + ) -> Tuple[float, str]: + """Score a response against the expected option label or text.""" + response_str = response_text(response) + pred_label = extract_last_choice(response_str) + gold_label = str(info.meta.get("gold_label", info.answer)) + gold_text = str(info.meta.get("gold_text", "")) + + if pred_label and pred_label.strip().lower() == gold_label.strip().lower(): + return 1.0, "The selected option is correct." + if gold_text and normalize_text(gold_text) in normalize_text(response_str): + return 1.0, "The selected option text is correct." + return 0.0, ( + f'Incorrect. Expected option "{gold_label}" ({gold_text}), got "{response_str}". ' + "Improve option-sensitive reasoning and produce the final answer in the expected MCQ format." + ) + + +def make_guide() -> Any: + """Return an instantiated MCQ guide.""" + return MCQGuide() diff --git a/benchmarks/hf_qa/musique.py b/benchmarks/hf_qa/musique.py new file mode 100644 index 0000000..d24a330 --- /dev/null +++ b/benchmarks/hf_qa/musique.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Tuple + +try: + from .common import HFQATask, best_alias_score, normalize_aliases, response_text +except ImportError: + from common import HFQATask, best_alias_score, normalize_aliases, response_text + +try: + from opto.trainer.guide import Guide +except Exception: + Guide = object # type: ignore[assignment,misc] + + +def format_context(raw: Any) -> str: + """Format MuSiQue paragraph rows into numbered context blocks.""" + parts: List[str] = [] + for index, paragraph in enumerate(raw or []): + title = paragraph.get("title", f"Paragraph {index + 1}") + text = paragraph.get("paragraph_text", paragraph.get("paragraph", "")) + parts.append(f"Paragraph {index + 1} ({title}):\n{text}") + return "\n\n".join(parts) + + +def row_to_tasks(row: Dict[str, Any], cfg: Dict[str, Any]) -> List[HFQATask]: + """Convert one MuSiQue row into a single multi-hop QA task.""" + question_field = cfg.get("question_field", "question") + context_field = cfg.get("context_field", "paragraphs") + if question_field not in row: + raise KeyError(f"MuSiQue row missing question field {question_field!r}") + + aliases = normalize_aliases( + [str(row.get("answer", ""))] + + [str(alias) for alias in (row.get("answer_aliases", []) or [])] + ) + return [ + HFQATask( + question=str(row[question_field]), + context=format_context(row.get(context_field, [])), + answer=aliases[0] if aliases else "", + meta={"aliases": aliases}, + ) + ] + + +def make_dataset(tasks: List[HFQATask]) -> Dict[str, List[Any]]: + """Return the default context-rich QA dataset shape.""" + return {"inputs": list(tasks), "infos": list(tasks)} + + +class MuSiQueGuide(Guide): # type: ignore[misc] + """Guide for alias-tolerant MuSiQue answer scoring.""" + + def get_feedback( + self, task: Any, response: Any, info: HFQATask, **kwargs: Any + ) -> Tuple[float, str]: + """Score a MuSiQue response against the gold aliases.""" + response_str = response_text(response) + aliases = info.meta.get("aliases", [info.answer]) + score = best_alias_score(response_str, aliases) + if score >= 0.999: + return 1.0, "Correct." + return score, ( + f"Incorrect. Gold aliases: {aliases}. Got: '{response_str}'. " + "Improve compositional multi-hop reasoning across the provided paragraphs." + ) + + +def make_guide() -> Any: + """Return an instantiated MuSiQue guide.""" + return MuSiQueGuide() diff --git a/benchmarks/hf_qa/qasper.py b/benchmarks/hf_qa/qasper.py new file mode 100644 index 0000000..b2f8722 --- /dev/null +++ b/benchmarks/hf_qa/qasper.py @@ -0,0 +1,172 @@ +from __future__ import annotations + +from typing import Any, Dict, Iterable, List, Tuple + +try: + from .common import HFQATask, best_alias_score, normalize_aliases, response_text +except ImportError: + from common import HFQATask, best_alias_score, normalize_aliases, response_text + +try: + from opto.trainer.guide import Guide +except Exception: + Guide = object # type: ignore[assignment,misc] + + +def format_context(raw: Any) -> str: + """Format Qasper full_text sections into a long-document context.""" + if not isinstance(raw, dict): + return str(raw or "") + + section_names = raw.get("section_name", []) + paragraphs = raw.get("paragraphs", []) + parts: List[str] = [] + for index, section_paragraphs in enumerate(paragraphs): + title = ( + section_names[index] + if index < len(section_names) + else f"Section {index + 1}" + ) + joined = "\n".join( + str(paragraph) for paragraph in section_paragraphs if str(paragraph).strip() + ) + if joined: + parts.append(f"{title}\n{joined}") + return "\n\n".join(parts) + + +def _iter_qasper_answers(raw: Any) -> Iterable[Dict[str, Any]]: + """Yield normalized answer dicts from Qasper's nested answer shape.""" + if isinstance(raw, list): + for item in raw: + yield from _iter_qasper_answers(item) + elif isinstance(raw, dict): + if "answer" in raw: + yield from _iter_qasper_answers(raw["answer"]) + else: + yield raw + + +def _aliases_and_evidence(raw_answers: Any) -> Tuple[List[str], List[str]]: + """Extract answer aliases and evidence snippets from Qasper answers.""" + aliases: List[str] = [] + evidence: List[str] = [] + for answer in _iter_qasper_answers(raw_answers): + if answer.get("unanswerable"): + aliases.append("unanswerable") + + spans = answer.get("extractive_spans", []) or [] + aliases.extend(str(span) for span in spans if str(span).strip()) + + free_form = str(answer.get("free_form_answer", "")).strip() + if free_form: + aliases.append(free_form) + + yes_no = answer.get("yes_no") + if isinstance(yes_no, bool): + aliases.append("yes" if yes_no else "no") + + evidence.extend( + str(item) + for item in (answer.get("evidence", []) or []) + if str(item).strip() + ) + + return normalize_aliases(aliases), evidence + + +def _converted_qasper_task(row: Dict[str, Any], cfg: Dict[str, Any]) -> HFQATask: + """Convert one row from script-free converted_qasper mirrors.""" + question_field = cfg.get("question_field", "input") + answer_field = cfg.get("answer_field", "output") + if question_field not in row: + raise KeyError(f"Qasper row missing question field {question_field!r}") + if answer_field not in row: + raise KeyError(f"Qasper row missing answer field {answer_field!r}") + + raw_input = str(row[question_field]) + question = raw_input + context = "" + if raw_input.startswith("Q:") and "\nText:" in raw_input: + question_part, context_part = raw_input.split("\nText:", 1) + question = question_part.removeprefix("Q:").strip() + context = context_part.strip() + + aliases = normalize_aliases([str(row[answer_field])]) + return HFQATask( + question=question, + context=context, + answer=aliases[0] if aliases else "", + meta={ + "aliases": aliases, + "question_id": str(row.get("pid", row.get("id", ""))), + "evidence": [], + }, + ) + + +def row_to_tasks(row: Dict[str, Any], cfg: Dict[str, Any]) -> List[HFQATask]: + """Expand one Qasper paper row into one task per question.""" + if "qas" not in row: + return [_converted_qasper_task(row, cfg)] + + context = format_context(row.get(cfg.get("context_field", "full_text"), {})) + qas = row.get("qas", {}) + if not isinstance(qas, dict): + raise TypeError( + f"Qasper row field 'qas' must be a dict, got {type(qas).__name__}" + ) + + questions = qas.get("question", []) + question_ids = qas.get("question_id", [""] * len(questions)) + answers = qas.get("answers", []) + + tasks: List[HFQATask] = [] + for index, question in enumerate(questions): + raw_answers = answers[index] if index < len(answers) else [] + aliases, evidence = _aliases_and_evidence(raw_answers) + tasks.append( + HFQATask( + question=str(question), + context=context, + answer=aliases[0] if aliases else "", + meta={ + "aliases": aliases, + "question_id": ( + question_ids[index] if index < len(question_ids) else "" + ), + "evidence": evidence, + }, + ) + ) + return tasks + + +def make_dataset(tasks: List[HFQATask]) -> Dict[str, List[Any]]: + """Return the default context-rich QA dataset shape.""" + return {"inputs": list(tasks), "infos": list(tasks)} + + +class QasperGuide(Guide): # type: ignore[misc] + """Guide for alias-tolerant Qasper answer scoring.""" + + def get_feedback( + self, task: Any, response: Any, info: HFQATask, **kwargs: Any + ) -> Tuple[float, str]: + """Score a Qasper response against gold aliases and evidence.""" + response_str = response_text(response) + aliases = info.meta.get("aliases", [info.answer]) + score = best_alias_score(response_str, aliases) + if score >= 0.999: + return 1.0, "Correct." + evidence = info.meta.get("evidence", [])[:2] + evidence_hint = f" Evidence hints: {evidence}." if evidence else "" + return score, ( + f"Incorrect. Expected one of {aliases}, got '{response_str}'." + f"{evidence_hint} Improve long-document scholarly QA and answer grounding." + ) + + +def make_guide() -> Any: + """Return an instantiated Qasper guide.""" + return QasperGuide() diff --git a/tests/m1/test_hf_qa_ext.py b/tests/m1/test_hf_qa_ext.py new file mode 100644 index 0000000..fb69004 --- /dev/null +++ b/tests/m1/test_hf_qa_ext.py @@ -0,0 +1,234 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +import benchmarks.hf_qa.hf_qa_loader as loader +from benchmarks.hf_qa import drop_qa, gsm8k, mcqa, musique, qasper +from benchmarks.hf_qa.common import HFQATask, response_text + + +def test_yaml_entries_exist() -> None: + """New HF task ids resolve to the expected task semantics.""" + expected = { + "musique": "musique", + "gsm8k": "gsm8k", + "arc_challenge": "mcqa", + "qasc": "mcqa", + "drop": "drop", + "quality": "mcqa", + "qasper": "qasper", + } + for task_id, task_type in expected.items(): + assert loader._load_task_config(task_id)["task_type"] == task_type + + +def test_partition_loader_same_split_nonoverlap( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Partitions sharing a split use disjoint contiguous slices.""" + calls = [] + + def fake_load(_cfg: dict, split: str, n: int) -> list[HFQATask]: + calls.append((split, n)) + return [ + HFQATask(question=f"{split}-{index}", context="", answer=str(index)) + for index in range(n) + ] + + monkeypatch.setattr(loader, "_load_hf_data", fake_load) + parts = loader._load_hf_partitions( + {}, + [("train", "train", 2), ("validate", "train", 2), ("test", "train", 1)], + ) + + assert calls == [("train", 5)] + assert [task.answer for task in parts["train"]] == ["0", "1"] + assert [task.answer for task in parts["validate"]] == ["2", "3"] + assert [task.answer for task in parts["test"]] == ["4"] + + +def test_resolve_count_rejects_invalid_values() -> None: + """Dataset sizes fail fast on unsafe values.""" + with pytest.raises(ValueError, match="non-negative"): + loader._resolve_count("num_train", -1, {}, 10) + with pytest.raises(TypeError, match="integer"): + loader._resolve_count("num_train", True, {}, 10) + + +def test_response_text_extracts_chat_message_content() -> None: + """LLM scoring uses assistant content rather than response object repr.""" + response = SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content="final answer"))] + ) + + assert response_text(response) == "final answer" + + +def test_gsm8k_guide_numeric_match_and_miss() -> None: + """GSM8K scoring extracts the final numeric answer.""" + row = {"question": "How many?", "answer": "Natalia sold 48+24 = 72. #### 72"} + task = gsm8k.row_to_tasks( + row, {"question_field": "question", "answer_field": "answer"} + )[0] + guide = gsm8k.make_guide() + + score, _ = guide.get_feedback(task.question, "Reasoning... #### 72", task) + assert score == 1.0 + score, _ = guide.get_feedback(task.question, "Reasoning... #### 71", task) + assert score == 0.0 + + +def test_mcqa_arc_like_row() -> None: + """MCQ rows with labeled choice dicts preserve labels and option text.""" + row = { + "question": "Which factor increased?", + "choices": {"label": ["A", "B"], "text": ["Shady areas", "Food sources"]}, + "answerKey": "B", + "combinedfact": "Food can increase populations.", + } + task = mcqa.row_to_tasks( + row, + { + "question_field": "question", + "options_field": "choices", + "answer_field": "answerKey", + "context_field": "combinedfact", + }, + )[0] + + assert "(A) Shady areas" in task.question + assert task.meta["gold_label"] == "B" + guide = mcqa.make_guide() + score, _ = guide.get_feedback(task.question, "I choose (B).", task) + assert score == 1.0 + + +def test_mcqa_quality_one_based_answer() -> None: + """QuALITY-style one-based answers map to A/B/C/D labels.""" + row = { + "question": "Why does he retire?", + "options": [ + "For honor", + "Because he has enough money", + "Because the ship broke", + "To become mayor", + ], + "answer": 2, + "article": "Long story...", + } + task = mcqa.row_to_tasks( + row, + { + "question_field": "question", + "options_field": "options", + "answer_field": "answer", + "answer_index_base": 1, + "context_field": "article", + }, + )[0] + + assert task.meta["gold_label"] == "B" + + +def test_mcqa_rejects_missing_options() -> None: + """MCQ conversion reports malformed rows instead of silently continuing.""" + with pytest.raises(ValueError, match="no options"): + mcqa.row_to_tasks({"question": "Q?", "answer": "A"}, {}) + + +def test_drop_aliases_and_guide() -> None: + """DROP scoring accepts aliases embedded in a sentence.""" + row = { + "question": "Who scored first?", + "passage": "Some passage", + "answers_spans": {"spans": ["Chaz Schilens", "Schilens"]}, + } + task = drop_qa.row_to_tasks( + row, + { + "question_field": "question", + "context_field": "passage", + "answer_field": "answers_spans", + }, + )[0] + guide = drop_qa.make_guide() + + score, _ = guide.get_feedback(task, "The first scorer was Chaz Schilens.", task) + assert score == 1.0 + + +def test_musique_alias_support() -> None: + """MuSiQue conversion keeps answer aliases for scoring.""" + row = { + "question": "Where did he die?", + "paragraphs": [{"title": "Doc", "paragraph_text": "He died in Copenhagen."}], + "answer": "Copenhagen", + "answer_aliases": ["Copenhagen", "Kobenhavn"], + } + task = musique.row_to_tasks( + row, {"question_field": "question", "context_field": "paragraphs"} + )[0] + + assert "Copenhagen" in task.meta["aliases"] + + +def test_qasper_row_expansion() -> None: + """Original nested Qasper rows expand to one task per question.""" + row = { + "full_text": {"section_name": ["Intro"], "paragraphs": [["A long paragraph."]]}, + "qas": { + "question": ["What is proposed?", "Is it supervised?"], + "question_id": ["q1", "q2"], + "answers": [ + [ + { + "answer": [ + { + "free_form_answer": "A new method", + "extractive_spans": [], + "evidence": ["A long paragraph."], + } + ] + } + ], + [ + { + "answer": [ + { + "free_form_answer": "", + "extractive_spans": [], + "yes_no": True, + "evidence": ["A long paragraph."], + } + ] + } + ], + ], + }, + } + + tasks = qasper.row_to_tasks(row, {}) + + assert len(tasks) == 2 + assert tasks[0].meta["question_id"] == "q1" + assert "A new method" in tasks[0].meta["aliases"] + assert "yes" in [alias.lower() for alias in tasks[1].meta["aliases"]] + + +def test_qasper_converted_row() -> None: + """Converted Qasper mirror rows split question and paper text.""" + row = { + "id": "paper", + "pid": "paper_0", + "input": "Q: What is proposed?\nText: Intro\nA long paragraph.", + "output": "A new method", + } + task = qasper.row_to_tasks( + row, {"question_field": "input", "answer_field": "output"} + )[0] + + assert task.question == "What is proposed?" + assert task.context.startswith("Intro") + assert task.answer == "A new method" diff --git a/tests/m2/test_noop_trainer.py b/tests/m2/test_noop_trainer.py new file mode 100644 index 0000000..fc6bffe --- /dev/null +++ b/tests/m2/test_noop_trainer.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +import shutil +from pathlib import Path +from uuid import uuid4 + +from trace_bench.config import RunConfig +from trace_bench.runner import BenchRunner + + +def test_noop_trainer_runs_through_bench_runner() -> None: + """NoOpTrainer should work when instantiated by the Trace-Bench runner.""" + repo_root = Path(__file__).resolve().parents[2] + run_root = repo_root / "pytest_tmp_global" / f"noop-{uuid4().hex}" / "runs" + cfg = RunConfig.from_dict( + { + "runs_dir": str(run_root), + "mode": "stub", + "resume": "none", + "tasks": [{"id": "trace_examples:greeting_stub"}], + "trainers": [{"id": "NoOpTrainer", "params_variants": [{"num_steps": 1}]}], + "seeds": [123], + } + ) + + try: + summary = BenchRunner(cfg).run(force=True) + assert len(summary.results) == 1 + assert summary.results[0]["status"] == "ok" + finally: + shutil.rmtree(run_root.parent, ignore_errors=True) diff --git a/tests/m3/test_llm_metadata.py b/tests/m3/test_llm_metadata.py index 2ab7fa0..4649cde 100644 --- a/tests/m3/test_llm_metadata.py +++ b/tests/m3/test_llm_metadata.py @@ -1,6 +1,15 @@ from __future__ import annotations -from trace_bench.runner import _current_llm_meta, _extract_token_usage +import copy +from typing import Any, Dict + +from trace_bench.runner import ( + _MaxCompletionTokensLLM, + _current_llm_meta, + _extract_token_usage, + _make_optimizer_litellm, + _uses_max_completion_tokens, +) def test_current_llm_meta_detects_openrouter(monkeypatch): @@ -12,7 +21,66 @@ def test_current_llm_meta_detects_openrouter(monkeypatch): def test_extract_token_usage_from_dict(): - usage = _extract_token_usage({"usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14}}) + usage = _extract_token_usage( + {"usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14}} + ) assert usage["prompt_tokens"] == 10 assert usage["completion_tokens"] == 4 assert usage["total_tokens"] == 14 + + +def test_uses_max_completion_tokens_for_gpt5_and_reasoning_models(): + assert _uses_max_completion_tokens("gpt-5.4-nano") + assert _uses_max_completion_tokens("openai/o3-mini") + assert not _uses_max_completion_tokens("gpt-4o-mini") + + +def test_max_completion_tokens_llm_translates_keyword(): + calls = [] + + def fake_llm(**kwargs: Any) -> str: + calls.append(kwargs) + return "ok" + + wrapped = _MaxCompletionTokensLLM(fake_llm) + + assert wrapped(max_tokens=7, response_format={"type": "json_object"}) == "ok" + assert calls == [ + {"max_completion_tokens": 7, "response_format": {"type": "json_object"}} + ] + + +def test_max_completion_tokens_llm_deepcopy_keeps_delegate() -> None: + class CallableLLM: + def __call__(self, **kwargs: Any) -> Dict[str, Any]: + return kwargs + + copied = copy.deepcopy(_MaxCompletionTokensLLM(CallableLLM())) + + assert copied(max_tokens=5) == {"max_completion_tokens": 5} + + +def test_make_optimizer_litellm_falls_back_without_mm_beta() -> None: + class OldLiteLLM: + def __init__(self, model: str) -> None: + self.model = model + + def __call__(self, **kwargs: Any) -> Dict[str, Any]: + return kwargs + + llm = _make_optimizer_litellm(OldLiteLLM, "gpt-5.4-nano") + + assert isinstance(llm, _MaxCompletionTokensLLM) + assert llm(max_tokens=3) == {"max_completion_tokens": 3} + + +def test_make_optimizer_litellm_uses_mm_beta_when_supported() -> None: + class NewLiteLLM: + def __init__(self, model: str, mm_beta: bool = True) -> None: + self.model = model + self.mm_beta = mm_beta + + llm = _make_optimizer_litellm(NewLiteLLM, "gpt-4o-mini") + + assert llm.model == "gpt-4o-mini" + assert llm.mm_beta is False diff --git a/trace_bench/runner.py b/trace_bench/runner.py index c5bc838..e1a19bb 100644 --- a/trace_bench/runner.py +++ b/trace_bench/runner.py @@ -234,6 +234,52 @@ def _current_llm_meta() -> Dict[str, str]: "llm_base_url": base, } +_MAX_COMPLETION_TOKEN_MODEL_PREFIXES = ("gpt-5", "o1", "o3", "o4") + + +def _uses_max_completion_tokens(model: str) -> bool: + """Return True when an OpenAI model rejects the legacy max_tokens field.""" + normalized = model.split("/", 1)[1] if model.startswith("openai/") else model + normalized = normalized.lower() + return normalized.startswith(_MAX_COMPLETION_TOKEN_MODEL_PREFIXES) + + +class _MaxCompletionTokensLLM: + """Adapter that maps max_tokens to max_completion_tokens for newer OpenAI models.""" + + def __init__(self, llm: Any) -> None: + self._llm = llm + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + """Forward calls after translating the token budget keyword.""" + if "max_tokens" in kwargs and "max_completion_tokens" not in kwargs: + kwargs = dict(kwargs) + kwargs["max_completion_tokens"] = kwargs.pop("max_tokens") + return self._llm(*args, **kwargs) + + def __getattr__(self, name: str) -> Any: + """Delegate metadata and usage attributes to the wrapped LLM.""" + if name.startswith("__"): + raise AttributeError(name) + try: + llm = object.__getattribute__(self, "_llm") + except AttributeError as exc: + raise AttributeError(name) from exc + return getattr(llm, name) + + +def _make_optimizer_litellm(litellm_cls: Any, model: str) -> Any: + """Build the LiteLLM object passed into OpenTrace optimizers.""" + try: + llm = litellm_cls(model=model, mm_beta=False) + except TypeError as exc: + if "mm_beta" not in str(exc): + raise + llm = litellm_cls(model=model) + if _uses_max_completion_tokens(model): + return _MaxCompletionTokensLLM(llm) + return llm + def _snapshot_model_state(model: Any) -> Dict[str, Any]: def _node_state(node: Any, idx: int | None = None) -> Dict[str, Any]: @@ -501,7 +547,7 @@ def _dummy_response(*_args, **_kwargs): # LiteLLM needs "gemini/" prefix to call Gemini via its SDK if litellm_model and "/" not in litellm_model and "gemini" in litellm_model.lower(): litellm_model = f"gemini/{litellm_model}" - opt_llm = LiteLLM(model=litellm_model, mm_beta=False) + opt_llm = _make_optimizer_litellm(LiteLLM, litellm_model) if isinstance(optimizer_kwargs, list): for item in optimizer_kwargs: item.setdefault("llm", opt_llm) diff --git a/trace_bench/trainers/noop_trainer.py b/trace_bench/trainers/noop_trainer.py index 35b88ec..a47f868 100644 --- a/trace_bench/trainers/noop_trainer.py +++ b/trace_bench/trainers/noop_trainer.py @@ -15,6 +15,7 @@ trainers: - id: NoOpTrainer """ + from __future__ import annotations from typing import Any, Dict, List @@ -48,6 +49,22 @@ class NoOpTrainer(_TrainerBase): EXTERNAL_REQUIRES: List[str] = [] # no extra dependencies needed + def __init__( + self, + agent: Any, + optimizer: Any = None, + logger: Any = None, + **kwargs: Any, + ) -> None: + """Store the agent and ignore the optimizer for no-op training.""" + try: + super().__init__(agent, logger=logger, **kwargs) + except TypeError: + super().__init__() + self.logger = logger + self.param = agent + self.optimizer = optimizer + def train( self, guide: Any, @@ -81,6 +98,9 @@ def train( # e.g. mylib.step(param=self.param, score=score, feedback=_feedback) if hasattr(self, "logger") and self.logger is not None: - self.logger.log({"step": step, "score": score}) + try: + self.logger.log("NoOpTrainer/score", score, step) + except TypeError: + self.logger.log({"step": step, "score": score}) return {"status": "ok", "scores": scores}