<commit_before>from lxml import etree
import abc
import re
import requests

import gensim.models as models
import gensim.models.word2vec as word2vec



# Abstract Rewriter class
# TODO: abstract out more functionality here
class Rewriter(object):
    # make it an abstract class
    __metaclass__ = abc.ABCMeta

    def rewrite(self, term):
        """
        Subclasses must implement!
        """
        raise NotImplementedError("Subclasses must implement!")


class ControlRewriter(Rewriter):
    """
    A rewriter that's basically a no-op. Just returns the term you give it.
    """

    def rewrite(self, term):
        return [term]



class WikipediaRewriter(Rewriter):
    """
    A class to rewrite queries using Wikipedia's category api.
    """
    WIKI_BASE = 'https://en.wikipedia.org/w/api.php?format=xml&action=query&prop=categories&titles='

    def clean_category(self, x):
        return x.replace('Category:','')

    def rewrite(self, term):
        """
        Given a base term, returns a list of related terms based on the Wikipedia
        category API.
        """
        api = self.WIKI_BASE + term
        r = requests.get(api)
        tree = etree.fromstring(r.text)
        # TODO: ignore wikipedia internal categories
        # https://en.wikipedia.org/wiki/Category:Tracking_categories
        wikipedia_results = [self.clean_category(x.get('title')) for x in tree.findall('.//cl')]

        # append the original term just for completeness
        return wikipedia_results + [term]


class Word2VecRewriter(Rewriter):
    """
    A class to rewrite queries using Word2Vec trained on a corpus gathered
    from Wikipedia.
    """

    # where the model will be stored
    MODEL_PATH = 'models/word2vec/word2vec'

    def __init__(self, create=False):
        """
        Initializes the rewriter. If create is True, this generates a new Word2Vec
        model (which takes a really long time to build.) If false, this loads
        an existing model we already saved.
        """

        if create:
            # generate a new Word2Vec model... takes a while!
            corpus = word2vec.Text8Corpus('datasets/enwik8')
            self.model = word2vec.Word2Vec(corpus, workers=8)
            self.model.save(self.MODEL_PATH)
        else:
            self.model = word2vec.Word2Vec.load(self.MODEL_PATH)


    def clean(self, string):
        """
        Cleans up a string of text you get from wikipedia, which is often formatted like
        `*[[Games`.
        """
        # remove [[, ]], *
        # TODO: also remove commas and punctuation
        # (though apostrophes and dashes can stay, i guess)
        return re.sub(r"(\[\[|\]\]|\*)", "", string)


    def rewrite(self, term):
        # try using the model to rewrite the term
        cleaned_results = []
        try:
            # most_similar returns an array of tuples... extract just the name, which is index 0
            # we also have to clean the text
            raw_results = self.model.most_similar(positive=[term], topn=10)
            cleaned_results = [self.clean(r[0]) for r in raw_results]
        except KeyError as k:
            # the word wasn't found in the model... must be too niche.<commit_msg>Trying to clean up the wikipedia api. Removing some of the obvious faulty category iterms. For cognition, we now get 4 results instead of 1 for udacity. Biochemistry is now 1 instead of 0 results.<commit_after>from lxml import etree
import abc
import re
import requests

import gensim.models as models
import gensim.models.word2vec as word2vec
import secure


# Abstract Rewriter class
# TODO: abstract out more functionality here
class Rewriter(object):
    # make it an abstract class
    __metaclass__ = abc.ABCMeta

    def rewrite(self, term):
        """
        Subclasses must implement!
        """
        raise NotImplementedError("Subclasses must implement!")


class ControlRewriter(Rewriter):
    """
    A rewriter that's basically a no-op. Just returns the term you give it.
    """

    def rewrite(self, term):
        return [term]



class WikipediaRewriter(Rewriter):
    """
    A class to rewrite queries using Wikipedia's category api.
    """
    WIKI_BASE = 'https://en.wikipedia.org/w/api.php?format=xml&action=query&prop=categories&titles='

    def clean_category(self, x):
        return x.replace('Category:','')

    def rewrite(self, term):
        """
        Given a base term, returns a list of related terms based on the Wikipedia
        category API.
        """
        api = self.WIKI_BASE + term
        r = requests.get(api)
        tree = etree.fromstring(r.text)
        # TODO: ignore wikipedia internal categories
        # https://en.wikipedia.org/wiki/Category:Tracking_categories
        wikipedia_results = [self.clean_category(x.get('title')) for x in tree.findall('.//cl')]

        # Words that identify a need to drop the category
        dropwords = ['articles','wikipedia','accuracy','articles','statements','magic','pages','authors','editors','appearances','redirects','cs1']
        dropwords.append(term)
        wikipedia_results = [w.split('Category:')[-1].lower() for w in wikipedia_results if not any(d in w.lower() for d in dropwords)]

        # append the original term just for completeness
        print wikipedia_results
        print ' '.join(wikipedia_results + [term])
        return wikipedia_results + [term]


class Word2VecRewriter(Rewriter):
    """
    A class to rewrite queries using Word2Vec trained on a corpus gathered
    from Wikipedia.
    """

    # where the model will be stored
    MODEL_PATH = secure.MODEL_PATH_BASE+'word2vec/word2vec'

    def __init__(self, create=False):
        """
        Initializes the rewriter. If create is True, this generates a new Word2Vec
        model (which takes a really long time to build.) If false, this loads
        an existing model we already saved.
        """

        if create:
            # generate a new Word2Vec model... takes a while!
            corpus = word2vec.Text8Corpus(secure.DATASET_PATH_BASE + 'enwik8')
            self.model = word2vec.Word2Vec(corpus, workers=8)
            self.model.save(self.MODEL_PATH)
        else:
            self.model = word2vec.Word2Vec.load(self.MODEL_PATH)


    def clean(self, string):
        """
        Cleans up a string of text you get from wikipedia, which is often formatted like
        `*[[Games`.
        """
        # remove [[, ]], *
        # TODO: also remove commas and punctuation
        # (though apostrophes and dashes can stay, i guess)
        return re.sub(r"(\[\[|\]\]|\*)", "", string)


    def rewrite(self, term):
        # try using the model to rewrite the term
        cleaned_results = []
        try:
            # most_similar returns an array of tuples... extract just the name, which is index 0
            # we also have to clean the text
            raw_results = self.model.most_similar(positive=[term], topn=10)
            cleaned_results = [self.clean(r[0]) for r in raw_results]
        except KeyError as k:
            # the word wasn't found in the model... must be too niche.
{'fraction_non_alphanumeric': 0.0982756203560914, 'fraction_numerical': 0.006448899481284172, 'mean_word_length': 2.5108267716535435, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 4}, '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': '7074922', 'n_tokens_mistral': 2092, 'n_tokens_neox': 1976, 'n_words': 676}