<commit_before>#!/usr/bin/env python
# -*- mode: cython -*-

cimport cpolycomp

import numpy as np
cimport numpy as np

# Direction of the Chebyshev transform
PCOMP_TD_DIRECT = 0
PCOMP_TD_INVERSE = 1

# Algorithm to use for polynomial compression
PCOMP_ALG_USE_CHEBYSHEV = 0
PCOMP_ALG_NO_CHEBYSHEV = 1

PCOMP_STAT_SUCCESS = 0 # All ok
PCOMP_STAT_INVALID_ENCODING = 1 # Decompression error
PCOMP_STAT_INVALID_BUFFER = 2 # Output buffer too small
PCOMP_STAT_INVALID_FIT = 3 # Least-square fit error

def __version__():
    cdef int major
    cdef int minor
    cpolycomp.pcomp_version(&major, &minor)

    return "{0}.{1}".format(major, minor)

################################################################################
# RLE compression

def rle_bufsize(input_size):
    return cpolycomp.pcomp_rle_bufsize(input_size)

def rle_compress_int8(np.ndarray[np.int8_t, ndim=1] values):
    cdef size_t real_size
    cdef int result
    cdef np.ndarray[np.int8_t, ndim=1] output = np.empty(rle_bufsize(values.size),
                                                         dtype='int8')

    result = cpolycomp.pcomp_compress_rle_int8(<np.int8_t *> output.data, &real_size,
                                               <np.int8_t *> values.data, values.size)
    return np.resize(output, real_size)

################################################################################
# Quantization

cdef class QuantParams:
    cdef cpolycomp.pcomp_quant_params_t* _c_params
    
    def __cinit__(self, element_size, bits_per_sample):
        self._c_params = cpolycomp.pcomp_init_quant_params(element_size, bits_per_sample)

    def element_size(self):
        return cpolycomp.pcomp_quant_element_size(self._c_params)

    def bits_per_sample(self):
        return cpolycomp.pcomp_quant_bits_per_sample(self._c_params)    

    def buf_size(self, input_size):
        return cpolycomp.pcomp_quant_bufsize(input_size, self._c_params)<commit_msg>Make "rle_compress" a polymorphic function<commit_after>#!/usr/bin/env python
# -*- mode: cython -*-

cimport cpolycomp

import numpy as np
cimport numpy as np

# Direction of the Chebyshev transform
PCOMP_TD_DIRECT = 0
PCOMP_TD_INVERSE = 1

# Algorithm to use for polynomial compression
PCOMP_ALG_USE_CHEBYSHEV = 0
PCOMP_ALG_NO_CHEBYSHEV = 1

PCOMP_STAT_SUCCESS = 0 # All ok
PCOMP_STAT_INVALID_ENCODING = 1 # Decompression error
PCOMP_STAT_INVALID_BUFFER = 2 # Output buffer too small
PCOMP_STAT_INVALID_FIT = 3 # Least-square fit error

def __version__():
    cdef int major
    cdef int minor
    cpolycomp.pcomp_version(&major, &minor)

    return "{0}.{1}".format(major, minor)

################################################################################
# RLE compression

def rle_bufsize(input_size):
    return cpolycomp.pcomp_rle_bufsize(input_size)

def rle_compress(np.ndarray values not None):
    cdef int result
    cdef size_t num_of_bytes = values.dtype.itemsize * rle_bufsize(values.size)
    cdef size_t real_size = num_of_bytes
    cdef np.ndarray[np.int8_t, ndim=1] output = np.empty(num_of_bytes,
                                                         dtype='int8')

    if values.dtype in (np.int8, np.uint8):
        result = cpolycomp.pcomp_compress_rle_int8(<np.int8_t *> &output.data[0],
                                                   &real_size,
                                                   <np.int8_t *> &values.data[0],
                                                   values.size)
    elif values.dtype in (np.int16, np.uint16):
        result = cpolycomp.pcomp_compress_rle_int16(<np.int16_t *> &output.data[0],
                                                   &real_size,
                                                   <np.int16_t *> &values.data[0],
                                                   values.size)
    elif values.dtype in (np.int32, np.uint32):
        result = cpolycomp.pcomp_compress_rle_int32(<np.int32_t *> &output.data[0],
                                                   &real_size,
                                                   <np.int32_t *> &values.data[0],
                                                   values.size)
    elif values.dtype in (np.int64, np.uint64):
        result = cpolycomp.pcomp_compress_rle_int64(<np.int64_t *> &output.data[0],
                                                   &real_size,
                                                   <np.int64_t *> &values.data[0],
                                                   values.size)

    output = np.resize(output, real_size * values.dtype.itemsize)
    return np.fromstring(output.tobytes(), values.dtype)

################################################################################
# Quantization

cdef class QuantParams:
    cdef cpolycomp.pcomp_quant_params_t* _c_params
    
    def __cinit__(self, element_size, bits_per_sample):
        self._c_params = cpolycomp.pcomp_init_quant_params(element_size, bits_per_sample)

    def element_size(self):
        return cpolycomp.pcomp_quant_element_size(self._c_params)

    def bits_per_sample(self):
        return cpolycomp.pcomp_quant_bits_per_sample(self._c_params)    

    def buf_size(self, input_size):
        return cpolycomp.pcomp_quant_bufsize(input_size, self._c_params)
{'fraction_non_alphanumeric': 0.13943313677002092, 'fraction_numerical': 0.01426669202967472, 'mean_word_length': 2.4254071661237786, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 30, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '567565', 'n_tokens_mistral': 1645, 'n_tokens_neox': 1591, 'n_words': 305}