Module rust_bert::pipelines::sequence_classification

Sequence classification pipeline (e.g. Sentiment Analysis)

More generic sequence classification pipeline, works with multiple models (Bert, Roberta)

use rust_bert::pipelines::sequence_classification::SequenceClassificationConfig;
use rust_bert::resources::{RemoteResource, Resource};
use rust_bert::distilbert::{DistilBertModelResources, DistilBertVocabResources, DistilBertConfigResources};
use rust_bert::pipelines::sequence_classification::SequenceClassificationModel;
use rust_bert::pipelines::common::ModelType;

//Load a configuration
let config = SequenceClassificationConfig::new(ModelType::DistilBert,
   Resource::Remote(RemoteResource::from_pretrained(DistilBertModelResources::DISTIL_BERT_SST2)),
   Resource::Remote(RemoteResource::from_pretrained(DistilBertVocabResources::DISTIL_BERT_SST2)),
   Resource::Remote(RemoteResource::from_pretrained(DistilBertConfigResources::DISTIL_BERT_SST2)),
   None, //merges resource only relevant with ModelType::Roberta
   true, //lowercase
);

//Create the model
let sequence_classification_model = SequenceClassificationModel::new(config)?;

let input = [
    "Probably my all-time favorite movie, a story of selflessness, sacrifice and dedication to a noble cause, but it's not preachy or boring.",
    "This film tried to be too many things all at once: stinging political satire, Hollywood blockbuster, sappy romantic comedy, family values promo...",
    "If you like original gut wrenching laughter you will like this movie. If you are young or old then you will love this movie, hell even my mom liked it.",
];
let output = sequence_classification_model.predict(&input);

(Example courtesy of IMDb)

Output:

let output =
[
   Label { text: String::from("POSITIVE"), score: 0.9986, id: 1, sentence: 0},
   Label { text: String::from("NEGATIVE"), score: 0.9985, id: 0, sentence: 1},
   Label { text: String::from("POSITIVE"), score: 0.9988, id: 1, sentence: 12},
]

Structs

Label	Label generated by a SequenceClassificationModel
SequenceClassificationConfig	Configuration for SequenceClassificationModel
SequenceClassificationModel	SequenceClassificationModel for Classification (e.g. Sentiment Analysis)

Enums

SequenceClassificationOption

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