Enum rust_bert::pipelines::sentence_embeddings::SentenceEmbeddingsOption

pub enum SentenceEmbeddingsOption {
    Bert(BertForSentenceEmbeddings),
    DistilBert(DistilBertForSentenceEmbeddings),
    Roberta(RobertaForSentenceEmbeddings),
    Albert(AlbertForSentenceEmbeddings),
    T5(T5ForSentenceEmbeddings),
}

Abstraction that holds one particular sentence embeddings model, for any of the supported models

Variants

Bert(BertForSentenceEmbeddings)

Bert for Sentence Embeddings

DistilBert(DistilBertForSentenceEmbeddings)

DistilBert for Sentence Embeddings

Roberta(RobertaForSentenceEmbeddings)

Roberta for Sentence Embeddings

Albert(AlbertForSentenceEmbeddings)

Albert for Sentence Embeddings

T5(T5ForSentenceEmbeddings)

T5 for Sentence Embeddings

Implementations

impl SentenceEmbeddingsOption

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

where P: Borrow<Path<'p>>,

Instantiate a new sentence embeddings transformer of the supplied type.

Arguments

• transformer_type - ModelType indicating the transformer 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. rust_model.ot)
• config - A configuration (the transformer model type of the configuration must be compatible with the value for transformer_type)

pub fn forward( &self, tokens_ids: &Tensor, tokens_masks: &Tensor ) -> Result<(Tensor, Option<Vec<Tensor>>), RustBertError>

Interface method to forward() of the particular transformer models.