unienv_data.integrations.huggingface¶
HFAsUniEnvDataset
¶
HFAsUniEnvDataset(hf_dataset: Dataset, space: Space[Any, BDeviceType, BDtypeType, BRNGType])
Bases: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
BACKEND_TO_FORMAT_MAP
class-attribute
instance-attribute
¶
BACKEND_TO_FORMAT_MAP = {'numpy': 'numpy', 'pytorch': 'torch', 'jax': 'jax'}
create
staticmethod
¶
create(hf_dataset: Dataset, backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType], device: Optional[BDeviceType] = None) -> HFAsUniEnvDataset[BArrayType, BDeviceType, BDtypeType, BRNGType]
get_flattened_at
¶
get_flattened_at(idx: Union[IndexableType, BArrayType]) -> BArrayType
Fetch samples as flattened backend arrays.
get_flattened_at_with_metadata
¶
get_flattened_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> Tuple[BArrayType, Optional[Dict[str, Any]]]
Fetch flattened samples together with optional per-sample metadata.
set_flattened_at
¶
set_flattened_at(idx: Union[IndexableType, BArrayType], value: BArrayType) -> None
Overwrite existing samples using flattened data.
append_flattened
¶
append_flattened(value: BArrayType) -> None
Append one flattened sample to the batch.
extend_flattened
¶
extend_flattened(value: BArrayType) -> None
Append a batched block of flattened samples.
set_at
¶
set_at(idx: Union[IndexableType, BArrayType], value: BatchT) -> None
Overwrite existing samples using structured data.
remove_at
¶
remove_at(idx: Union[IndexableType, BArrayType]) -> None
Remove one or more samples from the batch.
extend_from
¶
extend_from(other: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], chunk_size: int = 8, tqdm: bool = False) -> None
Copy data from another batch in bounded-size chunks.
get_slice
¶
get_slice(idx: Union[IndexableType, BArrayType]) -> BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a subset of indices.