File size: 5,311 Bytes
2abcc30 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | from __future__ import annotations
import json
import random
from pathlib import Path
from typing import Any
import pyarrow.parquet as pq
from albedo_eval_service.remote.dataset import EvalSample, apply_submit_protocol, load_manifest_samples
from albedo_eval_service.shared.dataset_manifest import load_manifest_file
from albedo_eval_service.shared.sampling import multi_source_manifest_sample_ids
from .constants import TOKENIZER_DIR
def load_samples(
dataset_root: Path,
*,
sample_ids: list[str] | None,
sample_count: int,
seed: str,
tokenizer_path: Path | None = None,
) -> list[EvalSample]:
tokenizer = str(tokenizer_path or TOKENIZER_DIR)
if not sample_ids:
sample_ids = pick_sample_ids(dataset_root, sample_count=sample_count, seed=seed)
samples = load_manifest_samples(
dataset_root=dataset_root,
sample_ids=sample_ids,
tokenizer_path=tokenizer,
enable_thinking=True,
)
return apply_submit_protocol(
samples,
salt=seed,
keep_original_ratio=0.25,
tokenizer_path=tokenizer,
enable_thinking=True,
)
def pick_sample_ids(dataset_root: Path, *, sample_count: int, seed: str) -> list[str]:
manifest_path = dataset_root / "manifest.json"
if manifest_path.is_file():
try:
manifest = load_manifest_file(manifest_path, expected_sha256="")
if "sources" in manifest:
return multi_source_manifest_sample_ids(
manifest, block_hash=seed, sample_count=sample_count
)
except Exception as exc:
print(f"official sampler unavailable ({exc}); falling back to lite sampler", flush=True)
return lite_sample_ids(dataset_root, sample_count=sample_count, seed=seed)
def lite_sample_ids(dataset_root: Path, *, sample_count: int, seed: str) -> list[str]:
"""Prefix sampler for a partial local tree (no pinned multi-source manifest)."""
shards = sorted(dataset_root.glob("*/data/train-*.parquet"))
if not shards:
raise FileNotFoundError(
f"no parquet shards under {dataset_root}/*/data/ — run: python -m local_eval download-lite-data"
)
rng = random.Random(seed)
candidates: list[str] = []
for shard in shards:
rel = shard.relative_to(dataset_root).as_posix()
table = pq.read_table(shard, columns=_message_columns(shard))
for row_idx, raw in enumerate(table.to_pydict()[table.column_names[0]]):
turns = _as_turns(raw)
assistant = [i for i, turn in enumerate(turns) if _role(turn) == "assistant"]
if len(assistant) < 3:
continue
for turn_idx in (1, 2, min(len(assistant) - 1, 4)):
if turn_idx < len(assistant):
candidates.append(f"{rel}:{row_idx}:{turn_idx}")
rng.shuffle(candidates)
if len(candidates) < sample_count:
raise ValueError(f"only {len(candidates)} lite prefixes available, need {sample_count}")
return candidates[:sample_count]
def leftover_observations(dataset_root: Path, sample_id: str) -> list[str]:
"""Gold user turns after the sampled cut — offline stand-in for repo-context."""
shard_name, row_idx, turn_idx = sample_id.rsplit(":", 2)
row = _read_row(dataset_root / shard_name, int(row_idx))
turns = _as_turns(row.get("messages") or row.get("turns") or row.get("conversation"))
assistant = [i for i, turn in enumerate(turns) if _role(turn) == "assistant"]
if int(turn_idx) >= len(assistant):
return []
start = assistant[int(turn_idx)] + 1
return [
_content(turn)
for turn in turns[start:]
if _role(turn) in {"user", "tool"} and _content(turn)
]
def _message_columns(shard: Path) -> list[str]:
schema = pq.read_schema(shard)
for name in ("messages", "turns", "conversation", "trajectory"):
if name in schema.names:
return [name]
return [schema.names[0]]
def _read_row(path: Path, row_idx: int) -> dict[str, Any]:
parquet = pq.ParquetFile(path)
seen = 0
for batch in parquet.iter_batches(batch_size=512):
if seen + batch.num_rows <= row_idx:
seen += batch.num_rows
continue
return {k: (v[0] if isinstance(v, list) and len(v) == 1 else v) for k, v in batch.slice(row_idx - seen, 1).to_pydict().items()}
raise IndexError(row_idx)
def _as_turns(value: Any) -> list[Any]:
parsed = value
if isinstance(value, str):
try:
parsed = json.loads(value)
except json.JSONDecodeError:
return []
if isinstance(parsed, dict):
for key in ("messages", "turns", "conversation"):
if isinstance(parsed.get(key), list):
return parsed[key]
return []
return parsed if isinstance(parsed, list) else []
def _role(turn: Any) -> str:
if not isinstance(turn, dict):
return ""
return str(turn.get("role") or turn.get("speaker") or turn.get("from") or "").lower()
def _content(turn: Any) -> str:
if not isinstance(turn, dict):
return str(turn or "")
for key in ("content", "text", "value", "message"):
value = turn.get(key)
if value:
return str(value)
return ""
|