1from __future__ import annotations
2
3import os
4import daft
5from daft import col
6from daft.aihc.common.udf import aihc_udf
7from daft.aihc.functions.multimodal.embedding.clip_embedding import ClipEmbedding
8
9if __name__ == "__main__":
10    if os.getenv("DAFT_RUNNER", "native") == "ray":
11        import ray
12        ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
13        daft.set_runner_ray()
14    daft.set_execution_config(actor_udf_ready_timeout=6000, min_cpu_per_task=0)
15
16    model_path = os.getenv("MODEL_PATH", "/opt/aihc/models")
17    model_name = "iic/multi-modal_clip-vit-base-patch16_zh"
18    model_version = "v1.0.1"
19    batch_size = 2
20    rank = 0
21    num_gpus = 1
22
23    samples = {"text": ["树苗", "蜻蜓", "鸽子", "雪人", None]}
24    ds = daft.from_pydict(samples)
25    ds = ds.with_column(
26        "embedding",
27        aihc_udf(
28            ClipEmbedding,
29            construct_args={
30                "content_type": "text",
31                "model_path": model_path,
32                "model_name": model_name,
33                "model_version": model_version,
34                "batch_size": batch_size,
35                "rank": rank,
36            },
37            num_gpus=num_gpus,
38            batch_size=1,
39            concurrency=1,
40        )(col("text")),
41    )
42
43    ds.show()
44#╭────────┬────────────────────────────────╮                                                                                                                                                                    
45#│ text   ┆ embedding                      │
46#│ ---    ┆ ---                            │
47#│ String ┆ List[Float32]                  │
48#╞════════╪════════════════════════════════╡
49#│ 树苗   ┆ [0.089660645, -0.022903442, 0… │
50#├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
51#│ 蜻蜓   ┆ [0.09259033, -0.037902832, 0.… │
52#├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
53#│ 鸽子   ┆ [0.14916992, -0.032196045, 0.… │
54#├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
55#│ 雪人   ┆ [0.13623047, -0.0008497238, 0… │
56#├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
57#│ None   ┆ None                           │
58#╰────────┴────────────────────────────────╯
59
60    samples = {
61        "image_path": ["https://{bucket}.bj.bcebos.com/rgb.jpg"]
62    }
63    ds = daft.from_pydict(samples)
64    ds = ds.with_column(
65        "embedding",
66        aihc_udf(
67            ClipEmbedding,
68            construct_args={
69                "content_type": "image_url",
70                "model_path": model_path,
71                "model_name": model_name,
72                "model_version": model_version,
73                "batch_size": batch_size,
74                "rank": rank,
75            },
76            num_gpus=num_gpus,
77            batch_size=1,
78        )(col("image_path")),
79    )
80
81    ds.show()
82#╭────────────────────────────────┬────────────────────────────────╮                                                                                                                                            
83#│ image_path                     ┆ embedding                      │                                                                                                                                            
84#│ ---                            ┆ ---                            │
85#│ String                         ┆ List[Float32]                  │
86#╞════════════════════════════════╪════════════════════════════════╡
87#│ https://{bucket}.bj.bcebos.com/┆ [-0.021224976, -0.078063965, … │
88#╰────────────────────────────────┴────────────────────────────────╯