<commit_before>import logging

from models import Dictionary, MessageWord, Word, MessageTopic
from msgvis.apps.corpus.models import Dataset


logger = logging.getLogger(__name__)

__all__ = ['get_twitter_context', 'get_chat_context', 'Dictionary']

_stoplist = None


def get_stoplist():
    global _stoplist
    if not _stoplist:
        from nltk.corpus import stopwords

        _stoplist = stopwords.words('english')
    return _stoplist


class DbTextIterator(object):
    def __init__(self, queryset):
        self.queryset = queryset
        self.current_position = 0
        self.current = None

    def __iter__(self):
        self.current_position = 0
        for msg in self.queryset.iterator():
            self.current = msg
            self.current_position += 1
            if self.current_position % 10000 == 0:
                logger.info("Iterating through database texts: item %d" % self.current_position)

            yield msg.text


class DbWordVectorIterator(object):
    def __init__(self, dictionary, freq_field='tfidf'):
        self.dictionary = dictionary
        self.freq_field = freq_field
        self.current_message_id = None
        self.current_vector = None

    def __iter__(self):
        qset = MessageWord.objects.filter(dictionary=self.dictionary).order_by('message')
        self.current_message_id = None
        self.current_vector = []
        current_position = 0
        for mw in qset.iterator():
            message_id = mw.message_id
            word_idx = mw.word_index
            freq = getattr(mw, self.freq_field)

            if self.current_message_id is None:
                self.current_message_id = message_id
                self.current_vector = []

            if self.current_message_id != message_id:
                yield self.current_vector
                self.current_vector = []
                self.current_message_id = message_id
                current_position += 1

                if current_position % 10000 == 0:
                    logger.info("Iterating through database word-vectors: item %d" % current_position)

            self.current_vector.append((word_idx, freq))

        # one more extra one
        yield self.current_vector

    def __len__(self):
        from django.db.models import Count

        count = MessageWord.objects \
            .filter(dictionary=self.dictionary) \
            .aggregate(Count('message', distinct=True))

        if count:
            return count['message__count']


class Tokenizer(object):
    def __init__(self, texts=None, stoplist=None):
        self.texts = texts
        self.stoplist = stoplist
        self.max_length = Word._meta.get_field('text').max_length
        if self.stoplist is None:
            self.stoplist = []

    def __iter__(self):
        if self.texts is None:
            raise RuntimeError("Tokenizer can only iterate if given texts")

        for text in self.texts:
            yield self.tokenize(text)

    def tokenize(self, text):
        words = []
        for word in self.split(text.lower()):
            if word not in self.stoplist:
                if len(word) >= self.max_length:
                    word = word[:self.max_length - 1]
                words.append(word)
        return words

    def split(self, text):
        return text.split()


class WordTokenizer(Tokenizer):
    def __init__(self, texts=None, stoplist=None):
        super(WordTokenizer, self).__init__(texts, stoplist)

        import nltk

        self._tokenize = nltk.word_tokenize

    def split(self, text):
        return self._tokenize(text)


class SimpleTokenizer(Tokenizer):
    def __init__(self, texts=None, stoplist=None):
        super(SimpleTokenizer, self).__init__(texts, stoplist)

        from nltk.tokenize import sent_tokenize, wordpunct_tokenize
        self._sent_tokenize = sent_tokenize
        self._tokenize = wordpunct_tokenize

        import re
        self._strip_punct = re.compile(r'[^\w\s]')

    def split(self, text):
        # split into sentence, then remove all punctuation
        sents = (self._strip_punct.sub('', sent) for sent in self._sent_tokenize(text))
        # then split on non-words
        return [token for sent in sents for token in self._tokenize(sent)]


class TopicContext(object):
    def __init__(self, name, queryset, tokenizer, minimum_frequency=2, stoplist=None):
        self.name = name
        self.queryset = queryset
        self.tokenizer = tokenizer
        self.stoplist = stoplist
        self.minimum_frequency = minimum_frequency

    def queryset_str(self):
        return str(self.queryset.query)

    def get_dict_settings(self):
        settings = dict(
            name=self.name,
            tokenizer=self.tokenizer.__name__,
            dataset=self.queryset_str(),
            stoplist=self.stoplist is not None,
            minimum_frequency=self.minimum_frequency
        )

        import json

        return json.dumps(settings, sort_keys=True)


    def find_dictionary(self):
        results = Dictionary.objects.filter(settings=self.get_dict_settings())
        return results.last()


    def build_dictionary(self):
        texts = DbTextIterator(self.queryset)

        tokenized_texts = self.tokenizer(texts, stoplist=self.stoplist)

        return Dictionary._create_from_texts(tokenized_texts=tokenized_texts,
                                             name=self.name,
                                             minimum_frequency=self.minimum_frequency,
                                             dataset=self.queryset.model.__name__,
                                             settings=self.get_dict_settings())

    def bows_exist(self, dictionary):
        return MessageWord.objects.filter(dictionary=dictionary).exists()


    def build_bows(self, dictionary):
        texts = DbTextIterator(self.queryset)
        tokenized_texts = self.tokenizer(texts, stoplist=self.stoplist)

        dictionary._vectorize_corpus(queryset=self.queryset,
                                     tokenizer=tokenized_texts)

    def build_lda(self, dictionary, num_topics=30, **kwargs):
        corpus = DbWordVectorIterator(dictionary)
        return dictionary._build_lda(self.name, corpus, num_topics=num_topics, **kwargs)

    def apply_lda(self, dictionary, model, lda=None):
        corpus = DbWordVectorIterator(dictionary)
        return dictionary._apply_lda(model, corpus, lda=lda)

    def evaluate_lda(self, dictionary, model, lda=None):
        corpus = DbWordVectorIterator(dictionary)
        return dictionary._evaluate_lda(model, corpus, lda=lda)


def standard_topic_pipeline(context, num_topics, **kwargs):
    dictionary = context.find_dictionary()
    if dictionary is None:
        dictionary = context.build_dictionary()

    if not context.bows_exist(dictionary):
        context.build_bows(dictionary)

    model, lda = context.build_lda(dictionary, num_topics=num_topics, **kwargs)
    context.apply_lda(dictionary, model, lda)
    context.evaluate_lda(dictionary, model, lda)


def default_topic_context(name, dataset_id):
    dataset = Dataset.objects.get(pk=dataset_id)
    queryset = dataset.message_set.filter(language__code='en')

    return TopicContext(name=name, queryset=queryset,
                        tokenizer=SimpleTokenizer,
                        stoplist=get_stoplist(),
                        minimum_frequency=4)<commit_msg>remove urls<commit_after>import logging

from models import Dictionary, MessageWord, Word, MessageTopic
from msgvis.apps.corpus.models import Dataset


logger = logging.getLogger(__name__)

__all__ = ['get_twitter_context', 'get_chat_context', 'Dictionary']

_stoplist = None


def get_stoplist():
    global _stoplist
    if not _stoplist:
        from nltk.corpus import stopwords

        _stoplist = stopwords.words('english')
    return _stoplist


class DbTextIterator(object):
    def __init__(self, queryset):
        self.queryset = queryset
        self.current_position = 0
        self.current = None

    def __iter__(self):
        self.current_position = 0
        for msg in self.queryset.iterator():
            self.current = msg
            self.current_position += 1
            if self.current_position % 10000 == 0:
                logger.info("Iterating through database texts: item %d" % self.current_position)

            yield msg.text


class DbWordVectorIterator(object):
    def __init__(self, dictionary, freq_field='tfidf'):
        self.dictionary = dictionary
        self.freq_field = freq_field
        self.current_message_id = None
        self.current_vector = None

    def __iter__(self):
        qset = MessageWord.objects.filter(dictionary=self.dictionary).order_by('message')
        self.current_message_id = None
        self.current_vector = []
        current_position = 0
        for mw in qset.iterator():
            message_id = mw.message_id
            word_idx = mw.word_index
            freq = getattr(mw, self.freq_field)

            if self.current_message_id is None:
                self.current_message_id = message_id
                self.current_vector = []

            if self.current_message_id != message_id:
                yield self.current_vector
                self.current_vector = []
                self.current_message_id = message_id
                current_position += 1

                if current_position % 10000 == 0:
                    logger.info("Iterating through database word-vectors: item %d" % current_position)

            self.current_vector.append((word_idx, freq))

        # one more extra one
        yield self.current_vector

    def __len__(self):
        from django.db.models import Count

        count = MessageWord.objects \
            .filter(dictionary=self.dictionary) \
            .aggregate(Count('message', distinct=True))

        if count:
            return count['message__count']


class Tokenizer(object):
    def __init__(self, texts, *filters):
        """
        Filters is a list of objects which can be used like sets
        to determine if a word should be removed: if word in filter, then word will
        be ignored.
        """
        self.texts = texts
        self.filters = filters
        self.max_length = Word._meta.get_field('text').max_length

    def __iter__(self):
        if self.texts is None:
            raise RuntimeError("Tokenizer can only iterate if given texts")

        for text in self.texts:
            yield self.tokenize(text)

    def tokenize(self, text):
        words = []

        for word in self.split(text.lower()):
            filter_out = False
            for f in self.filters:
                if word in f:
                    filter_out = True
                    break

            if filter_out:
                # skip this word
                continue

            if len(word) >= self.max_length:
                word = word[:self.max_length - 1]

            words.append(word)

        return words

    def split(self, text):
        return text.split()


class WordTokenizer(Tokenizer):
    def __init__(self, *args, **kwargs):
        super(WordTokenizer, self).__init__(*args, **kwargs)

        import nltk

        self._tokenize = nltk.word_tokenize

    def split(self, text):
        return self._tokenize(text)


class SimpleTokenizer(Tokenizer):
    def __init__(self, *args, **kwargs):
        super(SimpleTokenizer, self).__init__(*args, **kwargs)

        from nltk.tokenize import sent_tokenize, wordpunct_tokenize
        self._sent_tokenize = sent_tokenize
        self._tokenize = wordpunct_tokenize

        import re
        self._strip_punct = re.compile(r'[^\w\s]')

    def split(self, text):
        # split into sentence, then remove all punctuation
        sents = (self._strip_punct.sub('', sent) for sent in self._sent_tokenize(text))
        # then split on non-words
        return [token for sent in sents for token in self._tokenize(sent)]


class TopicContext(object):
    def __init__(self, name, queryset, tokenizer, filters, minimum_frequency=2):
        self.name = name
        self.queryset = queryset
        self.tokenizer = tokenizer
        self.filters = filters
        self.minimum_frequency = minimum_frequency

    def queryset_str(self):
        return str(self.queryset.query)

    def get_dict_settings(self):
        settings = dict(
            name=self.name,
            tokenizer=self.tokenizer.__name__,
            dataset=self.queryset_str(),
            filters=repr(self.filters),
            minimum_frequency=self.minimum_frequency
        )

        import json

        return json.dumps(settings, sort_keys=True)


    def find_dictionary(self):
        results = Dictionary.objects.filter(settings=self.get_dict_settings())
        return results.last()


    def build_dictionary(self):
        texts = DbTextIterator(self.queryset)

        tokenized_texts = self.tokenizer(texts, *self.filters)

        return Dictionary._create_from_texts(tokenized_texts=tokenized_texts,
                                             name=self.name,
                                             minimum_frequency=self.minimum_frequency,
                                             dataset=self.queryset.model.__name__,
                                             settings=self.get_dict_settings())

    def bows_exist(self, dictionary):
        return MessageWord.objects.filter(dictionary=dictionary).exists()


    def build_bows(self, dictionary):
        texts = DbTextIterator(self.queryset)
        tokenized_texts = self.tokenizer(texts, *self.filters)

        dictionary._vectorize_corpus(queryset=self.queryset,
                                     tokenizer=tokenized_texts)

    def build_lda(self, dictionary, num_topics=30, **kwargs):
        corpus = DbWordVectorIterator(dictionary)
        return dictionary._build_lda(self.name, corpus, num_topics=num_topics, **kwargs)

    def apply_lda(self, dictionary, model, lda=None):
        corpus = DbWordVectorIterator(dictionary)
        return dictionary._apply_lda(model, corpus, lda=lda)

    def evaluate_lda(self, dictionary, model, lda=None):
        corpus = DbWordVectorIterator(dictionary)
        return dictionary._evaluate_lda(model, corpus, lda=lda)


class LambdaWordFilter(object):
    def __init__(self, fn):
        self.fn = fn

    def __contains__(self, item):
        return self.fn(item)


def standard_topic_pipeline(context, num_topics, **kwargs):
    dictionary = context.find_dictionary()
    if dictionary is None:
        dictionary = context.build_dictionary()

    if not context.bows_exist(dictionary):
        context.build_bows(dictionary)

    model, lda = context.build_lda(dictionary, num_topics=num_topics, **kwargs)
    context.apply_lda(dictionary, model, lda)
    context.evaluate_lda(dictionary, model, lda)


def default_topic_context(name, dataset_id):
    dataset = Dataset.objects.get(pk=dataset_id)
    queryset = dataset.message_set.filter(language__code='en')

    filters = [
        set(get_stoplist()),
        LambdaWordFilter(lambda word: word.startswith('http') and len(word) > 4)
    ]

    return TopicContext(name=name, queryset=queryset,
                        tokenizer=SimpleTokenizer,
                        filters=filters,
                        minimum_frequency=4)
{'fraction_non_alphanumeric': 0.08687434827945777, 'fraction_numerical': 0.00293274244004171, 'mean_word_length': 2.3199913457377757, 'pattern_counts': {'":': 0, '<': 3, '<?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': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6173565', 'n_tokens_mistral': 4432, 'n_tokens_neox': 4163, 'n_words': 935}