Enum rust_bert::pipelines::sequence_classification::SequenceClassificationOption

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

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

Variants

Bert(BertForSequenceClassification)

Bert for Sequence Classification

DistilBert(DistilBertModelClassifier)

DistilBert 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

Bart(BartForSequenceClassification)

Bart for Sequence Classification

Implementations

impl SequenceClassificationOption

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

Instantiate a new sequence 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.