vector_utils.py 文件源码

def save_vector_cache(vectors, vector_out_file, filetype='', **kwargs):
    logging.info("Saving {} vectors to cache {}".format(len(vectors),vector_out_file))
    if (vector_out_file.endswith('.dill') or filetype == 'dill'):
        with open(vector_out_file, 'wb') as data_file:
            dill.dump(vectors, data_file, protocol=kwargs.get('dill_protocol', 3))
    elif (vector_out_file.endswith('.joblib') or filetype == 'joblib'):
        joblib.dump(vectors, vector_out_file, compress=kwargs.get('joblib_compression', 3),
                    protocol=kwargs.get('joblib_protocol', 3))
    elif (vector_out_file.endswith('.sqlite') or filetype == 'sqlite'):
        autocommit = kwargs.pop('autocommit', True)
        if (isinstance(vectors, SqliteDict)):
            vectors.commit()
        else:
            with SqliteDict(vector_out_file, autocommit=autocommit) as data_file:
                for key, value in vectors.items():
                    data_file[key] = value

                if (not autocommit):
                    data_file.commit()
    else:
        raise NotImplementedError