pipeline.py

from typing import List

import fasttext.util

class PreTrainedPipeline():
    def __init__(self, path=""):
        """
        Initialize model
        """
        self.model = fasttext.load_model(os.path.join(path, 'cc.en.300.bin'))

    def __call__(self, inputs: str) -> List[float]:
        """
        Args:
            inputs (:obj:`str`):
                a string to get the features of.
        Return:
            A :obj:`list` of floats: The features computed by the model.
        """
        return self.model.get_sentence_vector(inputs)