<commit_before>#cython: boundscheck=False
#cython: cdivision=True
# CTS code adapted from https://github.com/mgbellemare/SkipCTS
cimport cython
import numpy as np
cimport numpy as np

from libc.math cimport log, exp

from skimage.transform import resize


# Parameters of the CTS model. For clarity, we take these as constants.
cdef double PRIOR_STAY_PROB = 0.5
cdef double PRIOR_SPLIT_PROB = 0.5
cdef double LOG_PRIOR_STAY_PROB = log(PRIOR_STAY_PROB)
cdef double LOG_PRIOR_SPLIT_PROB = log(1.0 - PRIOR_STAY_PROB)
# Sampling parameter. The maximum number of rejections before we give up and
# sample from the root estimator.
cdef int MAX_SAMPLE_REJECTIONS = 25

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
cdef double get_prior(char* prior_name, int alphabet_size):
    if prior_name == <char*>'perks':
        return 1.0 / <double>alphabet_size
    elif prior_name == <char*>'jeffreys':
        return 0.5
    else: #use laplace prior
        return 1.0

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
cdef double log_add(double log_x, double log_y):
    """Given log x and log y, returns log(x + y)."""
    # Swap variables so log_y is larger.
    if log_x > log_y:
        log_x, log_y = log_y, log_x

    cdef double delta = log_y - log_x
    return log(1 + exp(delta)) + log_x if delta <= 50.0 else log_y


class Error(Exception):
    """Base exception for the `cts` module."""
    pass


cdef class Estimator:
    """The estimator for a CTS node.
    This implements a Dirichlet-multinomial model with specified prior. This
    class does not perform alphabet checking, and will return invalid
    probabilities if it is ever fed more than `model.alphabet_size` distinct
    symbols.
    Args:
        model: Reference to CTS model. We expected model.symbol_prior to be
            a `float`.
    """
    cdef double count_total
    cdef CTS _model
    cdef double[:] counts
    
    def __init__(self, CTS model):
        self.counts = np.ones([model.alphabet_size], dtype=np.double)*model.symbol_prior
        self.count_total = model.alphabet_size * model.symbol_prior
        self._model = model

    @cython.cdivision(True)
    cdef double prob(self, int symbol):
        """Returns the probability assigned to this symbol."""
        return self.counts[symbol] / self.count_total

    cdef double update(self, int symbol):
        """Updates our count for the given symbol."""
        cdef double prob = self.prob(symbol)
        cdef double log_prob = log(prob)
        self.counts[symbol] = self.counts[symbol] + 1.0
        self.count_total += 1.0
        return log_prob

    def __getstate__(self):
        return self.count_total, self._model, self.counts

    def __setstate__(self, state):
        self.count_total, self._model, self.counts = state
            
    
cdef class CTSNode:
    """A node in the CTS tree.
    Each node contains a base Dirichlet estimator holding the statistics for
    this particular context, and pointers to its children.
    """
    cdef double _log_stay_prob
    cdef double _log_split_prob
    cdef CTS _model
    cdef Estimator estimator
    cdef dict _children

    def __init__(self, CTS model):
        self._children = {}

        self._log_stay_prob = LOG_PRIOR_STAY_PROB
        self._log_split_prob = LOG_PRIOR_SPLIT_PROB

        # Back pointer to the CTS model object.
        self._model = model
        self.estimator = Estimator(model)

    cdef double update(self, int[:] context, int symbol):
        """Updates this node and its children.
        Recursively updates estimators for all suffixes of context. Each
        estimator is updated with the given symbol. Also updates the mixing
        weights.
        """
        lp_estimator = self.estimator.update(symbol)

        # If not a leaf node, recurse, creating nodes as needed.
        cdef CTSNode child
        cdef double lp_child
        cdef double lp_node
        if len(context) > 0:
            child = self.get_child(context[-1])
            lp_child = child.update(context[:-1], symbol)
            lp_node = self.mix_prediction(lp_estimator, lp_child)

            self.update_switching_weights(lp_estimator, lp_child)

            return lp_node
        else:
            self._log_stay_prob = 0.0
            return lp_estimator

    cdef double log_prob(self, int[:] context, int symbol):
        cdef double lp_estimator = log(self.estimator.prob(symbol))

        if len(context) > 0:
            child = self.get_child(context[-1])

            lp_child = child.log_prob(context[:-1], symbol)

            return self.mix_prediction(lp_estimator, lp_child)
        else:
            return lp_estimator


    cdef CTSNode get_child(self, int symbol, bint allocate=True):
        cdef CTSNode node = self._children.get(symbol, None)

        if node is None and allocate:
            node = CTSNode(self._model)
            self._children[symbol] = node

        return node

    cdef double mix_prediction(self, double lp_estimator, double lp_child):
        cdef double numerator = log_add(lp_estimator + self._log_stay_prob,
                                     lp_child + self._log_split_prob)
        cdef double denominator = log_add(self._log_stay_prob,
                                       self._log_split_prob)

        return numerator - denominator

    cdef void update_switching_weights(self, double lp_estimator, double lp_child):
        cdef log_alpha = self._model.log_alpha
        cdef log_1_minus_alpha = self._model.log_1_minus_alpha

        # Avoid numerical issues with alpha = 1. This reverts to straight up
        # weighting.
        if log_1_minus_alpha == 0:
            self._log_stay_prob += lp_estimator
            self._log_split_prob += lp_child

        else:<commit_msg>replace len with shape to eliminate more python overhead<commit_after>#cython: initializedcheck=False
#cython: boundscheck=False
#cython: nonecheck=False
#cython: cdivision=True
# CTS code adapted from https://github.com/mgbellemare/SkipCTS
cimport cython
import numpy as np
cimport numpy as np

from libc.math cimport log, exp

from skimage.transform import resize


# Parameters of the CTS model. For clarity, we take these as constants.
cdef double PRIOR_STAY_PROB = 0.5
cdef double PRIOR_SPLIT_PROB = 0.5
cdef double LOG_PRIOR_STAY_PROB = log(PRIOR_STAY_PROB)
cdef double LOG_PRIOR_SPLIT_PROB = log(1.0 - PRIOR_STAY_PROB)
# Sampling parameter. The maximum number of rejections before we give up and
# sample from the root estimator.
cdef int MAX_SAMPLE_REJECTIONS = 25

@cython.wraparound(False)
cdef double get_prior(char* prior_name, int alphabet_size):
    if prior_name == <char*>'perks':
        return 1.0 / <double>alphabet_size
    elif prior_name == <char*>'jeffreys':
        return 0.5
    else: #use laplace prior
        return 1.0

@cython.wraparound(False)
cdef double log_add(double log_x, double log_y):
    """Given log x and log y, returns log(x + y)."""
    # Swap variables so log_y is larger.
    if log_x > log_y:
        log_x, log_y = log_y, log_x

    cdef double delta = log_y - log_x
    return log(1 + exp(delta)) + log_x if delta <= 50.0 else log_y


class Error(Exception):
    """Base exception for the `cts` module."""
    pass

@cython.wraparound(False)
cdef class Estimator:
    """The estimator for a CTS node.
    This implements a Dirichlet-multinomial model with specified prior. This
    class does not perform alphabet checking, and will return invalid
    probabilities if it is ever fed more than `model.alphabet_size` distinct
    symbols.
    Args:
        model: Reference to CTS model. We expected model.symbol_prior to be
            a `float`.
    """
    cdef double count_total
    cdef CTS _model
    cdef double[:] counts
    
    def __init__(self, CTS model):
        self.counts = np.ones([model.alphabet_size], dtype=np.double)*model.symbol_prior
        self.count_total = model.alphabet_size * model.symbol_prior
        self._model = model

    cdef double prob(self, int symbol):
        """Returns the probability assigned to this symbol."""
        return self.counts[symbol] / self.count_total

    cdef double update(self, int symbol):
        """Updates our count for the given symbol."""
        cdef double prob = self.prob(symbol)
        cdef double log_prob = log(prob)
        self.counts[symbol] = self.counts[symbol] + 1.0
        self.count_total += 1.0
        return log_prob

    def __getstate__(self):
        return self.count_total, self._model, self.counts

    def __setstate__(self, state):
        self.count_total, self._model, self.counts = state
            
    
cdef class CTSNode:
    """A node in the CTS tree.
    Each node contains a base Dirichlet estimator holding the statistics for
    this particular context, and pointers to its children.
    """
    cdef double _log_stay_prob
    cdef double _log_split_prob
    cdef CTS _model
    cdef Estimator estimator
    cdef dict _children

    def __init__(self, CTS model):
        self._children = {}

        self._log_stay_prob = LOG_PRIOR_STAY_PROB
        self._log_split_prob = LOG_PRIOR_SPLIT_PROB

        # Back pointer to the CTS model object.
        self._model = model
        self.estimator = Estimator(model)

    cdef double update(self, int[:] context, int symbol):
        """Updates this node and its children.
        Recursively updates estimators for all suffixes of context. Each
        estimator is updated with the given symbol. Also updates the mixing
        weights.
        """
        lp_estimator = self.estimator.update(symbol)

        # If not a leaf node, recurse, creating nodes as needed.
        cdef CTSNode child
        cdef double lp_child
        cdef double lp_node
        if context.shape[0] > 0:
            child = self.get_child(context[-1])
            lp_child = child.update(context[:-1], symbol)
            lp_node = self.mix_prediction(lp_estimator, lp_child)

            self.update_switching_weights(lp_estimator, lp_child)

            return lp_node
        else:
            self._log_stay_prob = 0.0
            return lp_estimator

    cdef double log_prob(self, int[:] context, int symbol):
        cdef double lp_estimator = log(self.estimator.prob(symbol))

        if context.shape[0] > 0:
            child = self.get_child(context[-1])

            lp_child = child.log_prob(context[:-1], symbol)

            return self.mix_prediction(lp_estimator, lp_child)
        else:
            return lp_estimator


    cdef CTSNode get_child(self, int symbol, bint allocate=True):
        cdef CTSNode node = self._children.get(symbol, None)

        if node is None and allocate:
            node = CTSNode(self._model)
            self._children[symbol] = node

        return node

    cdef double mix_prediction(self, double lp_estimator, double lp_child):
        cdef double numerator = log_add(lp_estimator + self._log_stay_prob,
                                     lp_child + self._log_split_prob)
        cdef double denominator = log_add(self._log_stay_prob,
                                       self._log_split_prob)

        return numerator - denominator

    cdef void update_switching_weights(self, double lp_estimator, double lp_child):
        cdef double log_alpha = self._model.log_alpha
        cdef double log_1_minus_alpha = self._model.log_1_minus_alpha

        # Avoid numerical issues with alpha = 1. This reverts to straight up
        # weighting.
        if log_1_minus_alpha == 0:
            self._log_stay_prob += lp_estimator
            self._log_split_prob += lp_child

        else:
{'fraction_non_alphanumeric': 0.07704847704847705, 'fraction_numerical': 0.006177606177606178, 'mean_word_length': 2.6198757763975156, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 15, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3349348', 'n_tokens_mistral': 3678, 'n_tokens_neox': 3531, 'n_words': 1067}