Enum rust_bert::pipelines::token_classification::TokenClassificationOption

pub enum TokenClassificationOption {
    Bert(BertForTokenClassification),
    DistilBert(DistilBertForTokenClassification),
    Roberta(RobertaForTokenClassification),
    Electra(ElectraForTokenClassification),
}

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

Variants

Bert(BertForTokenClassification)

Bert for Token Classification

DistilBert(DistilBertForTokenClassification)

DistilBert for Token Classification

Roberta(RobertaForTokenClassification)

Roberta for Token Classification

Electra(ElectraForTokenClassification)

Electra for Token Classification

Implementations

impl TokenClassificationOption

pub fn new(model_type: ModelType, p: &Path, config: &ConfigOption) -> Self

Instantiate a new token 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 TokenClassificationOption