Enum rust_bert::pipelines::zero_shot_classification::ZeroShotClassificationOption

pub enum ZeroShotClassificationOption {
    Bart(BartForSequenceClassification),
    Deberta(DebertaForSequenceClassification),
    DebertaV2(DebertaV2ForSequenceClassification),
    Bert(BertForSequenceClassification),
    DistilBert(DistilBertModelClassifier),
    MobileBert(MobileBertForSequenceClassification),
    Roberta(RobertaForSequenceClassification),
    XLMRoberta(RobertaForSequenceClassification),
    Albert(AlbertForSequenceClassification),
    XLNet(XLNetForSequenceClassification),
    Longformer(LongformerForSequenceClassification),
}

Abstraction that holds one particular zero shot classification model, for any of the supported models

The models are using a classification architecture that should be trained on Natural Language Inference. The models should output a Tensor of size > 2 in the label dimension, with the first logit corresponding to contradiction and the last logit corresponding to entailment.

Variants

Bart(BartForSequenceClassification)

Bart for Sequence Classification

Deberta(DebertaForSequenceClassification)

DeBERTa for Sequence Classification

DebertaV2(DebertaV2ForSequenceClassification)

DeBERTaV2 for Sequence Classification

Bert(BertForSequenceClassification)

Bert for Sequence Classification

DistilBert(DistilBertModelClassifier)

DistilBert for Sequence Classification

MobileBert(MobileBertForSequenceClassification)

MobileBert for Sequence Classification

Roberta(RobertaForSequenceClassification)

Roberta for Sequence Classification

XLMRoberta(RobertaForSequenceClassification)

XLMRoberta for Sequence Classification

Albert(AlbertForSequenceClassification)

Albert for Sequence Classification

XLNet(XLNetForSequenceClassification)

XLNet for Sequence Classification

Longformer(LongformerForSequenceClassification)

Longformer for Sequence Classification

Implementations

impl ZeroShotClassificationOption

pub fn new(config: &ZeroShotClassificationConfig) -> Result<Self, RustBertError>

Instantiate a new zer-shot classification model of the supplied type.

Arguments

• ZeroShotClassificationConfig - Zero-shot classification pipeline configuration. The type of model created will be inferred from the ModelResources (Torch or ONNX) and ModelType (Architecture for Torch models) variants provided and

pub fn model_type(&self) -> ModelType

Returns the ModelType for this SequenceClassificationOption

pub fn forward_t( &self, input_ids: Option<&Tensor>, mask: Option<&Tensor>, token_type_ids: Option<&Tensor>, position_ids: Option<&Tensor>, input_embeds: Option<&Tensor>, train: bool ) -> Tensor

Interface method to forward_t() of the particular models.