Main Classes

SetFitModel

class setfit.SetFitModel

( model_body: Optional = None model_head: Union = None multi_target_strategy: Optional = None normalize_embeddings: bool = False labels: Optional = None model_card_data: Optional = None sentence_transformers_kwargs: Optional = None **kwargs )

A SetFit model with integration to the Hugging Face Hub.

Example:

>>> from setfit import SetFitModel
>>> model = SetFitModel.from_pretrained("tomaarsen/setfit-bge-small-v1.5-sst2-8-shot")
>>> model.predict([
...     "It's a charming and often affecting journey.",
...     "It's slow -- very, very slow.",
...     "A sometimes tedious film.",
... ])
['positive', 'negative', 'negative']

from_pretrained

< source >

( force_download: bool = False resume_download: Optional = None proxies: Optional = None token: Union = None cache_dir: Union = None local_files_only: bool = False revision: Optional = None **model_kwargs )

Parameters

pretrained_model_name_or_path (str, Path) —
- Either the model_id (string) of a model hosted on the Hub, e.g. bigscience/bloom.
- Or a path to a directory containing model weights saved using [~transformers.PreTrainedModel.save_pretrained], e.g., ../path/to/my_model_directory/.
revision (str, optional) — Revision of the model on the Hub. Can be a branch name, a git tag or any commit id. Defaults to the latest commit on main branch.
force_download (bool, optional, defaults to False) — Whether to force (re-)downloading the model weights and configuration files from the Hub, overriding the existing cache.
proxies (Dict[str, str], optional) — A dictionary of proxy servers to use by protocol or endpoint, e.g., {‘http’: ‘foo.bar:3128’, ‘http://hostname’: ‘foo.bar:4012’}. The proxies are used on every request.
token (str or bool, optional) — The token to use as HTTP bearer authorization for remote files. By default, it will use the token cached when running huggingface-cli login.
cache_dir (str, Path, optional) — Path to the folder where cached files are stored.
local_files_only (bool, optional, defaults to False) — If True, avoid downloading the file and return the path to the local cached file if it exists.
labels (List[str], optional) — If the labels are integers ranging from 0 to num_classes-1, then these labels indicate the corresponding labels.
model_card_data (SetFitModelCardData, optional) — A SetFitModelCardData instance storing data such as model language, license, dataset name, etc. to be used in the automatically generated model cards.
multi_target_strategy (str, optional) — The strategy to use with multi-label classification. One of “one-vs-rest”, “multi-output”, or “classifier-chain”.
use_differentiable_head (bool, optional) — Whether to load SetFit using a differentiable (i.e., Torch) head instead of Logistic Regression.
normalize_embeddings (bool, optional) — Whether to apply normalization on the embeddings produced by the Sentence Transformer body.
device (Union[torch.device, str], optional) — The device on which to load the SetFit model, e.g. “cuda:0”, “mps” or torch.device(“cuda”).
trust_remote_code (bool, defaults to False) — Whether or not to allow for custom Sentence Transformers models defined on the Hub in their own modeling files. This option should only be set to True for repositories you trust and in which you have read the code, as it will execute code present on the Hub on your local machine. Defaults to False.

Download a model from the Huggingface Hub and instantiate it.

Example:

>>> from setfit import SetFitModel
>>> model = SetFitModel.from_pretrained(
...     "sentence-transformers/paraphrase-mpnet-base-v2",
...     labels=["positive", "negative"],
... )

save_pretrained

< source >

( save_directory: Union config: Union = None repo_id: Optional = None push_to_hub: bool = False model_card_kwargs: Optional = None **push_to_hub_kwargs ) → str or None

Parameters

save_directory (str or Path) — Path to directory in which the model weights and configuration will be saved.
config (dict or DataclassInstance, optional) — Model configuration specified as a key/value dictionary or a dataclass instance.
push_to_hub (bool, optional, defaults to False) — Whether or not to push your model to the Huggingface Hub after saving it.
repo_id (str, optional) — ID of your repository on the Hub. Used only if push_to_hub=True. Will default to the folder name if not provided.
model_card_kwargs (Dict[str, Any], optional) — Additional arguments passed to the model card template to customize the model card. push_to_hub_kwargs — Additional key word arguments passed along to the ~ModelHubMixin.push_to_hub method.

Returns

str or None

url of the commit on the Hub if push_to_hub=True, None otherwise.

Save weights in local directory.

push_to_hub

< source >

( repo_id: str config: Union = None commit_message: str = 'Push model using huggingface_hub.' private: bool = False token: Optional = None branch: Optional = None create_pr: Optional = None allow_patterns: Union = None ignore_patterns: Union = None delete_patterns: Union = None model_card_kwargs: Optional = None )

Parameters

repo_id (str) — ID of the repository to push to (example: "username/my-model").
config (dict or DataclassInstance, optional) — Model configuration specified as a key/value dictionary or a dataclass instance.
commit_message (str, optional) — Message to commit while pushing.
private (bool, optional, defaults to False) — Whether the repository created should be private.
token (str, optional) — The token to use as HTTP bearer authorization for remote files. By default, it will use the token cached when running huggingface-cli login.
branch (str, optional) — The git branch on which to push the model. This defaults to "main".
create_pr (boolean, optional) — Whether or not to create a Pull Request from branch with that commit. Defaults to False.
allow_patterns (List[str] or str, optional) — If provided, only files matching at least one pattern are pushed.
ignore_patterns (List[str] or str, optional) — If provided, files matching any of the patterns are not pushed.
delete_patterns (List[str] or str, optional) — If provided, remote files matching any of the patterns will be deleted from the repo.
model_card_kwargs (Dict[str, Any], optional) — Additional arguments passed to the model card template to customize the model card.

Upload model checkpoint to the Hub.

Use allow_patterns and ignore_patterns to precisely filter which files should be pushed to the hub. Use delete_patterns to delete existing remote files in the same commit. See upload_folder reference for more details.

call

< source >

( inputs: Union batch_size: int = 32 as_numpy: bool = False use_labels: bool = True show_progress_bar: Optional = None ) → Union[torch.Tensor, np.ndarray, List[str], int, str]

Parameters

inputs (Union[str, List[str]]) — The input sentence or sentences to predict classes for.
batch_size (int, defaults to 32) — The batch size to use in encoding the sentences to embeddings. Higher often means faster processing but higher memory usage.
as_numpy (bool, defaults to False) — Whether to output as numpy array instead.
use_labels (bool, defaults to True) — Whether to try and return elements of SetFitModel.labels.
show_progress_bar (Optional[bool], defaults to None) — Whether to show a progress bar while encoding.

Returns

Union[torch.Tensor, np.ndarray, List[str], int, str]

A list of string labels with equal length to the inputs if use_labels is True and SetFitModel.labels has been defined. Otherwise a vector with equal length to the inputs, denoting to which class each input is predicted to belong. If the inputs is a single string, then the output is a single label as well.

Predict the various classes.

Example:

>>> model = SetFitModel.from_pretrained(...)
>>> model(["What a boring display", "Exhilarating through and through", "I'm wowed!"])
["negative", "positive", "positive"]
>>> model("That was cool!")
"positive"

label2id

< source >

( )

Return a mapping from string labels to integer IDs.

id2label

< source >

( )

Return a mapping from integer IDs to string labels.

create_model_card

< source >

( path: str model_name: Optional = '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.

encode

< source >

( inputs: List batch_size: int = 32 show_progress_bar: Optional = None ) → Union[torch.Tensor, np.ndarray]

Parameters

inputs (List[str]) — The input sentences to embed.
batch_size (int, defaults to 32) — The batch size to use in encoding the sentences to embeddings. Higher often means faster processing but higher memory usage.
show_progress_bar (Optional[bool], defaults to None) — Whether to show a progress bar while encoding.

Returns

Union[torch.Tensor, np.ndarray]

A matrix with shape [INPUT_LENGTH, EMBEDDING_SIZE], as a torch Tensor if this model has a differentiable Torch head, or otherwise as a numpy array.

Convert input sentences to embeddings using the SentenceTransformer body.

fit

< source >

( x_train: List y_train: Union num_epochs: int batch_size: Optional = None body_learning_rate: Optional = None head_learning_rate: Optional = None end_to_end: bool = False l2_weight: Optional = None max_length: Optional = None show_progress_bar: bool = True )

Parameters

x_train (List[str]) — A list of training sentences.
y_train (Union[List[int], List[List[int]]]) — A list of labels corresponding to the training sentences.
num_epochs (int) — The number of epochs to train for.
batch_size (int, optional) — The batch size to use.
body_learning_rate (float, optional) — The learning rate for the SentenceTransformer body in the AdamW optimizer. Disregarded if end_to_end=False.
head_learning_rate (float, optional) — The learning rate for the differentiable torch head in the AdamW optimizer.
end_to_end (bool, defaults to False) — If True, train the entire model end-to-end. Otherwise, freeze the SentenceTransformer body and only train the head.
l2_weight (float, optional) — The l2 weight for both the model body and head in the AdamW optimizer.
max_length (int, optional) — The maximum token length a tokenizer can generate. If not provided, the maximum length for the SentenceTransformer body is used.
show_progress_bar (bool, defaults to True) — Whether to display a progress bar for the training epochs and iterations.

Train the classifier head, only used if a differentiable PyTorch head is used.

freeze

< source >

( component: Optional = None )

Parameters

component (Literal["body", "head"], optional) — Either “body” or “head” to freeze that component. If no component is provided, freeze both. Defaults to None.

Freeze the model body and/or the head, preventing further training on that component until unfrozen.

generate_model_card

< source >

( ) → str

Returns

str

The model card string.

Generate and return a model card string based on the model card data.

predict

< source >

( inputs: Union batch_size: int = 32 as_numpy: bool = False use_labels: bool = True show_progress_bar: Optional = None ) → Union[torch.Tensor, np.ndarray, List[str], int, str]

Parameters

inputs (Union[str, List[str]]) — The input sentence or sentences to predict classes for.
batch_size (int, defaults to 32) — The batch size to use in encoding the sentences to embeddings. Higher often means faster processing but higher memory usage.
as_numpy (bool, defaults to False) — Whether to output as numpy array instead.
use_labels (bool, defaults to True) — Whether to try and return elements of SetFitModel.labels.
show_progress_bar (Optional[bool], defaults to None) — Whether to show a progress bar while encoding.

Returns

Union[torch.Tensor, np.ndarray, List[str], int, str]

Predict the various classes.

Example:

>>> model = SetFitModel.from_pretrained(...)
>>> model.predict(["What a boring display", "Exhilarating through and through", "I'm wowed!"])
["negative", "positive", "positive"]
>>> model.predict("That was cool!")
"positive"

predict_proba

< source >

( inputs: Union batch_size: int = 32 as_numpy: bool = False show_progress_bar: Optional = None ) → Union[torch.Tensor, np.ndarray]

Parameters

inputs (Union[str, List[str]]) — The input sentences to predict class probabilities for.
batch_size (int, defaults to 32) — The batch size to use in encoding the sentences to embeddings. Higher often means faster processing but higher memory usage.
as_numpy (bool, defaults to False) — Whether to output as numpy array instead.
show_progress_bar (Optional[bool], defaults to None) — Whether to show a progress bar while encoding.

Returns

Union[torch.Tensor, np.ndarray]

A matrix with shape [INPUT_LENGTH, NUM_CLASSES] denoting probabilities of predicting an input as a class. If the input is a string, then the output is a vector with shape [NUM_CLASSES,].

Predict the probabilities of the various classes.

Example:

>>> model = SetFitModel.from_pretrained(...)
>>> model.predict_proba(["What a boring display", "Exhilarating through and through", "I'm wowed!"])
tensor([[0.9367, 0.0633],
[0.0627, 0.9373],
[0.0890, 0.9110]], dtype=torch.float64)
>>> model.predict_proba("That was cool!")
tensor([0.8421, 0.1579], dtype=torch.float64)

to

< source >

( device: Union ) → 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.

Example:

>>> model = SetFitModel.from_pretrained(...)
>>> model.to("cpu")
>>> model(["cats are cute", "dogs are loyal"])

unfreeze

< source >

( component: Optional = None keep_body_frozen: Optional = None )

Parameters

component (Literal["body", "head"], optional) — Either “body” or “head” to unfreeze that component. If no component is provided, unfreeze both. Defaults to None.
keep_body_frozen (bool, optional) — Deprecated argument, use component instead.

Unfreeze the model body and/or the head, allowing further training on that component.

SetFitHead

class setfit.SetFitHead

< source >

( in_features: Optional = None out_features: int = 2 temperature: float = 1.0 eps: float = 1e-05 bias: bool = True device: Union = 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 make the model less confident and lower values make 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: Union temperature: Optional = 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 make the model less confident and lower values make it more confident. Will override the temperature given during initialization.

SetFitHead can accept embeddings in:

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

SetFitModelCardData

class setfit.SetFitModelCardData

< source >

( language: Union = None license: Optional = None tags: Optional = <factory> model_name: Optional = None model_id: Optional = None dataset_name: Optional = None dataset_id: Optional = None dataset_revision: Optional = None task_name: Optional = None st_id: Optional = None )

Parameters

language (Optional[Union[str, List[str]]]) — The model language, either a string or a list, e.g. “en” or [“en”, “de”, “nl”]
license (Optional[str]) — The license of the model, e.g. “apache-2.0”, “mit”, or “cc-by-nc-sa-4.0”
model_name (Optional[str]) — The pretty name of the model, e.g. “SetFit with mBERT-base on SST2”. If not defined, uses encoder_name/encoder_id and dataset_name/dataset_id to generate a model name.
model_id (Optional[str]) — The model ID when pushing the model to the Hub, e.g. “tomaarsen/span-marker-mbert-base-multinerd”.
dataset_name (Optional[str]) — The pretty name of the dataset, e.g. “SST2”.
dataset_id (Optional[str]) — The dataset ID of the dataset, e.g. “dair-ai/emotion”.
dataset_revision (Optional[str]) — The dataset revision/commit that was for training/evaluation.
st_id (Optional[str]) — The Sentence Transformers model ID.

A dataclass storing data used in the model card.

Install codecarbon to automatically track carbon emission usage and include it in your model cards.

Example:

>>> model = SetFitModel.from_pretrained(
...     "sentence-transformers/paraphrase-mpnet-base-v2",
...     labels=["negative", "positive"],
...     # Model card variables
...     model_card_data=SetFitModelCardData(
...         model_id="tomaarsen/setfit-paraphrase-mpnet-base-v2-sst2",
...         dataset_name="SST2",
...         dataset_id="sst2",
...         license="apache-2.0",
...         language="en",
...     ),
... )

to_dict

< source >

( )

to_yaml

< source >

( line_break = None )

AbsaModel

class setfit.AbsaModel

< source >

( aspect_extractor: AspectExtractor aspect_model: AspectModel polarity_model: PolarityModel )

call

< source >

( inputs: Union )

device

< source >

( )

from_pretrained

< source >

( model_id: str polarity_model_id: Optional = None spacy_model: Optional = None span_contexts: Tuple = (None, None) force_download: bool = None resume_download: bool = None proxies: Optional = None token: Union = None cache_dir: Optional = None local_files_only: bool = None use_differentiable_head: bool = None normalize_embeddings: bool = None **model_kwargs )

predict

< source >

( inputs: Union ) → Union[List[Dict[str, Any]], Dataset]

Parameters

inputs (Union[str, List[str], Dataset]) — Either a sentence, a list of sentences, or a dataset with columns text and span and optionally ordinal. This dataset contains gold aspects, and we only predict the polarities for them.

Returns

Union[List[Dict[str, Any]], Dataset]

Either a list of dictionaries with keys span and polarity if the input was a sentence or a list of sentences, or a dataset with columns text, span, ordinal, and pred_polarity.

Predicts aspects & their polarities of the given inputs.

Example:

>>> from setfit import AbsaModel
>>> model = AbsaModel.from_pretrained(
...     "tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants-aspect",
...     "tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants-polarity",
... )
>>> model.predict("The food and wine are just exquisite.")
[{'span': 'food', 'polarity': 'positive'}, {'span': 'wine', 'polarity': 'positive'}]

>>> from setfit import AbsaModel
>>> from datasets import load_dataset
>>> model = AbsaModel.from_pretrained(
...     "tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants-aspect",
...     "tomaarsen/setfit-absa-bge-small-en-v1.5-restaurants-polarity",
... )
>>> dataset = load_dataset("tomaarsen/setfit-absa-semeval-restaurants", split="train")
>>> model.predict(dataset)
Dataset({
features: ['text', 'span', 'label', 'ordinal', 'pred_polarity'],
num_rows: 3693
})

push_to_hub

< source >

( repo_id: str polarity_repo_id: Optional = None **kwargs )

to

< source >

( device: Union )

save_pretrained

< source >

( save_directory: Union polarity_save_directory: Union = None push_to_hub: bool = False **kwargs )

AspectModel

class setfit.AspectModel

< source >

( *args **kwargs )

call

< source >

( docs: List aspects_list: List )

device

< source >

( ) → torch.device

Returns

torch.device

The device that the model is on.

Get the Torch device that this model is on.

from_pretrained

< source >

( )

Parameters

pretrained_model_name_or_path (str, Path) —
- Either the model_id (string) of a model hosted on the Hub, e.g. bigscience/bloom.
- Or a path to a directory containing model weights saved using save_pretrained, e.g., ../path/to/my_model_directory/.
revision (str, optional) — Revision of the model on the Hub. Can be a branch name, a git tag or any commit id. Defaults to the latest commit on main branch.
force_download (bool, optional, defaults to False) — Whether to force (re-)downloading the model weights and configuration files from the Hub, overriding the existing cache.
proxies (Dict[str, str], optional) — A dictionary of proxy servers to use by protocol or endpoint, e.g., {'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}. The proxies are used on every request.
token (str or bool, optional) — The token to use as HTTP bearer authorization for remote files. By default, it will use the token cached when running huggingface-cli login.
cache_dir (str, Path, optional) — Path to the folder where cached files are stored.
local_files_only (bool, optional, defaults to False) — If True, avoid downloading the file and return the path to the local cached file if it exists.
labels (List[str], optional) — If the labels are integers ranging from 0 to num_classes-1, then these labels indicate the corresponding labels.
model_card_data (SetFitModelCardData, optional) — A SetFitModelCardData instance storing data such as model language, license, dataset name, etc. to be used in the automatically generated model cards.
model_card_data (SetFitModelCardData, optional) — A SetFitModelCardData instance storing data such as model language, license, dataset name, etc. to be used in the automatically generated model cards.
use_differentiable_head (bool, optional) — Whether to load SetFit using a differentiable (i.e., Torch) head instead of Logistic Regression.
normalize_embeddings (bool, optional) — Whether to apply normalization on the embeddings produced by the Sentence Transformer body.
span_context (int, defaults to 0) — The number of words before and after the span candidate that should be prepended to the full sentence. By default, 0 for Aspect models and 3 for Polarity models.
device (Union[torch.device, str], optional) — The device on which to load the SetFit model, e.g. "cuda:0", "mps" or torch.device("cuda").

Download a model from the Huggingface Hub and instantiate it.

predict

< source >

( inputs: Union batch_size: int = 32 as_numpy: bool = False use_labels: bool = True show_progress_bar: Optional = None ) → Union[torch.Tensor, np.ndarray, List[str], int, str]

Parameters

inputs (Union[str, List[str]]) — The input sentence or sentences to predict classes for.
batch_size (int, defaults to 32) — The batch size to use in encoding the sentences to embeddings. Higher often means faster processing but higher memory usage.
as_numpy (bool, defaults to False) — Whether to output as numpy array instead.
use_labels (bool, defaults to True) — Whether to try and return elements of SetFitModel.labels.
show_progress_bar (Optional[bool], defaults to None) — Whether to show a progress bar while encoding.

Returns

Union[torch.Tensor, np.ndarray, List[str], int, str]

Predict the various classes.

Example:

>>> model = SetFitModel.from_pretrained(...)
>>> model.predict(["What a boring display", "Exhilarating through and through", "I'm wowed!"])
["negative", "positive", "positive"]
>>> model.predict("That was cool!")
"positive"

push_to_hub

< source >

Parameters

repo_id (str) — ID of the repository to push to (example: "username/my-model").
config (dict or DataclassInstance, optional) — Model configuration specified as a key/value dictionary or a dataclass instance.
commit_message (str, optional) — Message to commit while pushing.
private (bool, optional, defaults to False) — Whether the repository created should be private.
token (str, optional) — The token to use as HTTP bearer authorization for remote files. By default, it will use the token cached when running huggingface-cli login.
branch (str, optional) — The git branch on which to push the model. This defaults to "main".
create_pr (boolean, optional) — Whether or not to create a Pull Request from branch with that commit. Defaults to False.
allow_patterns (List[str] or str, optional) — If provided, only files matching at least one pattern are pushed.
ignore_patterns (List[str] or str, optional) — If provided, files matching any of the patterns are not pushed.
delete_patterns (List[str] or str, optional) — If provided, remote files matching any of the patterns will be deleted from the repo.
model_card_kwargs (Dict[str, Any], optional) — Additional arguments passed to the model card template to customize the model card.

Upload model checkpoint to the Hub.

save_pretrained

< source >

( save_directory: Union config: Union = None repo_id: Optional = None push_to_hub: bool = False model_card_kwargs: Optional = None **push_to_hub_kwargs ) → str or None

Parameters

save_directory (str or Path) — Path to directory in which the model weights and configuration will be saved.
config (dict or DataclassInstance, optional) — Model configuration specified as a key/value dictionary or a dataclass instance.
push_to_hub (bool, optional, defaults to False) — Whether or not to push your model to the Huggingface Hub after saving it.
repo_id (str, optional) — ID of your repository on the Hub. Used only if push_to_hub=True. Will default to the folder name if not provided.
model_card_kwargs (Dict[str, Any], optional) — Additional arguments passed to the model card template to customize the model card. push_to_hub_kwargs — Additional key word arguments passed along to the ~ModelHubMixin.push_to_hub method.

Returns

str or None

url of the commit on the Hub if push_to_hub=True, None otherwise.

Save weights in local directory.

to

< source >

( device: Union ) → 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.

Example:

>>> model = SetFitModel.from_pretrained(...)
>>> model.to("cpu")
>>> model(["cats are cute", "dogs are loyal"])

PolarityModel

class setfit.PolarityModel

< source >

( span_context: int = 3 **kwargs )

call

< source >

( docs: List aspects_list: List )

device

< source >

( ) → torch.device

Returns

torch.device

The device that the model is on.

Get the Torch device that this model is on.

from_pretrained

< source >

( )

Parameters

pretrained_model_name_or_path (str, Path) —
- Either the model_id (string) of a model hosted on the Hub, e.g. bigscience/bloom.
- Or a path to a directory containing model weights saved using save_pretrained, e.g., ../path/to/my_model_directory/.
revision (str, optional) — Revision of the model on the Hub. Can be a branch name, a git tag or any commit id. Defaults to the latest commit on main branch.
force_download (bool, optional, defaults to False) — Whether to force (re-)downloading the model weights and configuration files from the Hub, overriding the existing cache.
proxies (Dict[str, str], optional) — A dictionary of proxy servers to use by protocol or endpoint, e.g., {'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}. The proxies are used on every request.
token (str or bool, optional) — The token to use as HTTP bearer authorization for remote files. By default, it will use the token cached when running huggingface-cli login.
cache_dir (str, Path, optional) — Path to the folder where cached files are stored.
local_files_only (bool, optional, defaults to False) — If True, avoid downloading the file and return the path to the local cached file if it exists.
labels (List[str], optional) — If the labels are integers ranging from 0 to num_classes-1, then these labels indicate the corresponding labels.
model_card_data (SetFitModelCardData, optional) — A SetFitModelCardData instance storing data such as model language, license, dataset name, etc. to be used in the automatically generated model cards.
model_card_data (SetFitModelCardData, optional) — A SetFitModelCardData instance storing data such as model language, license, dataset name, etc. to be used in the automatically generated model cards.
use_differentiable_head (bool, optional) — Whether to load SetFit using a differentiable (i.e., Torch) head instead of Logistic Regression.
normalize_embeddings (bool, optional) — Whether to apply normalization on the embeddings produced by the Sentence Transformer body.
span_context (int, defaults to 0) — The number of words before and after the span candidate that should be prepended to the full sentence. By default, 0 for Aspect models and 3 for Polarity models.
device (Union[torch.device, str], optional) — The device on which to load the SetFit model, e.g. "cuda:0", "mps" or torch.device("cuda").

Download a model from the Huggingface Hub and instantiate it.

predict

< source >

( inputs: Union batch_size: int = 32 as_numpy: bool = False use_labels: bool = True show_progress_bar: Optional = None ) → Union[torch.Tensor, np.ndarray, List[str], int, str]

Parameters

inputs (Union[str, List[str]]) — The input sentence or sentences to predict classes for.
batch_size (int, defaults to 32) — The batch size to use in encoding the sentences to embeddings. Higher often means faster processing but higher memory usage.
as_numpy (bool, defaults to False) — Whether to output as numpy array instead.
use_labels (bool, defaults to True) — Whether to try and return elements of SetFitModel.labels.
show_progress_bar (Optional[bool], defaults to None) — Whether to show a progress bar while encoding.

Returns

Union[torch.Tensor, np.ndarray, List[str], int, str]

Predict the various classes.

Example:

>>> model = SetFitModel.from_pretrained(...)
>>> model.predict(["What a boring display", "Exhilarating through and through", "I'm wowed!"])
["negative", "positive", "positive"]
>>> model.predict("That was cool!")
"positive"

push_to_hub

< source >

Parameters

repo_id (str) — ID of the repository to push to (example: "username/my-model").
config (dict or DataclassInstance, optional) — Model configuration specified as a key/value dictionary or a dataclass instance.
commit_message (str, optional) — Message to commit while pushing.
private (bool, optional, defaults to False) — Whether the repository created should be private.
token (str, optional) — The token to use as HTTP bearer authorization for remote files. By default, it will use the token cached when running huggingface-cli login.
branch (str, optional) — The git branch on which to push the model. This defaults to "main".
create_pr (boolean, optional) — Whether or not to create a Pull Request from branch with that commit. Defaults to False.
allow_patterns (List[str] or str, optional) — If provided, only files matching at least one pattern are pushed.
ignore_patterns (List[str] or str, optional) — If provided, files matching any of the patterns are not pushed.
delete_patterns (List[str] or str, optional) — If provided, remote files matching any of the patterns will be deleted from the repo.
model_card_kwargs (Dict[str, Any], optional) — Additional arguments passed to the model card template to customize the model card.

Upload model checkpoint to the Hub.

save_pretrained

< source >

( save_directory: Union config: Union = None repo_id: Optional = None push_to_hub: bool = False model_card_kwargs: Optional = None **push_to_hub_kwargs ) → str or None

Parameters

save_directory (str or Path) — Path to directory in which the model weights and configuration will be saved.
config (dict or DataclassInstance, optional) — Model configuration specified as a key/value dictionary or a dataclass instance.
push_to_hub (bool, optional, defaults to False) — Whether or not to push your model to the Huggingface Hub after saving it.
repo_id (str, optional) — ID of your repository on the Hub. Used only if push_to_hub=True. Will default to the folder name if not provided.
model_card_kwargs (Dict[str, Any], optional) — Additional arguments passed to the model card template to customize the model card. push_to_hub_kwargs — Additional key word arguments passed along to the ~ModelHubMixin.push_to_hub method.

Returns

str or None

url of the commit on the Hub if push_to_hub=True, None otherwise.

Save weights in local directory.

to

< source >

( device: Union ) → 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.

Example:

>>> model = SetFitModel.from_pretrained(...)
>>> model.to("cpu")
>>> model(["cats are cute", "dogs are loyal"])

< > Update on GitHub

SetFit

Main Classes

SetFitModel

class setfit.SetFitModel

from_pretrained

save_pretrained

push_to_hub

__call__

label2id

id2label

create_model_card

encode

fit

freeze

generate_model_card

predict

predict_proba

to

unfreeze

SetFitHead

class setfit.SetFitHead

forward

SetFitModelCardData

class setfit.SetFitModelCardData

to_dict

to_yaml

AbsaModel

class setfit.AbsaModel

__call__

device

from_pretrained

predict

push_to_hub

to

save_pretrained

AspectModel

class setfit.AspectModel

__call__

device

from_pretrained

predict

push_to_hub

save_pretrained

to

PolarityModel

class setfit.PolarityModel

__call__

device

from_pretrained

predict

push_to_hub

save_pretrained

to

call

call

call

call