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!

[universalmind303]

Hey @conceptofmind , That looks like a pretty efficient approach to me. Are you running into performance issues using this approach?

[conceptofmind]

I would not necessarily say performance issues, but I am always looking to optimize the current pipeline to be more efficient at scale haha. Did not know if there was a better way to do so than above.

[conceptofmind]

Or maybe if we want to return attention masks too

@daft.func.batch(return_dtype=DataType.struct({
    "input_ids": DataType.list(DataType.int64()),
    "attention_mask": DataType.list(DataType.int64()),
}))
def tokenize_batch(text_series: Series) -> Series:
    texts = text_series.to_pylist()
    encodings = tokenizer(texts, truncation=False, padding=False)
    combined = [
        {"input_ids": ids, "attention_mask": mask}
        for ids, mask in zip(encodings["input_ids"], encodings["attention_mask"])
    ]
    return Series.from_pylist(combined)

[kevinzwang]

I think the main way to further optimize would be to try tuning the batch_size parameter.

We also recently introduced dynamic batch sizing, where Daft can automatically optimize the batch size for performance. You can try it by enabling the feature:

daft.context.set_execution_config(enable_dynamic_batching=True)

[conceptofmind]

I think the main way to further optimize would be to try tuning the batch_size parameter.

We also recently introduced dynamic batch sizing, where Daft can automatically optimize the batch size for performance. You can try it by enabling the feature:

daft.context.set_execution_config(enable_dynamic_batching=True)

I will test this out tomorrow. Thanks

[conceptofmind]

Going to close this for now since I am just running with the automatic batching