torcharrow.DataFrame.to_tensor

DataFrame.to_tensor(conversion=None)

Convert to PyTorch containers (Tensor, PackedList, PackedMap, etc)

Parameters:

conversion (TensorConversion, or dict) – conversion can only be a dict type for DataFrame.to_tensor(). The dict maps from column name to the conversion methods. For column names not contained in dict, default PyTorch conversion will be used.

Examples

>>> import torcharrow as ta
>>> import torcharrow.pytorch as tap
>>> df = ta.dataframe({"label_ids": [0, 1], "token_ids": [[1, 2, 3, 4, 5], [101, 102]]})

>>> df
index    label_ids  token_ids
-------  -----------  ---------------
    0            0  [1, 2, 3, 4, 5]
    1            1  [101, 102]
dtype: Struct([Field('label_ids', int64), Field('token_ids', List(int64))]), count: 2, null_count: 0

>>> df.to_tensor({"token_ids": tap.PadSequence(padding_value=-1)})
TorchArrowStruct_0(
    label_ids=tensor([0, 1]),
    token_ids=tensor([
        [  1,   2,   3,   4,   5],
        [101, 102,  -1,  -1,  -1]]
    )
)