SetFitModel

class setfit.SetFitModel

( model_body: typing.Optional[sentence_transformers.SentenceTransformer.SentenceTransformer] = None model_head: typing.Union[setfit.modeling.SetFitHead, sklearn.linear_model._logistic.LogisticRegression, NoneType] = None multi_target_strategy: typing.Optional[str] = None l2_weight: float = 0.01 normalize_embeddings: bool = False )

A SetFit model with integration to the Hugging Face Hub.

create_model_card

< source >

( path: str model_name: typing.Optional[str] = 'SetFit Model' )

Parameters

path (str) — The path to save the model card to.
model_name (str, optional) — The name of the model. Defaults to SetFit Model.

Creates and saves a model card for a SetFit model.

to

< source >

( device: typing.Union[str, torch.device] ) → SetFitModel

Parameters

device (Union[str, torch.device]) — The identifier of the device to move the model to.

Returns

SetFitModel

Returns the original model, but now on the desired device.

Move this SetFitModel to device, and then return self. This method does not copy.

SetFitHead

class setfit.SetFitHead

< source >

( in_features: typing.Optional[int] = None out_features: int = 2 temperature: float = 1.0 eps: float = 1e-05 bias: bool = True device: typing.Union[torch.device, str, NoneType] = None multitarget: bool = False )

Parameters

in_features (int, optional) — The embedding dimension from the output of the SetFit body. If None, defaults to LazyLinear.
out_features (int, defaults to 2) — The number of targets. If set out_features to 1 for binary classification, it will be changed to 2 as 2-class classification.
temperature (float, defaults to 1.0) — A logits’ scaling factor. Higher values makes the model less confident and higher values makes it more confident.
eps (float, defaults to 1e-5) — A value for numerical stability when scaling logits.
bias (bool, optional, defaults to True) — Whether to add bias to the head.
device (torch.device, str, optional) — The device the model will be sent to. If None, will check whether GPU is available.
multitarget (bool, defaults to False) — Enable multi-target classification by making out_features binary predictions instead of a single multinomial prediction.

A SetFit head that supports multi-class classification for end-to-end training. Binary classification is treated as 2-class classification.

To be compatible with Sentence Transformers, we inherit Dense from: https://github.com/UKPLab/sentence-transformers/blob/master/sentence_transformers/models/Dense.py

forward

< source >

( features: typing.Union[typing.Dict[str, torch.Tensor], torch.Tensor] temperature: typing.Optional[float] = None )

Parameters

features (Dict[str, torch.Tensor] or torch.Tensor) -- The embeddings from the encoder. If using dict` format, make sure to store embeddings under the key: ‘sentence_embedding’ and the outputs will be under the key: ‘prediction’.
temperature (float, optional) — A logits’ scaling factor. Higher values makes the model less confident and higher values makes it more confident. Will override the temperature given during initialization.

SetFitHead can accept embeddings in:

Output format (dict) from Sentence-Transformers.
Pure torch.Tensor.