HFTextEmbedder

class probly.representer.embedder.huggingface.HFTextEmbedder(model: SentenceTransformerLike, *, batch_size: int = 32, normalize_embeddings: bool = True, encode_kwargs: Mapping[str, object] | None = None)[source]

Bases: Representer[Any, Any, Tensor, TextEmbedOutput]

Embed text generations using a Hugging Face sentence-transformers model.

Initialize the text embedder.

Parameters:

  • model – Sentence-transformer compatible embedding model.

  • batch_size – Number of texts to embed in one model call.

  • normalize_embeddings – Whether sentence-transformers should L2-normalize embeddings.

  • encode_kwargs – Additional keyword arguments forwarded to model.encode.

__call__(*args: In.args, **kwargs: In.kwargs) → R[source]

Alias for the represent method.

batch_size: int
encode_kwargs: dict[str, object]
classmethod from_model_name(model_name: str | None = None, *, cache_dir: str | PathLike[str] | None = None, model_kwargs: Mapping[str, object] | None = None, batch_size: int = 32, normalize_embeddings: bool = True, encode_kwargs: Mapping[str, object] | None = None) → Self[source]

Load a sentence-transformers model by name and initialize the embedder.

Parameters:

  • model_name – Hugging Face model name or local path. Defaults to DEFAULT_EMBEDDING_MODEL.

  • cache_dir – Optional Hugging Face cache directory.

  • model_kwargs – Additional keyword arguments forwarded to SentenceTransformer.

  • batch_size – Number of texts to embed in one model call.

  • normalize_embeddings – Whether sentence-transformers should L2-normalize embeddings.

  • encode_kwargs – Additional keyword arguments forwarded to model.encode.

Returns:

A text embedder backed by the loaded sentence-transformers model.

model: SentenceTransformerLike
normalize_embeddings: bool
predict(*args: In.args, **kwargs: In.kwargs) → R[source]

Predict the representation for a given input.

property predictor: SentenceTransformerLike

The underlying embedding model.

represent(generation: TextEmbedInput) → TextEmbedOutput[source]

Embed text generations.

Parameters:

generation – Text generation representation or sample.

Returns:

Embeddings with the same sample wrapping as generation.