<commit_before>from collections import defaultdict
from string import punctuation
from functools import lru_cache

import kenlm
import nltk
from nltk.tokenize import word_tokenize
from nltk.corpus import wordnet as wn

from qanta.pattern3 import pluralize
from qanta.extractors.abstract import AbstractFeatureExtractor
from qanta.util.constants import KEN_LM
from qanta.util.environment import data_path
from qanta.datasets.quiz_bowl import QuizBowlDataset
from qanta.wikipedia.cached_wikipedia import CachedWikipedia
from qanta import logging
from qanta.preprocess import format_guess

from clm.lm_wrapper import kTOKENIZER, LanguageModelBase


log = logging.get(__name__)


@lru_cache(maxsize=None)
def get_states():
    states = set()
    for ii in wn.synset("American_state.n.1").instance_hyponyms():
        for jj in ii.lemmas():
            name = jj.name()
            if len(name) > 2 and "_" not in name:
                states.add(name)
            elif name.startswith("New_"):
                states.add(name.replace("New_", ""))
    return states


def build_lm_data(output):
    cw = CachedWikipedia()

    with open(output, 'w') as o:
        dataset = QuizBowlDataset(1)
        training_data = dataset.training_data()
        train_pages = {format_guess(g) for g in training_data[1]}
        for i, page in enumerate(train_pages):
            content = cw[page].content
            for sentence in nltk.sent_tokenize(content):
                o.write("%s\n" % " ".join(kTOKENIZER(sentence.lower())))

            if i % 1000 == 0:
                log.info("%i\t%s" % (i, page))


def find_references(sentence, padding=5):
    tags = nltk.pos_tag(word_tokenize(sentence))
    tags.append(("END", "V"))
    states = get_states()

    references_found = []
    this_ref_start = -1
    for i, pair in enumerate(tags):
        word, tag = pair
        if word.lower() == 'this' or word.lower() == 'these':
            this_ref_start = i
        elif all(x in punctuation for x in word):
            continue
        elif word in states:
            continue
        elif this_ref_start >= 0 and tag.startswith('NN') and \
                not tags[i + 1][1].startswith('NN'):
            references_found.append((this_ref_start, i))
            this_ref_start = -1
        elif tag.startswith('V'):
            this_ref_start = -1

    for start, stop in references_found:
        yield (" ".join(LanguageModelBase.normalize_title('', x[0])
                        for x in tags[max(0, start - padding):start]),
               " ".join(LanguageModelBase.normalize_title('', x[0])
                        for x in tags[start:stop + 1]),
               " ".join(LanguageModelBase.normalize_title('', x[0])
                        for x in tags[stop + 1:stop + padding + 1]))


class Mentions(AbstractFeatureExtractor):
    def __init__(self):
        super().__init__()
        self.initialized = False
        self.refex_count = defaultdict(int)
        self.refex_lookup = defaultdict(set)
        self._lm = None
        self.generate_refexs()
        self.kenlm_path = data_path(KEN_LM)

    @property
    def name(self):
        return 'mentions'

    @property
    def lm(self):
        if not self.initialized:
            self._lm = kenlm.LanguageModel(self.kenlm_path)
            self.initialized = True
        return self._lm

    def score_guesses(self, guesses, text):
        # Find mentions if the text has changed
        pre = []
        ment = []
        suf = []
        for p, m, s in find_references(text):
            # Exclude too short mentions
            if len(m.strip()) > 3:
                pre.append(p.lower())
                suf.append(s.lower())
                ment.append(m.lower())<commit_msg>Move kenlm so its not required to be installed for import time<commit_after>from collections import defaultdict
from string import punctuation
from functools import lru_cache

import nltk
from nltk.tokenize import word_tokenize
from nltk.corpus import wordnet as wn

from qanta.pattern3 import pluralize
from qanta.extractors.abstract import AbstractFeatureExtractor
from qanta.util.constants import KEN_LM
from qanta.util.environment import data_path
from qanta.datasets.quiz_bowl import QuizBowlDataset
from qanta.wikipedia.cached_wikipedia import CachedWikipedia
from qanta import logging
from qanta.preprocess import format_guess

from clm.lm_wrapper import kTOKENIZER, LanguageModelBase


log = logging.get(__name__)


@lru_cache(maxsize=None)
def get_states():
    states = set()
    for ii in wn.synset("American_state.n.1").instance_hyponyms():
        for jj in ii.lemmas():
            name = jj.name()
            if len(name) > 2 and "_" not in name:
                states.add(name)
            elif name.startswith("New_"):
                states.add(name.replace("New_", ""))
    return states


def build_lm_data(output):
    cw = CachedWikipedia()

    with open(output, 'w') as o:
        dataset = QuizBowlDataset(1)
        training_data = dataset.training_data()
        train_pages = {format_guess(g) for g in training_data[1]}
        for i, page in enumerate(train_pages):
            content = cw[page].content
            for sentence in nltk.sent_tokenize(content):
                o.write("%s\n" % " ".join(kTOKENIZER(sentence.lower())))

            if i % 1000 == 0:
                log.info("%i\t%s" % (i, page))


def find_references(sentence, padding=5):
    tags = nltk.pos_tag(word_tokenize(sentence))
    tags.append(("END", "V"))
    states = get_states()

    references_found = []
    this_ref_start = -1
    for i, pair in enumerate(tags):
        word, tag = pair
        if word.lower() == 'this' or word.lower() == 'these':
            this_ref_start = i
        elif all(x in punctuation for x in word):
            continue
        elif word in states:
            continue
        elif this_ref_start >= 0 and tag.startswith('NN') and \
                not tags[i + 1][1].startswith('NN'):
            references_found.append((this_ref_start, i))
            this_ref_start = -1
        elif tag.startswith('V'):
            this_ref_start = -1

    for start, stop in references_found:
        yield (" ".join(LanguageModelBase.normalize_title('', x[0])
                        for x in tags[max(0, start - padding):start]),
               " ".join(LanguageModelBase.normalize_title('', x[0])
                        for x in tags[start:stop + 1]),
               " ".join(LanguageModelBase.normalize_title('', x[0])
                        for x in tags[stop + 1:stop + padding + 1]))


class Mentions(AbstractFeatureExtractor):
    def __init__(self):
        super().__init__()
        self.initialized = False
        self.refex_count = defaultdict(int)
        self.refex_lookup = defaultdict(set)
        self._lm = None
        self.generate_refexs()
        self.kenlm_path = data_path(KEN_LM)

    @property
    def name(self):
        return 'mentions'

    @property
    def lm(self):
        import kenlm
        if not self.initialized:
            self._lm = kenlm.LanguageModel(self.kenlm_path)
            self.initialized = True
        return self._lm

    def score_guesses(self, guesses, text):
        # Find mentions if the text has changed
        pre = []
        ment = []
        suf = []
        for p, m, s in find_references(text):
            # Exclude too short mentions
            if len(m.strip()) > 3:
                pre.append(p.lower())
                suf.append(s.lower())
                ment.append(m.lower())
{'fraction_non_alphanumeric': 0.10584958217270195, 'fraction_numerical': 0.006632179334129195, 'mean_word_length': 2.2754126846220677, 'pattern_counts': {'":': 0, '<': 3, '<?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': '333440', 'n_tokens_mistral': 2349, 'n_tokens_neox': 2266, 'n_words': 578}