mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-19 04:28:35 +00:00
feat(python)!: align Permutation.with_format("torch") with HuggingFace set_format("torch") (#3369)
Closes #3245. > **BREAKING CHANGE:** `with_format("torch")` no longer returns a list of stacked row tensors. It now returns per-row dicts so PyTorch's default `DataLoader` collate stacks them into `{col: tensor(B,)}`. Switch to `with_format("torch_row")` to keep the old shape. ### What changed `"torch"` now returns a list of per-row dicts (`[{col: tensor}, ...]`) at every indexed access path. The default `DataLoader` collate stacks them into a column-keyed batched dict, no custom `collate_fn` needed. The old shape is preserved under a new `"torch_row"` literal. `"torch_col"` is unchanged. The unbatching lives inside the transform (`batch_to_tensor_dict`), not `__getitems__`, so the shape survives pickling and works under `DataLoader(num_workers>0, multiprocessing_context="spawn")`. ### Format comparison | Format | `iter(batch_size=N)` | `__getitems__([0,1,2])` | `DataLoader` default collate | |---|---|---|---| | `"torch"` (new) | `list[{col: tensor}]` length N | `list[{col: tensor}]` length 3 | `{col: tensor(B,)}` | | `"torch_row"` (old `"torch"` behavior) | `list[tensor(n_cols,)]` length N | `list[tensor(n_cols,)]` length 3 | `tensor(B, n_cols)` | | `"torch_col"` (unchanged) | `tensor(n_cols, N)` | `tensor(n_cols, 3)` | needs `collate_fn=lambda x: x` | Output matches HuggingFace `Dataset.set_format("torch")` on container shape, keys, and values at every access path. The only divergence: HuggingFace downcasts `float64` to `torch.float32` by default, LanceDB preserves dtype. Verified by `scripts/verify_torch_format.py`. ### Migration ```python # Old default — column names lost, shape was tensor(B, n_cols) DataLoader(Permutation.identity(table).with_format("torch")) # New default — column names preserved DataLoader(Permutation.identity(table).with_format("torch")) # {col: tensor(B,)} # Keep old behavior DataLoader(Permutation.identity(table).with_format("torch_row")) # tensor(B, n_cols) ```
This commit is contained in:
@@ -935,14 +935,41 @@ def test_transform_fn(mem_db):
|
||||
try:
|
||||
import torch
|
||||
|
||||
torch_result = list(
|
||||
permutation.with_format("torch").iter(10, skip_last_batch=False)
|
||||
# "torch" returns a list of per-row dicts. Default DataLoader collate
|
||||
# stacks the per-row dicts back into a dict of batched tensors.
|
||||
torch_perm = permutation.with_format("torch")
|
||||
torch_batch = list(torch_perm.iter(10, skip_last_batch=False))[0]
|
||||
assert isinstance(torch_batch, list)
|
||||
assert len(torch_batch) == 10
|
||||
assert isinstance(torch_batch[0], dict)
|
||||
assert set(torch_batch[0].keys()) == {"id", "value"}
|
||||
assert isinstance(torch_batch[0]["id"], torch.Tensor)
|
||||
assert torch_batch[0]["id"].dtype == torch.int64
|
||||
|
||||
rows = torch_perm.__getitems__([0, 1, 2])
|
||||
assert isinstance(rows, list)
|
||||
assert len(rows) == 3
|
||||
assert isinstance(rows[0], dict)
|
||||
assert set(rows[0].keys()) == {"id", "value"}
|
||||
assert isinstance(rows[0]["id"], torch.Tensor)
|
||||
|
||||
# "torch_row" returns a list of tensors, one per row.
|
||||
torch_rows = list(
|
||||
permutation.with_format("torch_row").iter(10, skip_last_batch=False)
|
||||
)[0]
|
||||
assert isinstance(torch_result, list)
|
||||
assert len(torch_result) == 10
|
||||
assert isinstance(torch_result[0], torch.Tensor)
|
||||
assert torch_result[0].shape == (2,)
|
||||
assert torch_result[0].dtype == torch.int64
|
||||
assert isinstance(torch_rows, list)
|
||||
assert len(torch_rows) == 10
|
||||
assert isinstance(torch_rows[0], torch.Tensor)
|
||||
assert torch_rows[0].shape == (2,)
|
||||
assert torch_rows[0].dtype == torch.int64
|
||||
|
||||
# "torch_col" stacks columns into a single 2D tensor.
|
||||
torch_col = list(
|
||||
permutation.with_format("torch_col").iter(10, skip_last_batch=False)
|
||||
)[0]
|
||||
assert isinstance(torch_col, torch.Tensor)
|
||||
assert torch_col.shape == (2, 10)
|
||||
assert torch_col.dtype == torch.int64
|
||||
except ImportError:
|
||||
# Skip check if torch is not installed
|
||||
pass
|
||||
|
||||
@@ -148,15 +148,47 @@ def test_permutation_dataloader(mem_db):
|
||||
for batch in dataloader:
|
||||
assert batch["a"].size(0) == 10
|
||||
|
||||
permutation = permutation.with_format("torch")
|
||||
dataloader = torch.utils.data.DataLoader(permutation, batch_size=10, shuffle=True)
|
||||
# "torch" produces a list of per-row dicts per batch. The default
|
||||
# DataLoader collate stacks the per-row dicts back into a batched dict.
|
||||
torch_perm = permutation.with_format("torch")
|
||||
batch = next(torch_perm.iter(10, skip_last_batch=False))
|
||||
assert isinstance(batch, list)
|
||||
assert len(batch) == 10
|
||||
assert isinstance(batch[0], dict)
|
||||
assert isinstance(batch[0]["a"], torch.Tensor)
|
||||
rows = torch_perm.__getitems__([0, 1, 2])
|
||||
assert isinstance(rows, list)
|
||||
assert len(rows) == 3
|
||||
assert isinstance(rows[0], dict)
|
||||
assert isinstance(rows[0]["a"], torch.Tensor)
|
||||
dataloader = torch.utils.data.DataLoader(torch_perm, batch_size=10, shuffle=True)
|
||||
for batch in dataloader:
|
||||
assert isinstance(batch, dict)
|
||||
assert batch["a"].shape == (10,)
|
||||
# Spawn-based workers exercise the pickle round-trip path: the new
|
||||
# transform-as-list shape must survive pickling so workers produce the
|
||||
# same per-row dicts the parent does.
|
||||
spawn_loader = torch.utils.data.DataLoader(
|
||||
torch_perm,
|
||||
batch_size=10,
|
||||
num_workers=2,
|
||||
multiprocessing_context="spawn",
|
||||
)
|
||||
for batch in spawn_loader:
|
||||
assert isinstance(batch, dict)
|
||||
assert batch["a"].shape == (10,)
|
||||
|
||||
# "torch_row" returns a list of row tensors. Works with the default
|
||||
# DataLoader collate (stacks rows into 2D).
|
||||
row_perm = permutation.with_format("torch_row")
|
||||
dataloader = torch.utils.data.DataLoader(row_perm, batch_size=10, shuffle=True)
|
||||
for batch in dataloader:
|
||||
assert batch.size(0) == 10
|
||||
assert batch.size(1) == 1
|
||||
|
||||
permutation = permutation.with_format("torch_col")
|
||||
col_perm = permutation.with_format("torch_col")
|
||||
dataloader = torch.utils.data.DataLoader(
|
||||
permutation, collate_fn=lambda x: x, batch_size=10, shuffle=True
|
||||
col_perm, collate_fn=lambda x: x, batch_size=10, shuffle=True
|
||||
)
|
||||
for batch in dataloader:
|
||||
assert batch.size(0) == 1
|
||||
|
||||
Reference in New Issue
Block a user