wehub-resource-sync
124 lines
3.8 KiB
Python
124 lines
3.8 KiB
Python
"""ALFWorld task dataloader."""
|
|
from __future__ import annotations
|
|
|
|
from skillopt.datasets.base import BatchSpec, SplitDataLoader
|
|
|
|
|
|
class ALFWorldDataLoader(SplitDataLoader):
|
|
"""ALFWorld batch planner.
|
|
|
|
In split_dir mode, batches are fixed gamefile items so ablations differ
|
|
only in how the same training set is batched.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
split_dir: str = "",
|
|
data_path: str = "",
|
|
split_mode: str = "split_dir",
|
|
split_ratio: str = "2:1:7",
|
|
split_seed: int = 42,
|
|
split_output_dir: str = "",
|
|
seed: int = 42,
|
|
limit: int = 0,
|
|
train_size: int = 0,
|
|
**kwargs,
|
|
) -> None:
|
|
super().__init__(
|
|
split_dir=split_dir,
|
|
data_path=data_path,
|
|
split_mode=split_mode,
|
|
split_ratio=split_ratio,
|
|
split_seed=split_seed,
|
|
split_output_dir=split_output_dir,
|
|
seed=seed,
|
|
limit=limit,
|
|
)
|
|
self.train_size_override = int(train_size or 0)
|
|
|
|
@staticmethod
|
|
def _metadata_for_items(items: list[dict], split: str, phase: str) -> dict:
|
|
gamefiles = [str(item.get("gamefile") or "") for item in items]
|
|
if any(not gamefile for gamefile in gamefiles):
|
|
raise ValueError("ALFWorld split items must contain non-empty gamefile paths.")
|
|
eval_dataset = "train"
|
|
is_train = phase == "train"
|
|
first = gamefiles[0] if gamefiles else ""
|
|
if "/valid_seen/" in first:
|
|
eval_dataset = "eval_in_distribution"
|
|
is_train = False
|
|
elif "/valid_unseen/" in first:
|
|
eval_dataset = "eval_out_of_distribution"
|
|
is_train = False
|
|
return {
|
|
"eval_dataset": eval_dataset,
|
|
"is_train": is_train,
|
|
"gamefiles": gamefiles,
|
|
"result_ids": [str(item.get("id") or idx) for idx, item in enumerate(items)],
|
|
}
|
|
|
|
def get_train_size(self) -> int:
|
|
if self.train_size_override > 0:
|
|
return self.train_size_override
|
|
return super().get_train_size()
|
|
|
|
def build_train_batch(self, batch_size: int, seed: int, **kwargs) -> BatchSpec:
|
|
batch = super().build_train_batch(batch_size=batch_size, seed=seed, **kwargs)
|
|
items = list(batch.payload or [])
|
|
batch.metadata.update(self._metadata_for_items(items, "train", "train"))
|
|
return BatchSpec(
|
|
phase="train",
|
|
split="train",
|
|
seed=seed,
|
|
batch_size=len(items),
|
|
payload=items,
|
|
metadata=batch.metadata,
|
|
)
|
|
|
|
def plan_train_epoch(
|
|
self,
|
|
*,
|
|
epoch: int,
|
|
steps_per_epoch: int,
|
|
accumulation: int,
|
|
batch_size: int,
|
|
seed: int,
|
|
**kwargs,
|
|
) -> list[BatchSpec]:
|
|
batches = super().plan_train_epoch(
|
|
epoch=epoch,
|
|
steps_per_epoch=steps_per_epoch,
|
|
accumulation=accumulation,
|
|
batch_size=batch_size,
|
|
seed=seed,
|
|
**kwargs,
|
|
)
|
|
for batch in batches:
|
|
items = list(batch.payload or [])
|
|
batch.metadata.update(self._metadata_for_items(items, "train", "train"))
|
|
return batches
|
|
|
|
def build_eval_batch(
|
|
self,
|
|
env_num: int,
|
|
split: str,
|
|
seed: int,
|
|
**kwargs,
|
|
) -> BatchSpec:
|
|
batch = super().build_eval_batch(
|
|
env_num=env_num,
|
|
split=split,
|
|
seed=seed,
|
|
**kwargs,
|
|
)
|
|
items = list(batch.payload or [])
|
|
batch.metadata.update(self._metadata_for_items(items, split, "eval"))
|
|
return BatchSpec(
|
|
phase="eval",
|
|
split=split,
|
|
seed=seed,
|
|
batch_size=len(items),
|
|
payload=items,
|
|
metadata=batch.metadata,
|
|
)
|