HFSemanticClusterer

class probly.representer.semantic_clustering.huggingface.HFSemanticClusterer(model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, *, batch_size: int = 32, max_length: int | None = None, truncation: bool = True)[source]

Bases: Representer[Any, Any, Tensor, SemanticClusterOutput], ABC

Base semantic clusterer using a Hugging Face NLI model.

Initialize the semantic clusterer.

Parameters:

  • model – Hugging Face NLI sequence classification model.

  • tokenizer – Tokenizer associated with model.

  • batch_size – Number of NLI pairs to score in one model call.

  • max_length – Optional maximum pair sequence length passed to the tokenizer.

  • truncation – Whether pair inputs should be truncated to max_length or model defaults.

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

Alias for the represent method.

batch_size: int
classmethod from_model_name(model_name: str | None = None, *, cache_dir: str | PathLike[str] | None = None, force_download: bool = False, model_kwargs: Mapping[str, object] | None = None, tokenizer_kwargs: Mapping[str, object] | None = None, batch_size: int = 32, max_length: int | None = None, truncation: bool = True) → Self[source]

Load an NLI model by name and initialize the clusterer.

Parameters:

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

  • cache_dir – Optional Hugging Face cache directory.

  • force_download – Whether Hugging Face should re-download files even if cached.

  • model_kwargs – Additional keyword arguments forwarded to AutoModelForSequenceClassification.

  • tokenizer_kwargs – Additional keyword arguments forwarded to AutoTokenizer.

  • batch_size – Number of NLI pairs to score in one model call.

  • max_length – Optional maximum pair sequence length passed to the tokenizer.

  • truncation – Whether pair inputs should be truncated to max_length or model defaults.

Returns:

A semantic clusterer backed by the loaded NLI model and tokenizer.

max_length: int | None
model: PreTrainedModel
predict(*args: In.args, **kwargs: In.kwargs) → R[source]

Predict the representation for a given input.

property predictor: PreTrainedModel

The underlying NLI model used for pairwise entailment.

represent(generation: SemanticClusterInput, *, axis: int | None = None) → SemanticClusterOutput[source]

Cluster text generations into semantic equivalence classes.

Parameters:

  • generation – Text generation representation or sample.

  • axis – Comparison axis for raw TorchTextGeneration inputs. Ignored for samples.

Returns:

Sparse grouped logits whose final axis contains semantic cluster assignments.

tokenizer: PreTrainedTokenizerBase
truncation: bool