Function get_text_embedding

utils/src/model_utils.py:143–166

Generate text embeddings for a list of text strings. Args: text (list[str]): A list of text strings. model: The model used for generating embeddings. batchsize (int): The batch size for processing text. Returns: list: A list of text embeddings.

(
    text: list[str], model: BGEM3FlagModel, batchsize: int = 32
)

def get_text_embedding(
    text: list[str], model: BGEM3FlagModel, batchsize: int = 32
) -> list[torch.Tensor]:
    """
    Generate text embeddings for a list of text strings.

    Args:
        text (list[str]): A list of text strings.
        model: The model used for generating embeddings.
        batchsize (int): The batch size for processing text.

    Returns:
        list: A list of text embeddings.
    """
    if isinstance(text, str):
        return torch.tensor(model.encode(text)["dense_vecs"]).to(model.device)
    result = []
    for i in range(0, len(text), batchsize):
        result.extend(
            torch.tensor(model.encode(text[i : i + batchsize])["dense_vecs"]).to(
                model.device
            )
        )
    return result


def get_image_embedding(