Tokenize batch output dtype not matching

[hiskuDN]

The _tokenize_batch method in the dataloader.py fails when processing the batch[col_name] field, as its type is pyarrow.ListArray. The .to() function is called on this field, which results in an error because pyarrow.ListArray does not support the .to() method (unless I'm missing something in the implementation, I'm using the env setup provided in the repo)

Expected Behavior

The batch[col_name] field should be properly processed as a torch.Tensor or converted to one before calling .to().

Observed Behavior

The code throws an error when attempting to execute .to(self.gang.device) because pyarrow.ListArray does not have a .to() method.

Code Snippet

embs = [x.to(self.gang.device).to(dtype) for x in batch[col_name]]

Potential Solution

The batch[col_name] field may need to be explicitly converted to a torch.Tensor before calling .to(). For example:

embs = [torch.Tensor(x.as_py()).to(self.gang.device).to(dtype) for x in batch[col_name]]

[hasanyazarr]

Hi, I also encounter this problem and added a fix in pr #14.

[artemru]

can you show pa.Schema of the dataset you are using ?
normally, the correct conversion (pyarrow.ListArray[FixedSizeListArray] -> List[numpy.ndarray]) should happen around https://github.com/facebookresearch/large_concept_model/blob/main/lcm/datasets/parquet_utils.py#L181
If it's not the case, it could mean that the embeddings columns are not properly typed (it should be ListArray(FixedSizeListArray(dtype, 1024))).

[artemru]

if you're using pandas for data handling at some point, be aware of pandas-dev/pandas#59738

[hiskuDN]

@artemru yes, I am using pandas for data handling, I believe that might be the core issue. I'll see if I can follow the fix on pandas-dev/pandas#59738, if there is any.

[artemru]

I would recommend to use pyarrow directly or polars, where you can cast the dataset to the correct type with .cast method!

[artemru]

does it solve your issue ? do you think we can close it ?

[hiskuDN]

Ah yes yes, sorry for not closing it. I had an issue with data handling and I needed to cast it.