<commit_before>from libcpp.vector cimport vector
from libcpp.utility cimport pair
from cython.operator cimport dereference as deref, preincrement as inc

cdef extern from 'distributions/clustering.hpp':
    cppclass rng_t

    cppclass Assignments_cc "distributions::Clustering<int>::Assignments":
        Assignments_cc() nogil except +
        Assignments_cc(Assignments_cc &) nogil except +
        cppclass iterator:
            pair[int, int] & operator*() nogil
            iterator operator++() nogil
            bint operator==(iterator) nogil
            bint operator!=(iterator) nogil
        int & operator[](int) nogil
        iterator begin() nogil
        iterator end() nogil
        void clear() nogil
        bint empty() nogil

    cdef vector[int] count_assignments \
            "distributions::Clustering<int>::count_assignments" \
            (Assignments_cc & assignments) nogil

    cppclass PitmanYor_cc "distributions::Clustering<int>::PitmanYor":
        vector[int] sample_assignments(int size, rng_t & rng) nogil
        float score_counts(vector[int] & counts) nogil
        float score_add_value (
                int this_group_size,
                int total_group_count,
                int total_value_count) nogil
        float score_remove_value (
                int this_group_size,
                int total_group_count,
                int total_value_count) nogil


cdef class Assignments:
    cdef Assignments_cc * ptr
    def __cinit__(self):
        self.ptr = new Assignments_cc()
    def __dealloc__(self):
        del self.ptr

    def load(self, raw):
        self.ptr.clear()
        cdef int value_id
        cdef int group_id
        for value_id, group_id in raw.iteritems():
            deref(self.ptr)[value_id] = group_id

    def dump(self):
        cdef dict raw = {}
        cdef Assignments_cc.iterator i = self.ptr.begin()
        while i != self.ptr.end():
            deref(self.ptr)[deref(i).first] = deref(i).second
            inc(i)
        return raw


cdef class PitmanYor:
    cdef PitmanYor_cc * ptr
    def __cinit__(self):
        self.ptr = new PitmanYor_cc()
    def __dealloc__(self):
        del self.ptr<commit_msg>Implement cython wrapper of Pitman-Yor model<commit_after>from libcpp.vector cimport vector
from libcpp.utility cimport pair
from cython.operator cimport dereference as deref, preincrement as inc

cdef extern from 'distributions/clustering.hpp':
    cppclass rng_t

    cppclass Assignments_cc "distributions::Clustering<int>::Assignments":
        Assignments_cc() nogil except +
        cppclass iterator:
            pair[int, int] & operator*() nogil
            iterator operator++() nogil
            bint operator!=(iterator) nogil
        int & operator[](int) nogil

    cdef vector[int] count_assignments_cc \
            "distributions::Clustering<int>::count_assignments" \
            (Assignments_cc & assignments) nogil

    cppclass PitmanYor_cc "distributions::Clustering<int>::PitmanYor":
        float alpha
        float d
        vector[int] sample_assignments(int size, rng_t & rng) nogil
        float score_counts(vector[int] & counts) nogil
        float score_add_value (
                int group_size,
                int group_count,
                int sample_size) nogil
        float score_remove_value (
                int group_size,
                int group_count,
                int sample_size) nogil


cpdef list count_assignments(dict assignments):
    cdef Assignments_cc assignments_cc
    cdef int value_id
    cdef int group_id
    for value_id, group_id in assignments.iteritems():
        assignments_cc[value_id] = group_id
    cdef vector[int] counts_cc = count_assignments_cc(assignments_cc)
    cdef list counts = counts_cc
    return counts


cdef class PitmanYor:
    cdef PitmanYor_cc * ptr
    def __cinit__(self):
        self.ptr = new PitmanYor_cc()
    def __dealloc__(self):
        del self.ptr

    def __init__(self, float alpha=0.5, float d=0.0):
        self.ptr.alpha = alpha
        self.ptr.d = d

    property alpha:
        def __get__(self):
            return self.ptr.alpha
        def __set__(self, float value):
            self.ptr.alpha = value

    property d:
        def __get__(self):
            return self.ptr.d
        def __set__(self, float value):
            self.ptr.d = value

    def score_counts(self, list counts):
        cdef vector[int] counts_cc = counts
        cdef float score = self.ptr.score_counts(counts_cc)
        return score

    def score_add_value(
            self,
            int group_size,
            int group_count,
            int sample_size):
        return self.ptr.score_add_value(group_size, group_count, sample_size)

    def score_remove_value(
            self,
            int group_size,
            int group_count,
            int sample_size):
        return self.ptr.score_remove_value(
            group_size,
            group_count,
            sample_size)
{'fraction_non_alphanumeric': 0.07576668671076368, 'fraction_numerical': 0.0008017638805371818, 'mean_word_length': 2.2423651721897335, 'pattern_counts': {'":': 4, '<': 9, '<?xml version=': 0, '>': 9, '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': '4416699', 'n_tokens_mistral': 1520, 'n_tokens_neox': 1462, 'n_words': 374}