<commit_before>from sys import stdout
from math import log
import numpy as np



def _smpl_cat(d):
    """
    Old version of smpl_cat. For testing.
    Takes an array of probabilities d and returns a sample from the
    categorical distribution parameterized by d.
    """
    return np.random.multinomial(1, d).argmax()


def smpl_cat(d_cum):
    """
    Takes an array of cumurative probability distribution d and returns 
    a sample from the categorical distribution parameterized by d.
    """    
    r = np.random.random() * d_cum[-1]
    return np.searchsorted(d_cum, r)



class LDAGibbs(object):
    """
    An implementation of LDA using collapsed Gibbs sampling.

    References
    ----------    
    Griffiths, Tom. Gibbs sampling in the generative model of Latent Dirichlet Allocation.

    <NAME>. Distributed Gibbs Sampling of Latent Topic Models: The Gritty Details.

    Parameters
    ----------
    corpus : Corpus
        Source of observed data
    context_type : string
        Name of tokenization stored in `corpus` whose tokens will be
        treated as documents.
    K : int
        Number of topics. Default is `100`.
    alpha : float
        Parameter for the prior distribution of theta_d. Default is `0.01`.
    beta : float
        Parameter for the prior distribution of phi_d. Default is `0.01`.
    log_prob : boolean
        If `True`, compute the log probabilities of the corpus given
        the values of the latent variables at each iteration and
        records them in `log_prob`. Default is `True`.

    Attributes
    ----------
    W : list of integer arrays
        List of documents, which are extracted from the input Corpus object
    V : int
        Number of unique words in the corpus
    Z : list of integer arrays
        Topic assignments for every word coordinate in the corpus
    iterations : int
        Number of past iterations of the update rule
    doc_top : 2-dim floating point array
        Stores the unnormalized estimated posterior distribution over
        topics for each document in a D x K matrix
    top_word : 2-dim floating point array
        Stores the unnormalized estimated posterior distribution over
        terms for each topic in a K x V matrix
    sum_word_top : 1-dim floating point array
        Stores the sum of words over topics

    Methods
    -------
    train
        Takes an optional argument `itr`, which defaults to 1000, and
        updates the model `itr` times.
    update_z
        Takes a document index `d`, a word index `i` relative to that
        document and a word `w` and updates the model.
    z_dist
        Takes a document index `d` and a word `w` and computes the
        distribution over topics for `w` in `d`
    phi_k
        Takes a topic index `t` and returns the estimated posterior
        distribution over words for `t`
    phi_w
        Takes a word `w` and returns the estimated posterior
        distribution over topics for `w`
    theta_d
        Takes a document index `d` and returns the estimated posterior
        distribution over topics for `d`
    theta_k
        Takes a topic index `t` and returns the estimated posterior
        distribution over documents for `t`
    logp
        Compute the log probability of the corpus `W` given the
        estimated values of the latent variables `phi`, `theta` and
        `Z`

    """
    def __init__(self, corpus, context_type,
                 K=100, alpha = 0.01, beta = 0.01, 
                 log_prob=True):

        self.context_type = context_type
        self.K = K
        self.alpha = alpha
        self.beta = beta
        self.W = corpus.view_context(context_type)
        self.V = corpus.words.shape[0]
        self.iterations = 0

        if log_prob:
            self.log_prob = []
        
        # Initialize

        self.Z = [np.zeros_like(d) for d in self.W]
        self.doc_top = np.zeros((len(self.W), K)) + alpha
        self.top_word = np.zeros((K, self.V)) + beta
        self.sum_word_top = (self.V * beta) + np.zeros(K)

        for d, doc in enumerate(self.W):
            for i, w in enumerate(doc):
                self.update_z(d, i, w)


    def train(self, itr=1000, verbose=True):

        for t in xrange(self.iterations, self.iterations + itr):

            if verbose:
                stdout.write('\rIteration %d' % t)
                stdout.flush()

            self.iterations += 1
            
            if hasattr(self, 'log_prob'):
                self.log_prob.append((t, self.logp()))

            for d, doc in enumerate(self.W):
                for i, w in enumerate(doc):
                    z = self.Z[d][i]
                    self.doc_top[d, z] -= 1
                    self.top_word[z, w] -= 1
                    self.sum_word_top[z] -= 1
                    self.update_z(d, i, w)

        if verbose:
            stdout.write('\n')


    def _z_dist(self, d, w):
        """
        Old version of z_dist (returns non cumurative distribution). For testing. 
        """

        sum_word_top_inv = 1. / self.sum_word_top
        dist = (self.doc_top[d, :] *
                self.top_word[:, w] * sum_word_top_inv)
        nc = 1. / dist.sum()
        return dist * nc


    def z_dist(self, d, w):

        sum_word_top_inv = 1. / self.sum_word_top
        dist = (self.doc_top[d, :] * self.top_word[:, w]  * sum_word_top_inv)<commit_msg>fix multinom sampling<commit_after>from sys import stdout
from math import log
import numpy as np



def smpl_cat(d_cum):
    """
    Takes an array of cumurative probability distribution d and returns 
    a sample from the categorical distribution parameterized by d.
    """    
    r = np.random.random() * d_cum[-1]
    return np.searchsorted(d_cum, r)



class LDAGibbs(object):
    """
    An implementation of LDA using collapsed Gibbs sampling.

    References
    ----------    
    <NAME>. Gibbs sampling in the generative model of Latent Dirichlet Allocation.

    <NAME>. Distributed Gibbs Sampling of Latent Topic Models: The Gritty Details.

    Parameters
    ----------
    corpus : Corpus
        Source of observed data
    context_type : string
        Name of tokenization stored in `corpus` whose tokens will be
        treated as documents.
    K : int
        Number of topics. Default is `100`.
    alpha : float
        Parameter for the prior distribution of theta_d. Default is `0.01`.
    beta : float
        Parameter for the prior distribution of phi_d. Default is `0.01`.
    log_prob : boolean
        If `True`, compute the log probabilities of the corpus given
        the values of the latent variables at each iteration and
        records them in `log_prob`. Default is `True`.

    Attributes
    ----------
    W : list of integer arrays
        List of documents, which are extracted from the input Corpus object
    V : int
        Number of unique words in the corpus
    Z : list of integer arrays
        Topic assignments for every word coordinate in the corpus
    iterations : int
        Number of past iterations of the update rule
    doc_top : 2-dim floating point array
        Stores the unnormalized estimated posterior distribution over
        topics for each document in a D x K matrix
    top_word : 2-dim floating point array
        Stores the unnormalized estimated posterior distribution over
        terms for each topic in a K x V matrix
    sum_word_top : 1-dim floating point array
        Stores the sum of words over topics

    Methods
    -------
    train
        Takes an optional argument `itr`, which defaults to 1000, and
        updates the model `itr` times.
    update_z
        Takes a document index `d`, a word index `i` relative to that
        document and a word `w` and updates the model.
    z_dist
        Takes a document index `d` and a word `w` and computes the
        distribution over topics for `w` in `d`
    phi_k
        Takes a topic index `t` and returns the estimated posterior
        distribution over words for `t`
    phi_w
        Takes a word `w` and returns the estimated posterior
        distribution over topics for `w`
    theta_d
        Takes a document index `d` and returns the estimated posterior
        distribution over topics for `d`
    theta_k
        Takes a topic index `t` and returns the estimated posterior
        distribution over documents for `t`
    logp
        Compute the log probability of the corpus `W` given the
        estimated values of the latent variables `phi`, `theta` and
        `Z`

    """
    def __init__(self, corpus, context_type,
                 K=100, alpha = 0.01, beta = 0.01, 
                 log_prob=True):

        self.context_type = context_type
        self.K = K
        self.alpha = alpha
        self.beta = beta
        self.W = corpus.view_context(context_type)
        self.V = corpus.words.shape[0]
        self.iterations = 0

        if log_prob:
            self.log_prob = []
        
        # Initialize

        self.Z = [np.zeros_like(d) for d in self.W]
        self.doc_top = np.zeros((len(self.W), K)) + alpha
        self.top_word = np.zeros((K, self.V)) + beta
        self.sum_word_top = (self.V * beta) + np.zeros(K)

        for d, doc in enumerate(self.W):
            for i, w in enumerate(doc):
                self.update_z(d, i, w)


    def train(self, itr=1000, verbose=True):

        for t in xrange(self.iterations, self.iterations + itr):

            if verbose:
                stdout.write('\rIteration %d' % t)
                stdout.flush()

            self.iterations += 1
            
            if hasattr(self, 'log_prob'):
                self.log_prob.append((t, self.logp()))

            for d, doc in enumerate(self.W):
                for i, w in enumerate(doc):
                    z = self.Z[d][i]
                    self.doc_top[d, z] -= 1
                    self.top_word[z, w] -= 1
                    self.sum_word_top[z] -= 1
                    self.update_z(d, i, w)

        if verbose:
            stdout.write('\n')



    def z_dist(self, d, w):

        sum_word_top_inv = 1. / self.sum_word_top
        dist = (self.doc_top[d, :] * self.top_word[:, w]  * sum_word_top_inv)
{'fraction_non_alphanumeric': 0.08081600627697136, 'fraction_numerical': 0.007551981169085916, 'mean_word_length': 1.9217765042979942, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 6, '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': '475049', 'n_tokens_mistral': 3041, 'n_tokens_neox': 2837, 'n_words': 1106}