Sentence Transformer models run locally on your workers — there are no API calls and no
per-token costs. This makes them a good fit for large-scale embedding jobs where cost is a
concern.
Sentence Transformer models can run on CPU or GPU. Smaller models like all-MiniLM-L6-v2
work well on CPU, but larger models like bge-large-en-v1.5 benefit significantly from GPU
acceleration. Use the num_gpus parameter to request GPU resources for a worker:
# CPU-only (default) — suitable for lightweight modelssentence_transformer_udf(column="body", model="sentence-transformers/all-MiniLM-L6-v2")# GPU-accelerated — recommended for larger modelssentence_transformer_udf(column="body", model="BAAI/bge-large-en-v1.5", num_gpus=1.0)
Setting num_gpus to a fractional value (e.g., 0.5) tells the
Ray scheduler
to co-locate multiple workers on the same physical GPU. For example, two UDFs with
num_gpus=0.5 will be scheduled on a single GPU. Note that Ray does not enforce GPU memory
limits — it is your responsibility to ensure the combined models fit in GPU memory.