From fd6aed42c95b09a33c3ac515afbb6c21a8c30aa0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=85=B1=E7=89=9B=E8=82=89?= Date: Fri, 7 Aug 2026 02:06:06 +0800 Subject: [PATCH] feat(meta-reme): add dataset preparation layer with YAML-driven workspace init - Add meta-reme/run.py entry point loading a single --config YAML and preparing the workspace via the data_preparation module. - Add meta-reme/data_preparation/ (basic, lme, beam) normalizing BEAM and LongMemEval training cases into the shared CaseSpec/SessionSpec/QuerySpec contracts with reproducibility manifests and read-only case files. - Add per-benchmark config_meta_reme.yaml and run_meta.py launchers for beam and longmemeval that forward to the unified entry point. - Add tests/unit_sandbox/test_meta_reme_run.py covering case selection, read-only writes, empty-dir init, BEAM variant/time-anchor parsing, and config loading. - Update meta-reme/design.md to reflect the single-config CLI and the renamed data_preparation/ structure. --- benchmark/beam/config_meta_reme.yaml | 8 ++ benchmark/beam/run_meta.py | 30 +++++ benchmark/longmemeval/config_meta_reme.yaml | 7 + benchmark/longmemeval/run_meta.py | 30 +++++ meta-reme/data_preparation/__init__.py | 36 +++++ meta-reme/data_preparation/basic.py | 60 +++++++++ meta-reme/data_preparation/beam.py | 114 ++++++++++++++++ meta-reme/data_preparation/lme.py | 66 ++++++++++ meta-reme/design.md | 10 +- meta-reme/run.py | 133 +++++++++++++++++++ tests/unit_sandbox/test_meta_reme_run.py | 139 ++++++++++++++++++++ 11 files changed, 627 insertions(+), 6 deletions(-) create mode 100644 benchmark/beam/config_meta_reme.yaml create mode 100644 benchmark/beam/run_meta.py create mode 100644 benchmark/longmemeval/config_meta_reme.yaml create mode 100644 benchmark/longmemeval/run_meta.py create mode 100644 meta-reme/data_preparation/__init__.py create mode 100644 meta-reme/data_preparation/basic.py create mode 100644 meta-reme/data_preparation/beam.py create mode 100644 meta-reme/data_preparation/lme.py create mode 100644 meta-reme/run.py create mode 100644 tests/unit_sandbox/test_meta_reme_run.py diff --git a/benchmark/beam/config_meta_reme.yaml b/benchmark/beam/config_meta_reme.yaml new file mode 100644 index 00000000..78290e9f --- /dev/null +++ b/benchmark/beam/config_meta_reme.yaml @@ -0,0 +1,8 @@ +# Meta-ReMe preparation configuration for BEAM. +meta_workspace: "benchmark/beam/meta-workspace" + +dataset: + name: "beam" + source: "benchmark/beam/dataset/BEAM" + variant: "1M" # 100K | 500K | 1M | 10M + train_case_ids: [] # empty = all cases; or ["1", "2", "3"] diff --git a/benchmark/beam/run_meta.py b/benchmark/beam/run_meta.py new file mode 100644 index 00000000..90c35592 --- /dev/null +++ b/benchmark/beam/run_meta.py @@ -0,0 +1,30 @@ +"""Launch Meta-ReMe for BEAM using one YAML configuration.""" + +from __future__ import annotations + +import argparse +from pathlib import Path +import subprocess +import sys + +PROJECT_ROOT = Path(__file__).resolve().parents[2] + + +def main() -> None: + """Forward the selected YAML configuration to Meta-ReMe.""" + + parser = argparse.ArgumentParser(description="Run Meta-ReMe for BEAM") + parser.add_argument("--config", type=Path, required=True, help="Path to config_meta_reme.yaml") + args = parser.parse_args() + + command = [ + sys.executable, + str(PROJECT_ROOT / "meta-reme/run.py"), + "--config", + str(args.config.resolve()), + ] + subprocess.run(command, cwd=PROJECT_ROOT, check=True) + + +if __name__ == "__main__": + main() diff --git a/benchmark/longmemeval/config_meta_reme.yaml b/benchmark/longmemeval/config_meta_reme.yaml new file mode 100644 index 00000000..0ad548ce --- /dev/null +++ b/benchmark/longmemeval/config_meta_reme.yaml @@ -0,0 +1,7 @@ +# Meta-ReMe preparation configuration for LongMemEval. +meta_workspace: "benchmark/longmemeval/meta-workspace" + +dataset: + name: "longmemeval" + source: "benchmark/longmemeval/dataset/longmemeval_s_reme_cleaned.json" + train_case_ids: [] # empty = all cases; or a list of question IDs diff --git a/benchmark/longmemeval/run_meta.py b/benchmark/longmemeval/run_meta.py new file mode 100644 index 00000000..9d226df9 --- /dev/null +++ b/benchmark/longmemeval/run_meta.py @@ -0,0 +1,30 @@ +"""Launch Meta-ReMe for LongMemEval using one YAML configuration.""" + +from __future__ import annotations + +import argparse +from pathlib import Path +import subprocess +import sys + +PROJECT_ROOT = Path(__file__).resolve().parents[2] + + +def main() -> None: + """Forward the selected YAML configuration to Meta-ReMe.""" + + parser = argparse.ArgumentParser(description="Run Meta-ReMe for LongMemEval") + parser.add_argument("--config", type=Path, required=True, help="Path to config_meta_reme.yaml") + args = parser.parse_args() + + command = [ + sys.executable, + str(PROJECT_ROOT / "meta-reme/run.py"), + "--config", + str(args.config.resolve()), + ] + subprocess.run(command, cwd=PROJECT_ROOT, check=True) + + +if __name__ == "__main__": + main() diff --git a/meta-reme/data_preparation/__init__.py b/meta-reme/data_preparation/__init__.py new file mode 100644 index 00000000..f621ba51 --- /dev/null +++ b/meta-reme/data_preparation/__init__.py @@ -0,0 +1,36 @@ +"""Dataset-specific normalization for Meta-ReMe training data.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Iterable + +from data_preparation import beam, lme +from data_preparation.basic import write_normalized_dataset +from models import DatasetManifest + + +def prepare_dataset( + directory: Path, + dataset: str, + source: Path, + train_case_ids: Iterable[str] | None = None, + variant: str | None = None, +) -> DatasetManifest: + """Normalize the selected training cases into ``directory``.""" + + if dataset == "beam": + selected_variant = variant or "1M" + if selected_variant not in beam.VARIANTS: + raise ValueError(f"Unsupported BEAM variant: {selected_variant}") + selected = beam.selected_case_ids(source, selected_variant, train_case_ids) + cases = beam.iter_cases(source, selected_variant, selected) + return write_normalized_dataset(directory, dataset, selected_variant, cases) + if dataset == "longmemeval": + if variant is not None: + raise ValueError("LongMemEval does not accept a dataset variant") + return write_normalized_dataset(directory, dataset, None, lme.iter_cases(source, train_case_ids)) + raise ValueError(f"Unsupported dataset: {dataset}") + + +__all__ = ["prepare_dataset"] diff --git a/meta-reme/data_preparation/basic.py b/meta-reme/data_preparation/basic.py new file mode 100644 index 00000000..5bb26994 --- /dev/null +++ b/meta-reme/data_preparation/basic.py @@ -0,0 +1,60 @@ +"""Shared helpers for writing normalized Meta-ReMe datasets.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Iterable + +from models import CaseSpec, DatasetManifest, canonical_json_bytes, fingerprint + + +def select_case_ids(available: list[str], requested: Iterable[str] | None) -> list[str]: + """Return requested IDs after checking duplicates and availability.""" + + requested_ids = list(requested or []) + if len(requested_ids) != len(set(requested_ids)): + raise ValueError("Training case IDs must not contain duplicates") + if not requested_ids: + return available + unknown = sorted(set(requested_ids) - set(available)) + if unknown: + raise ValueError(f"Unknown training case IDs: {', '.join(unknown)}") + return requested_ids + + +def write_normalized_dataset( + directory: Path, + dataset: str, + variant: str | None, + cases: Iterable[CaseSpec], +) -> DatasetManifest: + """Write normalized cases and their reproducibility manifest.""" + + cases_directory = directory / "cases" + cases_directory.mkdir(parents=True) + case_ids: list[str] = [] + case_fingerprints: list[str] = [] + query_count = 0 + for index, case in enumerate(cases): + (cases_directory / f"{index:06d}.json").write_bytes(canonical_json_bytes(case) + b"\n") + case_ids.append(case.case_id) + case_fingerprints.append(fingerprint(case)) + query_count += len(case.queries) + normalized_fingerprint = fingerprint(case_fingerprints) + source_fingerprint = fingerprint( + { + "dataset": dataset, + "variant": variant, + "case_ids": case_ids, + "normalized_fingerprint": normalized_fingerprint, + }, + ) + manifest = DatasetManifest( + dataset=dataset, + source_fingerprint=source_fingerprint, + normalized_fingerprint=normalized_fingerprint, + case_count=len(case_ids), + query_count=query_count, + ) + (directory / "manifest.json").write_bytes(canonical_json_bytes(manifest) + b"\n") + return manifest diff --git a/meta-reme/data_preparation/beam.py b/meta-reme/data_preparation/beam.py new file mode 100644 index 00000000..bfbc53be --- /dev/null +++ b/meta-reme/data_preparation/beam.py @@ -0,0 +1,114 @@ +"""Normalize BEAM cases into the shared Meta-ReMe contracts.""" + +from __future__ import annotations + +from datetime import datetime +import json +from pathlib import Path +from typing import Any, Iterable, Iterator + +from data_preparation.basic import select_case_ids +from models import CaseSpec, QuerySpec, SessionSpec + +VARIANTS = ("100K", "500K", "1M", "10M") + + +def _parse_date(value: str) -> datetime: + for date_format in ("%B-%d-%Y", "%b-%d-%Y"): + try: + return datetime.strptime(value, date_format) + except ValueError: + continue + raise ValueError(f"Invalid BEAM time anchor: {value!r}") + + +def _sessions(raw_batches: list[dict[str, Any]], variant: str, case_id: str) -> list[SessionSpec]: + sessions: list[SessionSpec] = [] + for batch in raw_batches: + fallback_anchor = batch.get("time_anchor") or "January-1-2024" + previous_time: datetime | None = None + first_time: datetime | None = None + messages: list[dict[str, Any]] = [] + for turn in batch["turns"]: + anchor = next((message.get("time_anchor") for message in turn if message.get("time_anchor")), None) + current_time = _parse_date(anchor or fallback_anchor) if previous_time is None or anchor else previous_time + first_time = first_time or current_time + previous_time = current_time + for message in turn: + messages.append( + { + **message, + "name": message["role"], + "created_at": current_time.strftime("%Y-%m-%dT%H:%M:%S"), + }, + ) + batch_number = str(batch["batch_number"]) + sessions.append( + SessionSpec( + session_id=f"beam_{variant}_{case_id}_batch{batch_number}", + messages=messages, + metadata={"batch_number": batch["batch_number"], "date": first_time.strftime("%Y-%m-%d")}, + ), + ) + return sessions + + +def _queries(raw_questions: dict[str, list[dict[str, Any]]]) -> list[QuerySpec]: + queries: list[QuerySpec] = [] + for question_type in sorted(raw_questions): + for index, raw in enumerate(raw_questions[question_type], start=1): + golden_answer = raw.get("ideal_answer", raw.get("ideal_response", raw.get("answer"))) + if golden_answer is None: + golden_answer = raw.get("rubric", []) + metadata = { + key: value + for key, value in raw.items() + if key not in {"question", "ideal_answer", "ideal_response", "answer"} + } + metadata["question_type"] = question_type + queries.append( + QuerySpec( + query_id=f"{question_type}:{index}", + question=raw["question"], + golden_answer=golden_answer, + metadata=metadata, + ), + ) + return queries + + +def selected_case_ids(source: Path, variant: str, requested: Iterable[str] | None) -> list[str]: + """Select and validate case IDs for one BEAM variant.""" + + cases_root = source / "chats" / variant + if not cases_root.is_dir(): + raise FileNotFoundError(f"BEAM variant does not exist: {cases_root}") + available = sorted( + (path.name for path in cases_root.iterdir() if path.is_dir()), + key=lambda value: (not value.isdigit(), int(value) if value.isdigit() else value), + ) + return select_case_ids(available, requested) + + +def iter_cases(source: Path, variant: str, case_ids: Iterable[str]) -> Iterator[CaseSpec]: + """Yield normalized BEAM cases one at a time.""" + + cases_root = source / "chats" / variant + for case_id in case_ids: + case_root = cases_root / case_id + raw_chat = json.loads((case_root / "chat.json").read_text(encoding="utf-8")) + raw_questions = json.loads( + (case_root / "probing_questions/probing_questions.json").read_text(encoding="utf-8"), + ) + yield CaseSpec( + case_id=case_id, + sessions=_sessions(raw_chat, variant, case_id), + queries=_queries(raw_questions), + metadata={"dataset": "beam", "variant": variant}, + ) + + +def load_cases(source: Path, variant: str, train_case_ids: Iterable[str] | None) -> list[CaseSpec]: + """Load selected BEAM cases into the unified contracts.""" + + return list(iter_cases(source, variant, selected_case_ids(source, variant, train_case_ids))) diff --git a/meta-reme/data_preparation/lme.py b/meta-reme/data_preparation/lme.py new file mode 100644 index 00000000..1d0cb3dd --- /dev/null +++ b/meta-reme/data_preparation/lme.py @@ -0,0 +1,66 @@ +"""Normalize LongMemEval cases into the shared Meta-ReMe contracts.""" + +from __future__ import annotations + +from datetime import datetime +import json +from pathlib import Path +from typing import Iterable, Iterator + +from data_preparation.basic import select_case_ids +from models import CaseSpec, QuerySpec, SessionSpec + + +def iter_cases(source: Path, train_case_ids: Iterable[str] | None) -> Iterator[CaseSpec]: + """Yield selected LongMemEval question IDs as unified cases.""" + + raw_items = json.loads(source.read_text(encoding="utf-8")) + by_id = {str(item["question_id"]): item for item in raw_items} + for case_id in select_case_ids(list(by_id), train_case_ids): + item = by_id[case_id] + sessions = [] + for date, session_id, messages in zip( + item["haystack_dates"], + item["haystack_session_ids"], + item["haystack_sessions"], + strict=True, + ): + timestamp = datetime.strptime(date, "%Y/%m/%d (%a) %H:%M").strftime("%Y-%m-%dT%H:%M:%S") + normalized_messages = [ + {**message, "name": message["role"], "created_at": timestamp} for message in messages + ] + sessions.append( + SessionSpec(session_id=str(session_id), messages=normalized_messages, metadata={"date": date}), + ) + query_metadata = { + key: value + for key, value in item.items() + if key + not in { + "question_id", + "question", + "answer", + "haystack_dates", + "haystack_session_ids", + "haystack_sessions", + } + } + yield CaseSpec( + case_id=case_id, + sessions=sessions, + queries=[ + QuerySpec( + query_id=case_id, + question=item["question"], + golden_answer=item["answer"], + metadata=query_metadata, + ), + ], + metadata={"dataset": "longmemeval"}, + ) + + +def load_cases(source: Path, train_case_ids: Iterable[str] | None) -> list[CaseSpec]: + """Load selected LongMemEval cases into the unified contracts.""" + + return list(iter_cases(source, train_case_ids)) diff --git a/meta-reme/design.md b/meta-reme/design.md index 83e4cc97..5d58d022 100644 --- a/meta-reme/design.md +++ b/meta-reme/design.md @@ -12,9 +12,7 @@ Meta-ReMe 放置于 `meta-reme/`，入口为： ```bash python meta-reme/run.py \ - --dataset longmemeval \ - --workspace /path/to/workspace \ - --config meta-reme/configs/longmemeval.yaml + --config benchmark/longmemeval/config_meta_reme.yaml ``` 主要模块如下： @@ -32,9 +30,9 @@ meta-reme/ evaluator.py result_store.py bundle_builder.py - datasets/ - base.py - longmemeval.py + data_preparation/ + basic.py + lme.py beam.py ``` diff --git a/meta-reme/run.py b/meta-reme/run.py new file mode 100644 index 00000000..75fcd564 --- /dev/null +++ b/meta-reme/run.py @@ -0,0 +1,133 @@ +"""Initialize a Meta-ReMe workspace from a benchmark YAML configuration.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +import tempfile +from typing import Any, Iterable + +import yaml + +from data_preparation import prepare_dataset +from models import ( + BudgetSpec, + DatasetSpec, + DomainSpec, + ProposerSpec, + SandboxSpec, + ScopeSpec, +) +from workspace import Workspace + +PROJECT_ROOT = Path(__file__).resolve().parent.parent +DEFAULT_SOURCES = { + "beam": PROJECT_ROOT / "benchmark/beam/dataset/BEAM", + "longmemeval": PROJECT_ROOT / "benchmark/longmemeval/dataset/longmemeval_s_reme_cleaned.json", +} + + +def _domain_spec(dataset: str, source: Path, source_fingerprint: str) -> DomainSpec: + return DomainSpec( + dataset=DatasetSpec(name=dataset, source=str(source.resolve()), fingerprint=source_fingerprint), + bundle_target="beam" if dataset == "beam" else "lme", + benchmark_runner=f"benchmark.{'beam' if dataset == 'beam' else 'longmemeval'}.run", + scorer="mean_query_score", + scope=ScopeSpec(), + sandbox=SandboxSpec(image="reme:latest", timeout_seconds=3600), + proposer=ProposerSpec(model="not-configured"), + budget=BudgetSpec(max_proposals=1), + ) + + +def prepare_workspace( + meta_workspace: Path, + dataset: str, + train_case_ids: Iterable[str] | None = None, + dataset_variant: str | None = None, + dataset_source: Path | None = None, +) -> Workspace: + """Create/open a workspace and install only the selected training cases.""" + + meta_workspace = Path(meta_workspace).resolve() + source = Path(dataset_source or DEFAULT_SOURCES.get(dataset, "")).resolve() + with tempfile.TemporaryDirectory(prefix="meta-reme-dataset-") as temporary: + normalized = Path(temporary) + manifest = prepare_dataset( + normalized, + dataset, + source, + train_case_ids=train_case_ids, + variant=dataset_variant, + ) + spec = _domain_spec(dataset, source, manifest.source_fingerprint) + is_empty_directory = meta_workspace.is_dir() and not any(meta_workspace.iterdir()) + workspace = ( + Workspace.create(meta_workspace, spec) + if not meta_workspace.exists() or is_empty_directory + else Workspace.open(meta_workspace, spec) + ) + with workspace.acquire_lock(): + if not workspace.path("datasets/search/manifest.json").exists(): + workspace.install_search_dataset(normalized) + return workspace + + +def _mapping(value: Any, label: str) -> dict[str, Any]: + if not isinstance(value, dict): + raise ValueError(f"{label} must be a YAML mapping") + return value + + +def _project_path(value: Any, label: str) -> Path: + if not isinstance(value, str) or not value: + raise ValueError(f"{label} must be a non-empty path string") + path = Path(value).expanduser() + return path if path.is_absolute() else PROJECT_ROOT / path + + +def load_config(config_path: Path) -> dict[str, Any]: + """Load and validate the dataset-preparation portion of a Meta-ReMe config.""" + + raw = _mapping(yaml.safe_load(config_path.read_text(encoding="utf-8")), "config") + dataset_config = _mapping(raw.get("dataset"), "dataset") + dataset = dataset_config.get("name") + if dataset not in DEFAULT_SOURCES: + raise ValueError(f"Unsupported dataset: {dataset}") + case_ids = dataset_config.get("train_case_ids", []) + if not isinstance(case_ids, list) or any(not isinstance(case_id, (str, int)) for case_id in case_ids): + raise ValueError("dataset.train_case_ids must be a list of strings or integers") + source_value = dataset_config.get("source") + source = DEFAULT_SOURCES[dataset] if source_value is None else _project_path(source_value, "dataset.source") + return { + "meta_workspace": _project_path(raw.get("meta_workspace"), "meta_workspace"), + "dataset": dataset, + "train_case_ids": [str(case_id) for case_id in case_ids], + "dataset_variant": dataset_config.get("variant"), + "dataset_source": source, + } + + +def parse_args() -> argparse.Namespace: + """Parse the single configuration-file argument.""" + + parser = argparse.ArgumentParser(description="Initialize Meta-ReMe and normalize its training set") + parser.add_argument("--config", type=Path, required=True, help="Path to config_meta_reme.yaml") + return parser.parse_args() + + +def main() -> None: + """Prepare the configured workspace and print a short dataset summary.""" + + args = parse_args() + workspace = prepare_workspace(**load_config(args.config)) + manifest = json.loads(workspace.path("datasets/search/manifest.json").read_text(encoding="utf-8")) + print( + f"Meta-ReMe workspace ready: {workspace.root} " + f"({manifest['case_count']} training cases, {manifest['query_count']} queries)", + ) + + +if __name__ == "__main__": + main() diff --git a/tests/unit_sandbox/test_meta_reme_run.py b/tests/unit_sandbox/test_meta_reme_run.py new file mode 100644 index 00000000..bf32eb25 --- /dev/null +++ b/tests/unit_sandbox/test_meta_reme_run.py @@ -0,0 +1,139 @@ +"""Focused tests for Meta-ReMe training-data preparation.""" + +# pylint: disable=missing-function-docstring + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +import sys + +import yaml + +META_REME = Path(__file__).resolve().parents[2] / "meta-reme" +sys.path.insert(0, str(META_REME)) +SPEC = importlib.util.spec_from_file_location("meta_reme_run", META_REME / "run.py") +assert SPEC and SPEC.loader +run = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(run) + +from data_preparation import beam # noqa: E402 pylint: disable=wrong-import-position + + +def _lme_item(case_id: str) -> dict: + return { + "question_id": case_id, + "question_type": "single-session-user", + "question": f"question {case_id}", + "question_date": "2023/05/21 (Sun) 02:21", + "answer": f"answer {case_id}", + "original_answer": f"answer {case_id}", + "answer_session_ids": [f"session-{case_id}"], + "original_answer_session_ids": [f"session-{case_id}"], + "haystack_dates": ["2023/05/20 (Sat) 02:21"], + "haystack_session_ids": [f"session-{case_id}"], + "haystack_sessions": [[{"role": "user", "content": f"memory {case_id}"}]], + } + + +def test_prepare_workspace_stores_only_selected_longmemeval_cases(tmp_path: Path) -> None: + source = tmp_path / "longmemeval.json" + source.write_text(json.dumps([_lme_item("train"), _lme_item("held-out")]), encoding="utf-8") + workspace_root = tmp_path / "meta-workspace" + + workspace = run.prepare_workspace( + workspace_root, + "longmemeval", + train_case_ids=["train"], + dataset_source=source, + ) + + manifest = json.loads(workspace.path("datasets/search/manifest.json").read_text(encoding="utf-8")) + case_files = list(workspace.path("datasets/search/cases").glob("*.json")) + stored_case = json.loads(case_files[0].read_text(encoding="utf-8")) + assert manifest["case_count"] == 1 + assert manifest["query_count"] == 1 + assert stored_case["case_id"] == "train" + stored_text = "".join(path.read_text(encoding="utf-8") for path in case_files) + assert "held-out" not in stored_text + assert not case_files[0].stat().st_mode & 0o222 + + assert ( + run.prepare_workspace( + workspace_root, + "longmemeval", + train_case_ids=["train"], + dataset_source=source, + ).root + == workspace.root + ) + + +def test_prepare_workspace_initializes_an_existing_empty_directory(tmp_path: Path) -> None: + source = tmp_path / "longmemeval.json" + source.write_text(json.dumps([_lme_item("train")]), encoding="utf-8") + workspace_root = tmp_path / "meta-workspace" + workspace_root.mkdir() + + workspace = run.prepare_workspace(workspace_root, "longmemeval", dataset_source=source) + + assert workspace.path(".meta-reme-workspace.json").is_file() + assert workspace.path("datasets/search/manifest.json").is_file() + + +def test_load_beam_cases_respects_variant_and_selection(tmp_path: Path) -> None: + case_root = tmp_path / "chats/100K/7" + (case_root / "probing_questions").mkdir(parents=True) + (case_root / "chat.json").write_text( + json.dumps( + [ + { + "batch_number": 1, + "turns": [ + [ + {"role": "user", "content": "remember", "time_anchor": "March-15-2024"}, + {"role": "assistant", "content": "okay"}, + ], + ], + }, + ], + ), + encoding="utf-8", + ) + (case_root / "probing_questions/probing_questions.json").write_text( + json.dumps({"information_extraction": [{"question": "what?", "ideal_answer": "remember"}]}), + encoding="utf-8", + ) + + cases = beam.load_cases(tmp_path, "100K", ["7"]) + + assert [case.case_id for case in cases] == ["7"] + assert cases[0].metadata["variant"] == "100K" + assert cases[0].queries[0].query_id == "information_extraction:1" + assert cases[0].sessions[0].messages[1]["created_at"] == "2024-03-15T00:00:00" + + +def test_load_config_reads_all_preparation_arguments(tmp_path: Path) -> None: + config_path = tmp_path / "config_meta_reme.yaml" + config_path.write_text( + yaml.safe_dump( + { + "meta_workspace": "benchmark/beam/meta-workspace", + "dataset": { + "name": "beam", + "source": "benchmark/beam/dataset/BEAM", + "variant": "100K", + "train_case_ids": [1, "2"], + }, + }, + ), + encoding="utf-8", + ) + + config = run.load_config(config_path) + + assert config["meta_workspace"] == run.PROJECT_ROOT / "benchmark/beam/meta-workspace" + assert config["dataset_source"] == run.PROJECT_ROOT / "benchmark/beam/dataset/BEAM" + assert config["dataset_variant"] == "100K" + assert config["train_case_ids"] == ["1", "2"]