base.py 文件源码

def _load_model(self):
        """The method to download the model from S3 and load to the memory.

        Args:
            saved (bool): wether to save the model files or just load it to the memory.

        Returns:
            gensim.models.doc2vec.Doc2Vec: The word-embedding model object.
        """
        try:
            model = Doc2Vec.load(os.path.join(LOCAL_CACHE_DIRECTORY, self.model_name))
            return model

        except:
            files  = list_files(self.s3_conn, self.s3_path)
            if not self.saved:
                with tempfile.TemporaryDirectory() as td:
                    for f in files:
                        filepath = os.path.join(td, f)
                        if not os.path.exists(filepath):
                            logging.warning('calling download from %s to %s', self.s3_path + f, filepath)
                            download(self.s3_conn, filepath, os.path.join(self.s3_path, f))
                    model = Doc2Vec.load(os.path.join(td, self.model_name))

            else:
                if not os.path.isdir(LOCAL_CACHE_DIRECTORY):
                    os.mkdir(LOCAL_CACHE_DIRECTORY)
                for f in files:
                    filepath = os.path.join(LOCAL_CACHE_DIRECTORY, f)
                    if not os.path.exists(filepath) and self.saved:
                        logging.warning('calling download from %s to %s', self.s3_path + f, filepath)
                        download(self.s3_conn, filepath, os.path.join(self.s3_path, f))
                model = Doc2Vec.load(os.path.join(LOCAL_CACHE_DIRECTORY, self.model_name))

            return model