vader_predict.py

class VaderSentiment(Base):
    """Predict fine-grained sentiment classes using Vader."""
    def __init__(self, model_file: str=None) -> None:
        super().__init__()
        from nltk.sentiment.vader import SentimentIntensityAnalyzer
        self.vader = SentimentIntensityAnalyzer()

    def score(self, text: str) -> float:
        return self.vader.polarity_scores(text)['compound']

    def predict(self, train_file: None, test_file: str, lower_case: bool) -> pd.DataFrame:
        "Return DataFrame with a new column of predicted labels"
        df = self.read_data(test_file, lower_case)
        df['score'] = df['text'].apply(self.score)
        # Convert float score to category based on binning
        df['pred'] = pd.cut(df['score'], bins=5, labels=[1, 2, 3, 4, 5])
        df = df.drop('score', axis=1)
        return df