fasttext_predict.py

class FastTextSentiment(Base):
    """Predict fine-grained sentiment scores using FastText"""
    def __init__(self, model_file: str=None) -> None:
        super().__init__()
        import fasttext
        self.model = fasttext.load_model(model_file)

    def score(self, text: str) -> int:
        # Predict just the top label (hence 1 index below)
        labels, probabilities = self.model.predict(text, 1)
        pred = int(labels[0][-1])
        return pred

    def predict(self, train_file: None, test_file: str, lower_case: bool) -> pd.DataFrame:
        df = self.read_data(test_file, lower_case)
        df['pred'] = df['text'].apply(self.score)
        return df