Optimal way to use batch func for vectorized operations

[conceptofmind]

Hi all,

I was wondering what the optimal way to use batch functions for vectorized operations is?

Say we want to count the total number of tokens in a dataset:

import daft
from daft import col, DataType, Series
from daft.functions import list_count, sum
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(
    "common-pile/comma-v0.1-2t", 
    local_files_only=True
)

@daft.func.batch(return_dtype=DataType.list(DataType.int64()))
def tokenize_batch(text_series: Series) -> Series:
    texts = text_series.to_pylist()
    encodings = tokenizer(texts, truncation=False, padding=False)
    return Series.from_pylist(encodings["input_ids"])


ddf = daft.read_json("/common-pile-comma_v0.1_training_dataset/arxiv_abstracts")

ddf = ddf.with_column("tokens", tokenize_batch(col("text")))

ddf = ddf.with_column("num_tokens", list_count(col("tokens")))

total_tokens = ddf.agg(sum(col("num_tokens")).alias("total_tokens")).collect()

print("Total tokens:", total_tokens)

Is this the correct way to do so?

Thanks!