<commit_before>from copy import copy
import os
import tempfile

cimport cython

import numpy as np
cimport numpy as np

cimport flann
from flann cimport (flann_index_t as index_t,
                    flann_algorithm_t as algorithm_t,
                    flann_centers_init_t as centers_init_t,
                    flann_distance_t as distance_t,
                    flann_log_level_t as log_level_t,
                    FLANNParameters as Parameters,
                    DEFAULT_FLANN_PARAMETERS as DEFAULT_PARAMETERS,
                    flann_build_index_float as build_index_float)

################################################################################
### Enum conversions

str_to_algorithm_t = {
    'linear': flann.FLANN_INDEX_LINEAR,
    'kdtree': flann.FLANN_INDEX_KDTREE,
    'kmeans': flann.FLANN_INDEX_KMEANS,
    'composite': flann.FLANN_INDEX_COMPOSITE,
    'kdtree_single': flann.FLANN_INDEX_KDTREE_SINGLE,
    'hierarchical': flann.FLANN_INDEX_HIERARCHICAL,
    'lsh': flann.FLANN_INDEX_LSH,
    'saved': flann.FLANN_INDEX_SAVED,
    'autotuned': flann.FLANN_INDEX_AUTOTUNED,
}
algorithm_t_to_str = {v: k for k, v in str_to_algorithm_t.iteritems()}

str_to_centers_init_t = {
    'random': flann.FLANN_CENTERS_RANDOM,
    'gonzales': flann.FLANN_CENTERS_GONZALES,
    'kmeanspp': flann.FLANN_CENTERS_KMEANSPP,
}
centers_init_t_to_str = {v: k for k, v in str_to_centers_init_t.iteritems()}

str_to_log_level_t = {
    'none': flann.FLANN_LOG_NONE,
    'fatal': flann.FLANN_LOG_FATAL,
    'error': flann.FLANN_LOG_ERROR,
    'warn': flann.FLANN_LOG_WARN,
    'info': flann.FLANN_LOG_INFO,
    'debug': flann.FLANN_LOG_DEBUG,
}
log_level_t_to_str = {v: k for k, v in str_to_log_level_t.iteritems()}

str_to_distance_t = {
    'euclidean': flann.FLANN_DIST_EUCLIDEAN,
    'l2': flann.FLANN_DIST_L2,
    'manhattan': flann.FLANN_DIST_MANHATTAN,
    'l1': flann.FLANN_DIST_L1,
    'minkowski': flann.FLANN_DIST_MINKOWSKI,
    'max_dist': flann.FLANN_DIST_MAX,
    'hik': flann.FLANN_DIST_HIST_INTERSECT,
    'hellinger': flann.FLANN_DIST_HELLINGER,
    'chi_square': flann.FLANN_DIST_CHI_SQUARE,
    'kullback_leibler': flann.FLANN_DIST_KULLBACK_LEIBLER,
    'hamming': flann.FLANN_DIST_HAMMING,
    'hamming_lut': flann.FLANN_DIST_HAMMING_LUT,
    'hamming_popcnt': flann.FLANN_DIST_HAMMING_POPCNT,
    'l2_simple': flann.FLANN_DIST_L2_SIMPLE,
}
distance_t_to_str = {
    flann.FLANN_DIST_EUCLIDEAN: 'euclidean',
    flann.FLANN_DIST_MANHATTAN: 'manhattan',
    flann.FLANN_DIST_MINKOWSKI: 'minkowski',
    flann.FLANN_DIST_MAX: 'max_dist',
    flann.FLANN_DIST_HIST_INTERSECT: 'hik',
    flann.FLANN_DIST_HELLINGER: 'hellinger',
    flann.FLANN_DIST_CHI_SQUARE: 'chi_square',
    flann.FLANN_DIST_KULLBACK_LEIBLER: 'kullback_leibler',
    flann.FLANN_DIST_HAMMING: 'hamming',
    flann.FLANN_DIST_HAMMING_LUT: 'hamming_lut',
    flann.FLANN_DIST_HAMMING_POPCNT: 'hamming_popcnt',
    flann.FLANN_DIST_L2_SIMPLE: 'l2_simple',
}

def set_distance_type(str distance_type, int order = 0):
    cdef distance_t dist = str_to_distance_t[distance_type]
    flann.flann_set_distance_type(dist, order)


################################################################################
### Wrapping parameters struct

cdef class FLANNParameters:
    # attributes in the pxd: _this
    def __cinit__(self):
        self._this = DEFAULT_PARAMETERS

    def __init__(self, **kwargs):
        for k, v in kwargs.iteritems():
            setattr(self, k, v)

    def __getitem__(self, k):
        try:
            return getattr(self, k)
        except AttributeError as e:
            raise KeyError(*e.args)

    def __setitem__(self, k, v):
        try:
            setattr(self, k, v)
        except AttributeError as e:
            raise KeyError(*e.args)

    def update(self, dict={}, **kwargs):
        for k, v in dict.iteritems():
            setattr(self, k, v)
        for k, v in kwargs.iteritems():
            setattr(self, k, v)

    def as_dict(self):
        cdef dict d = dict(self._this)
        d['algorithm'] = self.algorithm
        d['centers_init'] = self.centers_init
        d['log_level'] = self.log_level
        return d

    def __repr__(self):
        return "<FLANNParameters {}>".format(self.as_dict())

    property algorithm:
        # the algorithm to use
        def __get__(self):
            return algorithm_t_to_str[self._this.algorithm]
        def __set__(self, str algorithm):
            cdef algorithm_t algo = str_to_algorithm_t[algorithm]
            self._this.algorithm = algo

    # search time parameters
    property checks:  # TODO: also support unlimited, autotuned
        # how many leaves (features) to check in one search
        def __get__(self): return self._this.checks
        def __set__(self, int checks): self._this.checks = checks
    property eps:
        # eps parameter for eps-knn search
        def __get__(self): return self._this.eps
        def __set__(self, float eps): self._this.eps = eps
    property sorted:
        # indicates if results returned by radius search should be sorted or not
        def __get__(self): return bool(self._this.sorted)
        def __set__(self, bint sorted):
            self._this.sorted = sorted
    property max_neighbors:
        # limits the max number of neighbors should be returned by radius search
        def __get__(self): return self._this.max_neighbors
        def __set__(self, int max_neighbors):
            self._this.max_neighbors = max_neighbors
    property cores:
        # number of parallel cores to use for searching
        def __get__(self): return self._this.cores
        def __set__(self, int cores): self._this.cores = cores

    # KD tree params
    property trees:
        # number of randomized trees to use (for kdtree)
        def __get__(self): return self._this.trees
        def __set__(self, int trees): self._this.trees = trees
    property leaf_max_size:
        def __get__(self): return self._this.leaf_max_size
        def __set__(self, int leaf_max_size):
            self._this.leaf_max_size = leaf_max_size

    # kmeans index params
    property branching:
        # branching factor (for kmeans tree)
        def __get__(self): return self._this.branching
        def __set__(self, int branching): self._this.branching = branching
    property iterations:
        # max iterations to perform in one kmeans clustering (kmeans tree)
        def __get__(self): return self._this.iterations
        def __set__(self, int iterations): self._this.iterations = iterations
    property centers_init:
        # algorithm used for picking the initial cluster centers for kmeans tree
        def __get__(self):
            return centers_init_t_to_str[self._this.centers_init]
        def __set__(self, str centers_init):
            cdef centers_init_t init = str_to_centers_init_t[centers_init]
            self._this.centers_init = init
    property cb_index:
        # cluster boundary index. Used when searching the kmeans tree
        def __get__(self): return self._this.cb_index
        def __set__(self, float cb_index): self._this.cb_index = cb_index

    # autotuned index parameters
    property target_precision:
        # precision desired (used for autotuning, -1 otherwise) */
        def __get__(self): return self._this.target_precision
        def __set__(self, float target_precision):
            self._this.target_precision = target_precision
    property build_weight:
        # build tree time weighting factor
        def __get__(self): return self._this.build_weight
        def __set__(self, float build_weight):
            self._this.build_weight = build_weight
    property memory_weight:
        # index memory weigthing factor
        def __get__(self): return self._this.memory_weight
        def __set__(self, float memory_weight):
            self._this.memory_weight = memory_weight
    property sample_fraction:
        # what fraction of the dataset to use for autotuning
        def __get__(self): return self._this.sample_fraction
        def __set__(self, float sample_fraction):
            self._this.sample_fraction = sample_fraction

    # LSH parameters
    property table_number_:
        # The number of hash tables to use
        def __get__(self): return self._this.table_number_
        def __set__(self, int table_number_):
            self._this.table_number_ = table_number_
    property key_size_:
        # The length of the key in the hash tables
        def __get__(self): return self._this.key_size_
        def __set__(self, int key_size_): self._this.key_size_ = key_size_
    property multi_probe_level_:
        # Number of levels to use in multi-probe LSH, 0 for standard LSH
        def __get__(self): return self._this.multi_probe_level_
        def __set__(self, int multi_probe_level_):
            self._this.multi_probe_level_ = multi_probe_level_

    # other parameters
    property log_level:
        # determines the verbosity of each flann function
        def __get__(self): return log_level_t_to_str[self._this.log_level]
        def __set__(self, str log_level):
            cdef log_level_t level = str_to_log_level_t[log_level]
            self._this.log_level = level
    property random_seed:
        # random seed to use
        def __get__(self): return self._this.random_seed
        def __set__(self, long random_seed):
            self._this.random_seed = random_seed


################################################################################
### Wrapping index class

# TODO: template across different types
cdef class FLANNIndex:
    # attributes in the pxd: _this, _data, speedup, params

    def __cinit__(self):
        self._this = NULL
        # TODO: _data?
        self.speedup = -1
        self.params = FLANNParameters()

    def __init__(self, **kwargs):
        self.params.update(**kwargs)
        self._rn_gen = np.random.RandomState()
        self._rn_gen.seed()

    def free_index(self):
        self._free_index()

    cdef void _free_index(self) nogil:
        flann.flann_free_index_float(self._this, &self.params._this)
        self._this = NULL
        with gil:  # TODO: how to handle this more appropriately?
            self._data = np.empty((0, 0), dtype=np.float32)

    def __dealloc__(self):
        if self._this is not NULL:
            self._free_index()

    property data:
        def __get__(self):
            try:
                return np.asarray(self._data)
            except AttributeError:
                return None

    ############################################################################
    ### general helpers

    cdef _ensure_random_seed(self, kwargs):
        if 'random_seed' not in kwargs:
            kwargs['random_seed'] = self._rn_gen.randint(2 ** 30)

    cpdef _check_array(self, array, int dim=2):
        array = np.require(array,
                    requirements=['C_CONTIGUOUS', 'ALIGNED'],
                    dtype=np.float32)
        if dim == 2 and array.ndim == 1:
            array = array.reshape(1, array.size)
        if array.ndim != dim:
            raise ValueError("expected a {}d array".format(dim))<commit_msg>don't constantly re-seed the random number generator, it's slow<commit_after>from copy import copy
import os
import tempfile

cimport cython

import numpy as np
cimport numpy as np

cimport flann
from flann cimport (flann_index_t as index_t,
                    flann_algorithm_t as algorithm_t,
                    flann_centers_init_t as centers_init_t,
                    flann_distance_t as distance_t,
                    flann_log_level_t as log_level_t,
                    FLANNParameters as Parameters,
                    DEFAULT_FLANN_PARAMETERS as DEFAULT_PARAMETERS,
                    flann_build_index_float as build_index_float)

################################################################################
### Enum conversions

str_to_algorithm_t = {
    'linear': flann.FLANN_INDEX_LINEAR,
    'kdtree': flann.FLANN_INDEX_KDTREE,
    'kmeans': flann.FLANN_INDEX_KMEANS,
    'composite': flann.FLANN_INDEX_COMPOSITE,
    'kdtree_single': flann.FLANN_INDEX_KDTREE_SINGLE,
    'hierarchical': flann.FLANN_INDEX_HIERARCHICAL,
    'lsh': flann.FLANN_INDEX_LSH,
    'saved': flann.FLANN_INDEX_SAVED,
    'autotuned': flann.FLANN_INDEX_AUTOTUNED,
}
algorithm_t_to_str = {v: k for k, v in str_to_algorithm_t.iteritems()}

str_to_centers_init_t = {
    'random': flann.FLANN_CENTERS_RANDOM,
    'gonzales': flann.FLANN_CENTERS_GONZALES,
    'kmeanspp': flann.FLANN_CENTERS_KMEANSPP,
}
centers_init_t_to_str = {v: k for k, v in str_to_centers_init_t.iteritems()}

str_to_log_level_t = {
    'none': flann.FLANN_LOG_NONE,
    'fatal': flann.FLANN_LOG_FATAL,
    'error': flann.FLANN_LOG_ERROR,
    'warn': flann.FLANN_LOG_WARN,
    'info': flann.FLANN_LOG_INFO,
    'debug': flann.FLANN_LOG_DEBUG,
}
log_level_t_to_str = {v: k for k, v in str_to_log_level_t.iteritems()}

str_to_distance_t = {
    'euclidean': flann.FLANN_DIST_EUCLIDEAN,
    'l2': flann.FLANN_DIST_L2,
    'manhattan': flann.FLANN_DIST_MANHATTAN,
    'l1': flann.FLANN_DIST_L1,
    'minkowski': flann.FLANN_DIST_MINKOWSKI,
    'max_dist': flann.FLANN_DIST_MAX,
    'hik': flann.FLANN_DIST_HIST_INTERSECT,
    'hellinger': flann.FLANN_DIST_HELLINGER,
    'chi_square': flann.FLANN_DIST_CHI_SQUARE,
    'kullback_leibler': flann.FLANN_DIST_KULLBACK_LEIBLER,
    'hamming': flann.FLANN_DIST_HAMMING,
    'hamming_lut': flann.FLANN_DIST_HAMMING_LUT,
    'hamming_popcnt': flann.FLANN_DIST_HAMMING_POPCNT,
    'l2_simple': flann.FLANN_DIST_L2_SIMPLE,
}
distance_t_to_str = {
    flann.FLANN_DIST_EUCLIDEAN: 'euclidean',
    flann.FLANN_DIST_MANHATTAN: 'manhattan',
    flann.FLANN_DIST_MINKOWSKI: 'minkowski',
    flann.FLANN_DIST_MAX: 'max_dist',
    flann.FLANN_DIST_HIST_INTERSECT: 'hik',
    flann.FLANN_DIST_HELLINGER: 'hellinger',
    flann.FLANN_DIST_CHI_SQUARE: 'chi_square',
    flann.FLANN_DIST_KULLBACK_LEIBLER: 'kullback_leibler',
    flann.FLANN_DIST_HAMMING: 'hamming',
    flann.FLANN_DIST_HAMMING_LUT: 'hamming_lut',
    flann.FLANN_DIST_HAMMING_POPCNT: 'hamming_popcnt',
    flann.FLANN_DIST_L2_SIMPLE: 'l2_simple',
}

def set_distance_type(str distance_type, int order = 0):
    cdef distance_t dist = str_to_distance_t[distance_type]
    flann.flann_set_distance_type(dist, order)


################################################################################
### Wrapping parameters struct

cdef class FLANNParameters:
    # attributes in the pxd: _this
    def __cinit__(self):
        self._this = DEFAULT_PARAMETERS

    def __init__(self, **kwargs):
        for k, v in kwargs.iteritems():
            setattr(self, k, v)

    def __getitem__(self, k):
        try:
            return getattr(self, k)
        except AttributeError as e:
            raise KeyError(*e.args)

    def __setitem__(self, k, v):
        try:
            setattr(self, k, v)
        except AttributeError as e:
            raise KeyError(*e.args)

    def update(self, dict={}, **kwargs):
        for k, v in dict.iteritems():
            setattr(self, k, v)
        for k, v in kwargs.iteritems():
            setattr(self, k, v)

    def as_dict(self):
        cdef dict d = dict(self._this)
        d['algorithm'] = self.algorithm
        d['centers_init'] = self.centers_init
        d['log_level'] = self.log_level
        return d

    def __repr__(self):
        return "<FLANNParameters {}>".format(self.as_dict())

    property algorithm:
        # the algorithm to use
        def __get__(self):
            return algorithm_t_to_str[self._this.algorithm]
        def __set__(self, str algorithm):
            cdef algorithm_t algo = str_to_algorithm_t[algorithm]
            self._this.algorithm = algo

    # search time parameters
    property checks:  # TODO: also support unlimited, autotuned
        # how many leaves (features) to check in one search
        def __get__(self): return self._this.checks
        def __set__(self, int checks): self._this.checks = checks
    property eps:
        # eps parameter for eps-knn search
        def __get__(self): return self._this.eps
        def __set__(self, float eps): self._this.eps = eps
    property sorted:
        # indicates if results returned by radius search should be sorted or not
        def __get__(self): return bool(self._this.sorted)
        def __set__(self, bint sorted):
            self._this.sorted = sorted
    property max_neighbors:
        # limits the max number of neighbors should be returned by radius search
        def __get__(self): return self._this.max_neighbors
        def __set__(self, int max_neighbors):
            self._this.max_neighbors = max_neighbors
    property cores:
        # number of parallel cores to use for searching
        def __get__(self): return self._this.cores
        def __set__(self, int cores): self._this.cores = cores

    # KD tree params
    property trees:
        # number of randomized trees to use (for kdtree)
        def __get__(self): return self._this.trees
        def __set__(self, int trees): self._this.trees = trees
    property leaf_max_size:
        def __get__(self): return self._this.leaf_max_size
        def __set__(self, int leaf_max_size):
            self._this.leaf_max_size = leaf_max_size

    # kmeans index params
    property branching:
        # branching factor (for kmeans tree)
        def __get__(self): return self._this.branching
        def __set__(self, int branching): self._this.branching = branching
    property iterations:
        # max iterations to perform in one kmeans clustering (kmeans tree)
        def __get__(self): return self._this.iterations
        def __set__(self, int iterations): self._this.iterations = iterations
    property centers_init:
        # algorithm used for picking the initial cluster centers for kmeans tree
        def __get__(self):
            return centers_init_t_to_str[self._this.centers_init]
        def __set__(self, str centers_init):
            cdef centers_init_t init = str_to_centers_init_t[centers_init]
            self._this.centers_init = init
    property cb_index:
        # cluster boundary index. Used when searching the kmeans tree
        def __get__(self): return self._this.cb_index
        def __set__(self, float cb_index): self._this.cb_index = cb_index

    # autotuned index parameters
    property target_precision:
        # precision desired (used for autotuning, -1 otherwise) */
        def __get__(self): return self._this.target_precision
        def __set__(self, float target_precision):
            self._this.target_precision = target_precision
    property build_weight:
        # build tree time weighting factor
        def __get__(self): return self._this.build_weight
        def __set__(self, float build_weight):
            self._this.build_weight = build_weight
    property memory_weight:
        # index memory weigthing factor
        def __get__(self): return self._this.memory_weight
        def __set__(self, float memory_weight):
            self._this.memory_weight = memory_weight
    property sample_fraction:
        # what fraction of the dataset to use for autotuning
        def __get__(self): return self._this.sample_fraction
        def __set__(self, float sample_fraction):
            self._this.sample_fraction = sample_fraction

    # LSH parameters
    property table_number_:
        # The number of hash tables to use
        def __get__(self): return self._this.table_number_
        def __set__(self, int table_number_):
            self._this.table_number_ = table_number_
    property key_size_:
        # The length of the key in the hash tables
        def __get__(self): return self._this.key_size_
        def __set__(self, int key_size_): self._this.key_size_ = key_size_
    property multi_probe_level_:
        # Number of levels to use in multi-probe LSH, 0 for standard LSH
        def __get__(self): return self._this.multi_probe_level_
        def __set__(self, int multi_probe_level_):
            self._this.multi_probe_level_ = multi_probe_level_

    # other parameters
    property log_level:
        # determines the verbosity of each flann function
        def __get__(self): return log_level_t_to_str[self._this.log_level]
        def __set__(self, str log_level):
            cdef log_level_t level = str_to_log_level_t[log_level]
            self._this.log_level = level
    property random_seed:
        # random seed to use
        def __get__(self): return self._this.random_seed
        def __set__(self, long random_seed):
            self._this.random_seed = random_seed


################################################################################
### Wrapping index class

# TODO: template across different types
cdef class FLANNIndex:
    # attributes in the pxd: _this, _data, speedup, params

    def __cinit__(self):
        self._this = NULL
        # TODO: _data?
        self.speedup = -1
        self.params = FLANNParameters()

    def __init__(self, **kwargs):
        self.params.update(**kwargs)

    def free_index(self):
        self._free_index()

    cdef void _free_index(self) nogil:
        flann.flann_free_index_float(self._this, &self.params._this)
        self._this = NULL
        with gil:  # TODO: how to handle this more appropriately?
            self._data = np.empty((0, 0), dtype=np.float32)

    def __dealloc__(self):
        if self._this is not NULL:
            self._free_index()

    property data:
        def __get__(self):
            try:
                return np.asarray(self._data)
            except AttributeError:
                return None

    ############################################################################
    ### general helpers

    cdef _ensure_random_seed(self, kwargs):
        if 'random_seed' not in kwargs:
            kwargs['random_seed'] = np.random.randint(2 ** 30)

    cpdef _check_array(self, array, int dim=2):
        array = np.require(array,
                    requirements=['C_CONTIGUOUS', 'ALIGNED'],
                    dtype=np.float32)
        if dim == 2 and array.ndim == 1:
            array = array.reshape(1, array.size)
        if array.ndim != dim:
            raise ValueError("expected a {}d array".format(dim))
{'fraction_non_alphanumeric': 0.11013295005389867, 'fraction_numerical': 0.0022457779374775422, 'mean_word_length': 2.8999824837975128, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5065142', 'n_tokens_mistral': 7313, 'n_tokens_neox': 6944, 'n_words': 1666}