# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The Geneva Authors

"""
Sentence Transformer UDF for text embedding.

This UDF generates embeddings from the caption text using sentence-transformers.
It piggybacks off the "caption" column generated by the BLIP UDF.
"""

from geneva.udfs.embeddings import sentence_transformer_udf

# Create the UDF that embeds captions
# This piggybacks off the "caption" column generated by BLIP
# Using dimension=384 for lazy loading (all-MiniLM-L6-v2 model)
caption_embedding = sentence_transformer_udf(
    model="sentence-transformers/all-MiniLM-L6-v2",
    column="caption",  # Use captions from BLIP
    normalize=True,
    trust_remote_code=True,
    dimension=384,  # Specify dimension to enable lazy loading
)
