<commit_before>__author__ = 'jamiebrew'

from ngram import Ngram
import string
import math
import operator
import re
from collections import Counter

"""
a corpus represents information about a text as a tree indexed by string
	* each entry in the tree is an ngram object
	* the key for a multi-word ngram is the space-separated words in the ngram
"""
class Corpus(object):

    def __init__(self, text, name, max_ngram_size = 2, sort_attribute = "frequency", foresight = 0, hindsight = 2, wordcount_criterion = 1):
        self.wordcount = 0
        self.wordcount_criterion = wordcount_criterion
        self.foresight = foresight
        self.hindsight = hindsight
        self.max_ngram_size = max_ngram_size
        self.tree = self.make_tree(text)
        self.name = name
        self.sort_attribute = sort_attribute

    # returns the n words occurring most frequently in a string
    def top_words(self, n, text):
        text = text.translate(string.maketrans("",""), string.punctuation)
        words = re.findall('[a-z]+', text.lower())
        wordcounts = Counter(words)
        sorted_wordcounts = list(reversed(sorted(wordcounts.items(), key=operator.itemgetter(1))))
        return sorted_wordcounts[0:n]

    # the list of words in a given list of sentences that meet a given wordcount criterion
    def short_list(self, sentence_list, wordcount_criterion):
        # all unique words mapped to wordcount
        unique_words = {}

        # make wordcount dictionary
        for sentence in sentence_list:
            for word in sentence:
                if word not in unique_words:
                    unique_words[word] = 1
                else:
                    unique_words[word] += 1
        shortlist = set()
        for word in unique_words:
            if unique_words[word] >= wordcount_criterion:
                shortlist.add(word)
        print "Unique words in corpus:", len(unique_words)
        print "Words occurring at least %s times: %s" % (wordcount_criterion, len(shortlist))
        return shortlist

    # constructs the tree of ngrams' likelihood of following other ngrams
    def make_tree(self, str):
        sentences = self.make_sentences(str)
        shortlist = self.short_list(sentences, self.wordcount_criterion)

        T = {}
        # go through each sentence, add each word to the dictionary, incrementing length each time
        for sentence in sentences:
            sentence = ['START_SENTENCE'] + sentence
            for ngram_size in range(0, self.max_ngram_size+1):
                for start in range(0, len(sentence)):
                    end = start + ngram_size
                    if end <= len(sentence):
                        words_to_add = sentence[start:end]
                        if set(words_to_add) < shortlist and len(words_to_add) > 0: # checks that all of the words in the ngram pass criterion
                            new_ngram = " ".join(words_to_add)
                            self.add_ngram(new_ngram, T)
                            if ngram_size==1:
                                self.wordcount += 1

                            # add dictionaries of words following this ngram
                            for word_position in range(end, end+self.hindsight):
                                if word_position < len(sentence):
                                    reach = word_position - end
                                    target = T[new_ngram].after[reach]
                                    word = sentence[word_position]
                                    if word in shortlist:
                                        self.add_ngram(word,target)

                            # add dictionaries of words preceding this ngram
                            for word_position in range(start-1, start-self.foresight-1, -1):
                                if word_position >= 0:
                                   reach = start - word_position
                                   target = T[new_ngram].before[reach-1]
                                   word = sentence[word_position]
                                   if word in shortlist:
                                        self.add_ngram(word,target)

        T = self.calculate_frequencies(T)
        T = self.calculate_sig_scores(T)
        return T

    # splits the text into sentences, lowercases them and cleans punctuation
    def make_sentences(self, str):
        # split at period followed by newline or space, or question mark, or exclamation point
        sentences = str.split('.\n' or '. ' or '?' or '!')

        for sentence in range (0,len(sentences)):
            sentences[sentence] = sentences[sentence].strip('\n')
            sentences[sentence] = sentences[sentence].translate(string.maketrans("",""), string.punctuation.replace("'","")) # removes all punctuation except apostrophe
            sentences[sentence] = sentences[sentence].lower()
            sentences[sentence] = sentences[sentence].split()
        return sentences

    # adds an ngram to a given tree
    def add_ngram(self, str, tree):
        if str in tree:
            tree[str].count += 1
        else:
            tree[str] = Ngram(str, self.hindsight, self.foresight)

    def lookup_ngram(self, ngram, tree):
        if ngram in tree:
            return tree[ngram]

    # given a tree that has kept count for each word, finds and stores normalized frequencies
    def calculate_frequencies(self,T):
        for ngram_key in T:
            ngram = T[ngram_key]
            ngram.frequency = ngram.count/float(self.wordcount)
            for dict in ngram.before:
                for before_key in dict:
                    before_ngram = dict[before_key]
                    before_ngram.frequency = before_ngram.count/float(ngram.count)
            for dict in ngram.after:
                for after_key in dict:
                    after_ngram = dict[after_key]
                    after_ngram.frequency = after_ngram.count/float(ngram.count)
        return T

    # given a tree that has frequencies for each word, computes and stores the significance scores
    def calculate_sig_scores(self,T):
        for n_gram_key in T:
            n_gram = T[n_gram_key]
            for dict in n_gram.before:
                for before_key in dict:
                    before_n_gram = dict[before_key]
                    before_n_gram.sig_score = (before_n_gram.frequency/T[before_key].frequency) * math.log(n_gram.frequency+1,10)
            for dict in n_gram.after:
                for after_key in dict:
                    after_n_gram = dict[after_key]
                    after_n_gram.sig_score = (after_n_gram.frequency/T[after_key].frequency) * math.log(n_gram.frequency+1,10)
        return T

    # ranks all ngrams in a tree by the specified attribute. returns a list
    def sort_ngrams(self, tree, sort_att):
        to_sort = {}
        for ngram_key in tree:
            ngram = tree[ngram_key]
            if sort_att == 'frequency':
                to_sort[ngram_key] = ngram.frequency
            elif sort_att == 'sigscore':
                to_sort[ngram_key] = ngram.sig_score
            elif sort_att == 'count':
                to_sort[ngram_key] = ngram.count
        return sorted(to_sort, key=operator.itemgetter(0))
<commit_msg>Remove unused methods, use instance variable of wordcount_criterion<commit_after>__author__ = 'jamiebrew'

from ngram import Ngram
import string
import math
import operator
import re
from collections import Counter

"""
a corpus represents information about a text as a tree indexed by string
	* each entry in the tree is an ngram object
	* the key for a multi-word ngram is the space-separated words in the ngram
"""
class Corpus(object):

    def __init__(self, text, name='', max_ngram_size=2, sort_attribute='frequency', foresight=0, hindsight=2, wordcount_criterion=1):
        self.wordcount = 0
        self.wordcount_criterion = wordcount_criterion
        self.foresight = foresight
        self.hindsight = hindsight
        self.max_ngram_size = max_ngram_size
        self.tree = self.make_tree(text)
        self.name = name
        self.sort_attribute = sort_attribute

    # the list of words in a given list of sentences that meet a given wordcount criterion
    def short_list(self, sentence_list):
        # all unique words mapped to wordcount
        unique_words = {}

        # make wordcount dictionary
        for sentence in sentence_list:
            for word in sentence:
                if word not in unique_words:
                    unique_words[word] = 1
                else:
                    unique_words[word] += 1
        shortlist = set()
        for word in unique_words:
            if unique_words[word] >= self.wordcount_criterion:
                shortlist.add(word)
        print "Unique words in corpus:", len(unique_words)
        print "Words occurring at least %s times: %s" % (self.wordcount_criterion, len(shortlist))
        return shortlist

    # constructs the tree of ngrams' likelihood of following other ngrams
    def make_tree(self, str):
        sentences = self.make_sentences(str)
        shortlist = self.short_list(sentences)

        T = {}
        # go through each sentence, add each word to the dictionary, incrementing length each time
        for sentence in sentences:
            sentence = ['START_SENTENCE'] + sentence
            for ngram_size in range(0, self.max_ngram_size+1):
                for start in range(0, len(sentence)):
                    end = start + ngram_size
                    if end <= len(sentence):
                        words_to_add = sentence[start:end]
                        if set(words_to_add) < shortlist and len(words_to_add) > 0: # checks that all of the words in the ngram pass criterion
                            new_ngram = " ".join(words_to_add)
                            self.add_ngram(new_ngram, T)
                            if ngram_size==1:
                                self.wordcount += 1

                            # add dictionaries of words following this ngram
                            for word_position in range(end, end+self.hindsight):
                                if word_position < len(sentence):
                                    reach = word_position - end
                                    target = T[new_ngram].after[reach]
                                    word = sentence[word_position]
                                    if word in shortlist:
                                        self.add_ngram(word,target)

                            # add dictionaries of words preceding this ngram
                            for word_position in range(start-1, start-self.foresight-1, -1):
                                if word_position >= 0:
                                    reach = start - word_position
                                    target = T[new_ngram].before[reach-1]
                                    word = sentence[word_position]
                                    if word in shortlist:
                                        self.add_ngram(word,target)

        T = self.calculate_frequencies(T)
        T = self.calculate_sig_scores(T)
        return T

    # splits the text into sentences, lowercases them and cleans punctuation
    def make_sentences(self, str):
        # split at period followed by newline or space, or question mark, or exclamation point
        sentences = str.split('.\n' or '. ' or '?' or '!')

        for sentence in range (0,len(sentences)):
            sentences[sentence] = sentences[sentence].strip('\n')
            sentences[sentence] = sentences[sentence].translate(string.maketrans("",""), string.punctuation.replace("'","")) # removes all punctuation except apostrophe
            sentences[sentence] = sentences[sentence].lower()
            sentences[sentence] = sentences[sentence].split()
        return sentences

    # adds an ngram to a given tree
    def add_ngram(self, str, tree):
        if str in tree:
            tree[str].count += 1
        else:
            tree[str] = Ngram(str, self.hindsight, self.foresight)

    # given a tree that has kept count for each word, finds and stores normalized frequencies
    def calculate_frequencies(self,T):
        for ngram_key in T:
            ngram = T[ngram_key]
            ngram.frequency = ngram.count/float(self.wordcount)
            for dict in ngram.before:
                for before_key in dict:
                    before_ngram = dict[before_key]
                    before_ngram.frequency = before_ngram.count/float(ngram.count)
            for dict in ngram.after:
                for after_key in dict:
                    after_ngram = dict[after_key]
                    after_ngram.frequency = after_ngram.count/float(ngram.count)
        return T

    # given a tree that has frequencies for each word, computes and stores the significance scores
    def calculate_sig_scores(self,T):
        for n_gram_key in T:
            n_gram = T[n_gram_key]
            for dict in n_gram.before:
                for before_key in dict:
                    before_n_gram = dict[before_key]
                    before_n_gram.sig_score = (before_n_gram.frequency/T[before_key].frequency) * math.log(n_gram.frequency+1,10)
            for dict in n_gram.after:
                for after_key in dict:
                    after_n_gram = dict[after_key]
                    after_n_gram.sig_score = (after_n_gram.frequency/T[after_key].frequency) * math.log(n_gram.frequency+1,10)
        return T

{'fraction_non_alphanumeric': 0.07523557126030624, 'fraction_numerical': 0.0040488810365135455, 'mean_word_length': 1.7958427660012348, 'pattern_counts': {'":': 0, '<': 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': 35, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3561734', 'n_tokens_mistral': 3735, 'n_tokens_neox': 3533, 'n_words': 1043}