[][src]Module rust_bert::pipelines::sentiment

Sentiment Analysis pipeline

Predicts the binary sentiment for a sentence. DistilBERT model finetuned on SST-2. All resources for this model can be downloaded using the Python utility script included in this repository.

  1. Set-up a Python virtual environment and install dependencies (in ./requirements.txt)
  2. Run the conversion script python /utils/download-dependencies_sst2_sentiment.py. The dependencies will be downloaded to the user's home directory, under ~/rustbert/distilbert_sst2
use rust_bert::pipelines::sentiment::SentimentClassifier;
let device = Device::cuda_if_available();
let sentiment_classifier = SentimentClassifier::new(vocab_path,
                                                    config_path,
                                                    weights_path, device)?;
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 = sentiment_classifier.predict(&input);

(Example courtesy of IMDb)

Output:

[
   Sentiment { polarity: Positive, score: 0.998 },
   Sentiment { polarity: Negative, score: 0.992 },
   Sentiment { polarity: Positive, score: 0.999 }
]

Structs

Sentiment

Sentiment returned by the model.

SentimentClassifier

SentimentClassifier to perform sentiment analysis

Enums

SentimentPolarity

Enum with the possible sentiment polarities. Note that the pre-trained SST2 model does not include neutral sentiment.