Enum rust_bert::pipelines::zero_shot_classification::ZeroShotClassificationOption

pub enum ZeroShotClassificationOption {
    Bart(BartForSequenceClassification),
    Bert(BertForSequenceClassification),
    DistilBert(DistilBertModelClassifier),
    MobileBert(MobileBertForSequenceClassification),
    Roberta(RobertaForSequenceClassification),
    XLMRoberta(RobertaForSequenceClassification),
    Albert(AlbertForSequenceClassification),
    XLNet(XLNetForSequenceClassification),
}

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

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

Implementations

impl ZeroShotClassificationOption

pub fn new<'p, P>(
    model_type: ModelType,
    p: P,
    config: &ConfigOption
) -> Result<Self, RustBertError> where
    P: Borrow<Path<'p>>, 

Instantiate a new zero shot classification model of the supplied type.

Arguments

• model_type - ModelType indicating the model type to load (must match with the actual data to be loaded)
• p - tch::nn::Path path to the model file to load (e.g. model.ot)
• config - A configuration (the model type of the configuration must be compatible with the value for model_type)

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.