admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-30 21:19:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use generic names in multidimensional add_features test

Browse Source
This commit is contained in:
felixmin
2026-03-21 12:43:12 +01:00
committed by CarolinePascal
parent e0e03887e0
commit 910aac4c03
1 changed files with 8 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/datasets/test_dataset_tools.py
+8 -8
View File
@@ -381,7 +381,7 @@ def test_add_features_with_callable(sample_dataset, tmp_path):
def test_add_features_with_multidimensional_values(sample_dataset, tmp_path):
"""Test adding a multi-dimensional per-frame feature."""
num_frames = sample_dataset.meta.total_frames
latent_values = np.random.randn(num_frames, 1, 4).astype(np.float32)
feature_values = np.random.randn(num_frames, 1, 4).astype(np.float32)
feature_info = {
"dtype": "float32",
@@ -389,7 +389,7 @@ def test_add_features_with_multidimensional_values(sample_dataset, tmp_path):
"names": None,
}
features = {
"latent_vectors": (latent_values, feature_info),
"multi_dim_feature": (feature_values, feature_info),
}
with (
@@ -397,19 +397,19 @@ def test_add_features_with_multidimensional_values(sample_dataset, tmp_path):
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(tmp_path / "with_latents")
mock_snapshot_download.return_value = str(tmp_path / "with_multi_dim_feature")
new_dataset = add_features(
dataset=sample_dataset,
features=features,
output_dir=tmp_path / "with_latents",
output_dir=tmp_path / "with_multi_dim_feature",
)
assert "latent_vectors" in new_dataset.meta.features
assert "multi_dim_feature" in new_dataset.meta.features
sample_item = new_dataset[0]
assert "latent_vectors" in sample_item
assert isinstance(sample_item["latent_vectors"], torch.Tensor)
assert tuple(sample_item["latent_vectors"].shape) == (1, 4)
assert "multi_dim_feature" in sample_item
assert isinstance(sample_item["multi_dim_feature"], torch.Tensor)
assert tuple(sample_item["multi_dim_feature"].shape) == (1, 4)
def test_add_existing_feature(sample_dataset, tmp_path):