<commit_before>import itertools
import os
import os.path

from nltk.tag.stanford import NERTagger
import wget

from iepy.models import PreProcessSteps, Entity, EntityOccurrence
from iepy.preprocess import BasePreProcessStepRunner
from iepy.utils import DIRS, unzip_file

stanford_ner_name = 'stanford-ner-2014-01-04'
download_url_base = 'http://nlp.stanford.edu/software/'


class NonTokenizingNERTagger(NERTagger):

    @property
    def _cmd(self):
        old = super(NonTokenizingNERTagger, self)._cmd
        old = old + ["-tokenizerFactory", "edu.stanford.nlp.process.WhitespaceTokenizer"]
        return old


class NERRunner(BasePreProcessStepRunner):
    """Wrapper to insert a generic callable sentence NER tagger into the pipeline.
    """
    # TODO: rename to ner
    step = PreProcessSteps.nerc

    def __init__(self, ner, override=False):
        self.override = override
        self.ner = ner

    def __call__(self, doc):
        # this step does not necessarily requires PreProcessSteps.tagging:
        if not doc.was_preprocess_done(PreProcessSteps.sentencer):
            return
        if not self.override and doc.was_preprocess_done(PreProcessSteps.nerc):
            #print 'Already done'
            return

        entities = self.execute(doc)

        doc.set_preprocess_result(PreProcessSteps.nerc, entities)
        doc.save()

    def execute(self, doc):
        entities = []
        # Apply the ner algorithm which takes a list of sentences and returns
        # a list of sentences, each being a list of NER-tokens, each of which is
        # a pairs (tokenstring, class)
        ner_sentences = self.ner(doc.get_sentences())
        # Flatten the nested list above into just a list of kinds
        ner_kinds = (k for s in ner_sentences for (_, k) in s)
        
        # We build a large iterator z that goes over tuples like the following:
        #  (offset, (token, kind))
        # offset just goes incrementally from 0
        
        z = itertools.chain(
            enumerate(zip(doc.tokens, ner_kinds)),
            # Add a sentinel last token to simplify last iteration of loop below
            [(len(doc.tokens), (None, 'INVALID'))]
        )

        # Traverse z, looking for changes in the kind field. If there is a
        # change of kind, we have a new set of contiguous tokens; if the kind
        # of those isn't "O" (which means "other"), record the occurrence
        #
        # offset keeps the start of the current token run; last_kind keeps the kind.
        last_kind = 'O'
        offset = 0
        for i, (token, kind) in z:
            if kind != last_kind:
                if last_kind != 'O':
                    # Found a new entity in offset:i
                    name = ' '.join(doc.tokens[offset:i])
                    entities.append(EntityOccurrence.build(name, last_kind.lower(), name, offset, i))
                # Restart offset counter at each change of entity type
                offset = i
            last_kind = kind

        return entities


class StanfordNERRunner(NERRunner):
<commit_msg>Sanity check over tokenization length<commit_after>import itertools
import os
import os.path

from nltk.tag.stanford import NERTagger
import wget

from iepy.models import PreProcessSteps, Entity, EntityOccurrence
from iepy.preprocess import BasePreProcessStepRunner
from iepy.utils import DIRS, unzip_file

stanford_ner_name = 'stanford-ner-2014-01-04'
download_url_base = 'http://nlp.stanford.edu/software/'


class NonTokenizingNERTagger(NERTagger):

    @property
    def _cmd(self):
        old = super(NonTokenizingNERTagger, self)._cmd
        old = old + ["-tokenizerFactory", "edu.stanford.nlp.process.WhitespaceTokenizer"]
        return old


class NERRunner(BasePreProcessStepRunner):
    """Wrapper to insert a generic callable sentence NER tagger into the pipeline.
    """
    # TODO: rename to ner
    step = PreProcessSteps.nerc

    def __init__(self, ner, override=False):
        self.override = override
        self.ner = ner

    def __call__(self, doc):
        # this step does not necessarily requires PreProcessSteps.tagging:
        if not doc.was_preprocess_done(PreProcessSteps.sentencer):
            return
        if not self.override and doc.was_preprocess_done(PreProcessSteps.nerc):
            #print 'Already done'
            return

        entities = self.execute(doc)

        doc.set_preprocess_result(PreProcessSteps.nerc, entities)
        doc.save()

    def execute(self, doc):
        entities = []
        # Apply the ner algorithm which takes a list of sentences and returns
        # a list of sentences, each being a list of NER-tokens, each of which is
        # a pairs (tokenstring, class)
        ner_sentences = self.ner(doc.get_sentences())
        # Flatten the nested list above into just a list of kinds
        ner_kinds = (k for s in ner_sentences for (_, k) in s)
        
        # We build a large iterator z that goes over tuples like the following:
        #  (offset, (token, kind))
        # offset just goes incrementally from 0
        
        z = itertools.chain(
            enumerate(zip(doc.tokens, ner_kinds)),
            # Add a sentinel last token to simplify last iteration of loop below
            [(len(doc.tokens), (None, 'INVALID'))]
        )

        # Traverse z, looking for changes in the kind field. If there is a
        # change of kind, we have a new set of contiguous tokens; if the kind
        # of those isn't "O" (which means "other"), record the occurrence
        #
        # offset keeps the start of the current token run; last_kind keeps the kind.
        last_kind = 'O'
        offset = 0
        for i, (token, kind) in z:
            if kind != last_kind:
                if last_kind != 'O':
                    # Found a new entity in offset:i
                    name = ' '.join(doc.tokens[offset:i])
                    entities.append(EntityOccurrence.build(name, last_kind.lower(), name, offset, i))
                # Restart offset counter at each change of entity type
                offset = i
            last_kind = kind

        # Just a sanity check: verify that all NER tokens were consumed
        try:
            next(ner_kinds)
            assert False, "ner_kinds should have been completely consumed"
        except StopIteration:
            # Actually the stop iteration is the expected result here
            pass

        return entities


class StanfordNERRunner(NERRunner):

{'fraction_non_alphanumeric': 0.0767097966728281, 'fraction_numerical': 0.0030807147258163892, 'mean_word_length': 2.3314520266803487, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3159274', 'n_tokens_mistral': 1858, 'n_tokens_neox': 1793, 'n_words': 627}