1from __future__ import annotations
2
3import os
4
5import daft
6from daft import col
7
8from daft.aihc.common.udf import aihc_udf
9from daft.aihc.functions.text.token_count_calculator import TokenCountCalculator
10
11if __name__ == "__main__":
12    if os.getenv("DAFT_RUNNER", "native") == "ray":
13        import ray
14        ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
15        daft.set_runner_ray()
16    daft.set_execution_config(actor_udf_ready_timeout=6000, min_cpu_per_task=0)
17
18samples = {
19        "text": [
20            "Hello world, this is a tokenizer test.",
21            "百度智能云 AIHC 数据处理算子。",
22            "",
23        ]
24    }
25    ds = daft.from_pydict(samples)
26    ds = ds.with_column(
27        "num_token",
28        aihc_udf(
29            TokenCountCalculator,
30            construct_args={
31                "model_path": "/path/to/models",
32                "model_name": "EleutherAI/pythia-6.9b-deduped",
33            },
34            num_cpus=1,
35            concurrency=1,
36            batch_size=64,
37        )(col("text")),
38    )
39    ds.show()