Struct rust_bert::pipelines::sequence_classification::SequenceClassificationModel

pub struct SequenceClassificationModel { /* private fields */ }

SequenceClassificationModel for Classification (e.g. Sentiment Analysis)

Implementations

impl SequenceClassificationModel

pub fn new( config: SequenceClassificationConfig ) -> Result<SequenceClassificationModel, RustBertError>

Build a new SequenceClassificationModel

Arguments

• config - SequenceClassificationConfig object containing the resource references (model, vocabulary, configuration) and device placement (CPU/GPU)

Example

use rust_bert::pipelines::sequence_classification::SequenceClassificationModel;

let model = SequenceClassificationModel::new(Default::default())?;

pub fn new_with_tokenizer( config: SequenceClassificationConfig, tokenizer: TokenizerOption ) -> Result<SequenceClassificationModel, RustBertError>

Build a new SequenceClassificationModel with a provided tokenizer.

Arguments

• config - SequenceClassificationConfig object containing the resource references (model, vocabulary, configuration) and device placement (CPU/GPU)
• tokenizer - TokenizerOption tokenizer to use for sequence classification.

Example

use rust_bert::pipelines::common::{ModelType, TokenizerOption};
use rust_bert::pipelines::sequence_classification::SequenceClassificationModel;
let tokenizer = TokenizerOption::from_file(
    ModelType::Bert,
    "path/to/vocab.txt",
    None,
    false,
    None,
    None,
)?;
let model = SequenceClassificationModel::new_with_tokenizer(Default::default(), tokenizer)?;

pub fn get_tokenizer(&self) -> &TokenizerOption

Get a reference to the model tokenizer.

pub fn get_tokenizer_mut(&mut self) -> &mut TokenizerOption

Get a mutable reference to the model tokenizer.

pub fn predict<'a, S>(&self, input: S) -> Vec<Label>

where S: AsRef<[&'a str]>,

Classify texts

Arguments

• input - &[&str] Array of texts to classify.

Returns

• Vec<Label> containing labels for input texts

Example

let sequence_classification_model =  SequenceClassificationModel::new(Default::default())?;
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);

pub fn predict_multilabel( &self, input: &[&str], threshold: f64 ) -> Result<Vec<Vec<Label>>, RustBertError>

Multi-label classification of texts

Arguments

• input - &[&str] Array of texts to classify.
• threshold - f64 threshold above which a label will be considered true by the classifier

Returns

• Vec<Vec<Label>> containing a vector of true labels for each input text

Example

let sequence_classification_model =  SequenceClassificationModel::new(Default::default())?;
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_multilabel(&input, 0.5);