File size: 6,051 Bytes
319eb16 | 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 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | """Utilities for indexing and looking up AgiBotWorld task JSONs.
This module builds and uses an index mapping episode_id -> task_json_path,
so downstream scripts can quickly resolve which task file contains a given
episode. The index is stored as a JSON file on disk for fast re-use.
"""
from __future__ import annotations
import json
import os
from collections.abc import Iterable
from glob import glob
# Default locations used by the downloader
DEFAULT_DATASET_ROOT: str = os.path.join("datasets", "tmp", "agibotworld_alpha")
DEFAULT_TASK_INFO_DIR: str = os.path.join(DEFAULT_DATASET_ROOT, "task_info")
DEFAULT_INDEX_PATH: str = os.path.join(DEFAULT_DATASET_ROOT, "episode_to_task_index.json")
def _task_json_files(task_info_dir: str) -> Iterable[str]:
"""Yield task JSON files within the given directory.
Parameters
- task_info_dir: Directory that contains task_*.json files
"""
pattern = os.path.join(task_info_dir, "task_*.json")
for file_path in sorted(glob(pattern)):
# Guard against directories or non-files matching the pattern
if os.path.isfile(file_path):
yield file_path
def build_episode_to_task_index(
task_info_dir: str = DEFAULT_TASK_INFO_DIR,
output_index_path: str = DEFAULT_INDEX_PATH,
verbose: bool = True,
) -> dict[str, str]:
"""Build an index mapping episode_id -> task_json_path and save to disk.
Keys are stored as strings in the JSON index for portability, but callers
may freely pass `int` episode ids to the lookup helpers provided here.
Returns the mapping in-memory as well.
"""
if not os.path.isdir(task_info_dir):
raise FileNotFoundError(f"Task info directory not found: {os.path.abspath(task_info_dir)}")
episode_to_task: dict[str, str] = {}
duplicate_episode_ids: list[str] = []
for json_path in _task_json_files(task_info_dir):
try:
with open(json_path, encoding="utf-8") as f:
entries = json.load(f)
except Exception as exc: # pragma: no cover - defensive logging
if verbose:
print(f"[agibot_helper] Skipping unreadable file {json_path}: {exc}")
continue
if not isinstance(entries, list):
if verbose:
print(f"[agibot_helper] Unexpected JSON structure in {json_path}; expected a list.")
continue
for item in entries:
if not isinstance(item, dict) or "episode_id" not in item:
continue
episode_id_val = item["episode_id"]
# Normalize to string for JSON index stability
episode_key = str(episode_id_val)
if episode_key in episode_to_task:
duplicate_episode_ids.append(episode_key)
episode_to_task[episode_key] = json_path
# Ensure output directory exists
os.makedirs(os.path.dirname(output_index_path), exist_ok=True)
with open(output_index_path, "w", encoding="utf-8") as f:
json.dump(episode_to_task, f)
if verbose:
total = len(episode_to_task)
dupes = len(set(duplicate_episode_ids))
print(
f"[agibot_helper] Built episode->task index with {total} episodes"
+ (f" ({dupes} duplicate ids overwritten)" if dupes else "")
+ f" at {os.path.abspath(output_index_path)}"
)
return episode_to_task
def load_episode_to_task_index(index_path: str = DEFAULT_INDEX_PATH) -> dict[str, str]:
"""Load the episode->task index from disk.
Raises FileNotFoundError if the index doesn't exist.
"""
if not os.path.isfile(index_path):
raise FileNotFoundError(
f"Episode index not found: {os.path.abspath(index_path)}. Build it with build_episode_to_task_index()."
)
with open(index_path, encoding="utf-8") as f:
return json.load(f)
def ensure_episode_index(
task_info_dir: str = DEFAULT_TASK_INFO_DIR,
index_path: str = DEFAULT_INDEX_PATH,
verbose: bool = False,
) -> dict[str, str]:
"""Return an in-memory episode->task index, building it if missing."""
if os.path.isfile(index_path):
return load_episode_to_task_index(index_path)
return build_episode_to_task_index(task_info_dir, index_path, verbose=verbose)
def find_task_json_for_episode(
episode_id: int | str,
task_info_dir: str = DEFAULT_TASK_INFO_DIR,
index_path: str = DEFAULT_INDEX_PATH,
) -> str:
"""Return the path to the task JSON file containing the given episode.
If the index file is missing, it will be built automatically.
Raises KeyError if the episode id is not present in the index.
"""
index = ensure_episode_index(task_info_dir=task_info_dir, index_path=index_path)
key = str(episode_id)
if key not in index:
raise KeyError(f"Episode id {episode_id} not found in index at {os.path.abspath(index_path)}")
return index[key]
def get_episode_record(
episode_id: int | str,
task_info_dir: str = DEFAULT_TASK_INFO_DIR,
index_path: str = DEFAULT_INDEX_PATH,
) -> tuple[str, dict]:
"""Return (json_path, episode_record) for the given episode id.
This reads the relevant task JSON file and returns the full dictionary
entry for the requested episode. Raises KeyError if the episode isn't
present.
"""
json_path = find_task_json_for_episode(episode_id=episode_id, task_info_dir=task_info_dir, index_path=index_path)
with open(json_path, encoding="utf-8") as f:
entries = json.load(f)
key = int(episode_id)
for item in entries:
if isinstance(item, dict) and int(item.get("episode_id", -1)) == key:
return json_path, item
raise KeyError(f"Episode id {episode_id} not found inside file {os.path.abspath(json_path)}")
__all__ = [
"DEFAULT_DATASET_ROOT",
"DEFAULT_INDEX_PATH",
"DEFAULT_TASK_INFO_DIR",
"build_episode_to_task_index",
"ensure_episode_index",
"find_task_json_for_episode",
"get_episode_record",
"load_episode_to_task_index",
]
|