Skip to content

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'}

is_mutable class-attribute instance-attribute

is_mutable = False

hf_dataset instance-attribute

hf_dataset = hf_dataset

single_space instance-attribute

single_space = single_space

single_metadata_space instance-attribute

single_metadata_space = single_metadata_space

backend property

backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]

device property

device: Optional[BDeviceType]

create staticmethod

create(hf_dataset: Dataset, backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType], device: Optional[BDeviceType] = None) -> HFAsUniEnvDataset[BArrayType, BDeviceType, BDtypeType, BRNGType]

get_at_with_metadata

get_at_with_metadata(idx)

get_at

get_at(idx)

get_column

get_column(nested_keys)

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.

append

append(value: BatchT) -> None

Append one structured sample to the batch.

extend

extend(value: BatchT) -> None

Append a batched block of structured samples.

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.

close

close() -> None