<commit_before>import json
import urllib
import codecs


class LitTagger:

    def __init__(self, label, src_filename):
        self.label = label
        self.src_filename = src_filename
        #names = set()
        #for filename in src_filenames:
        #    f = open(filename)
        #    namelist = f.read().strip().split('\n')
        #    names.update(namelist)
        #self.names = frozenset(names)
        f = open(src_filename)
        namelist = f.read().strip().split('\n')
        self.names = frozenset(namelist)
        
        # compute prefix closure
        prefixes = set()
        for name in self.names:
            sname = name.split()
            prefixes.update([' '.join(sname[:i]) for i in range(1, len(sname) + 1)])
        
        self.prefixes = frozenset(prefixes)
    
    def tag(self, sent):
        result = []
        i = 0
        while i < len(sent):
            j = i + 1
            prev_segment = segment = ' '.join(sent[i:j])
            #print 'check: ', segment
            while segment in self.prefixes and j <= len(sent):
                j += 1
                prev_segment = segment
                segment = ' '.join(sent[i:j])
            if prev_segment in self.names:
                # label sent[i:j - 1]
                for k in range(i, j - 1):
                    result.append((sent[k], self.label))
                i = j - 1
            else:
                result.append((sent[i], 'O'))
                i += 1
        
        return result


def download_freebase_type(type_name, dest_filename, normalizer=None, aliases=False):
    if not normalizer:
        normalizer = lambda x: x
    
    # https://developers.google.com/freebase/v1/mql-overview
    service_url = 'https://www.googleapis.com/freebase/v1/mqlread'
    #query = [{'type': '/medicine/disease', 'name': None}]
    query = [{'type': type_name, 'name': None}]
    if aliases:
        query[0]['/common/topic/alias'] = []
    
    f = codecs.open(dest_filename, 'w', encoding='utf8')

    params = { 'query': json.dumps(query) }
    url = service_url + '?' + urllib.urlencode(params) + '&cursor'
    response = json.loads(urllib.urlopen(url).read())
    cursor = response['cursor']
    
    # write results
    for result in response['result']:
        name = normalizer(result['name'])
        f.write(name + '\n')
        if aliases:
            for name in result['/common/topic/alias']:
                name = normalizer(name)
                f.write(name + '\n')    
    
    while cursor:
        params = { 'query': json.dumps(query), 'cursor': cursor }
        url = service_url + '?' + urllib.urlencode(params)
        response = json.loads(urllib.urlopen(url).read())
        cursor = response['cursor']
        
        # write results
        for result in response['result']:
            if result['name']:
                name = normalizer(result['name'])
                f.write(name + '\n')
            if aliases:
                for name in result['/common/topic/alias']:
                    name = normalizer(name)
                    f.write(name + '\n')<commit_msg>new entities() method to return indexes of entities. tagger runner based on the tag method, and alternative tagger runner based on the entities method.<commit_after>import json
import urllib
import codecs

from iepy.models import PreProcessSteps, Entity, EntityOccurrence
from iepy.ner import NERRunner
from iepy.preprocess import BasePreProcessStepRunner


class LitTagger:

    def __init__(self, label, src_filename):
        self.label = label
        self.src_filename = src_filename
        #names = set()
        #for filename in src_filenames:
        #    f = open(filename)
        #    namelist = f.read().strip().split('\n')
        #    names.update(namelist)
        #self.names = frozenset(names)
        f = open(src_filename)
        namelist = f.read().strip().split('\n')
        self.names = frozenset(namelist)
        
        # compute prefix closure
        prefixes = set()
        for name in self.names:
            sname = name.split()
            prefixes.update([' '.join(sname[:i]) for i in range(1, len(sname) + 1)])
        
        self.prefixes = frozenset(prefixes)
    
    def tag(self, sent):
        """Tagger with output a la Stanford (no start/end markers).
        """
        entities = self.entities(sent)
        # dummy entity for nicer code:
        entities.append((len(sent), len(sent)))
        next_entity = entities.pop(0)
        result = []
        for i, t in enumerate(sent):
            #print i, next_entity
            if i >= next_entity[1]:
                # assert entities
                next_entity = entities.pop(0)
            
            if i < next_entity[0]:
                result.append((t, 'O'))
            elif i < next_entity[1]:
                result.append((t, self.label))
        
        return result

    def entities(self, sent):
        """Return entities as a list of pairs (offset, offset_end).
        """
        result = []
        i = 0
        while i < len(sent):
            j = i + 1
            prev_segment = segment = ' '.join(sent[i:j])
            #print 'check: ', segment
            while segment in self.prefixes and j <= len(sent):
                j += 1
                prev_segment = segment
                segment = ' '.join(sent[i:j])
            if prev_segment in self.names:
                # label sent[i:j - 1]
                #for k in range(i, j - 1):
                #    result.append((sent[k], self.label))
                result.append((i, j - 1))
                i = j - 1
            else:
                #result.append((sent[i], 'O'))
                i += 1
        
        return result


class LitTaggerRunner(NERRunner):

    def __init__(self, label, src_filename, override=False):
        lit_tagger = LitTagger()
        callable_lit_tagger = lambda x: lit_tagger.tag(x)
        NERRunner.__init__(self, callable_lit_tagger, override)


class LitTaggerRunner2(BasePreProcessStepRunner):

    def __init__(self, label, src_filename, override=False):
        self.label = label
        self.lit_tagger = LitTagger(label, src_filename)
    
    def __call__(self, doc):
        # this step does not requires PreProcessSteps.tagging:
        if not doc.was_preprocess_done(PreProcessSteps.sentencer):
            return
        
        entities = []
        sent_offset = 0
        for sent in doc.get_sentences():
            sent_entities = self.lit_tagger.entities(sent)
            
            for (i, j) in sent_entities:
                name = ' '.join(sent[i:j])
                kind = self.label.lower() # XXX: should be in models.ENTITY_KINDS
                entity, created = Entity.objects.get_or_create(key=name, 
                            defaults={'canonical_form': name, 'kind': kind})
                entity_oc = EntityOccurrence(entity=entity, 
                                        offset=sent_offset + i, 
                                        offset_end=sent_offset + j)
                entities.append(entity_oc)
            
            sent_offset += len(sent)


def download_freebase_type(type_name, dest_filename, normalizer=None, aliases=False):
    if not normalizer:
        normalizer = lambda x: x
    
    # https://developers.google.com/freebase/v1/mql-overview
    service_url = 'https://www.googleapis.com/freebase/v1/mqlread'
    #query = [{'type': '/medicine/disease', 'name': None}]
    query = [{'type': type_name, 'name': None}]
    if aliases:
        query[0]['/common/topic/alias'] = []
    
    f = codecs.open(dest_filename, 'w', encoding='utf8')

    params = { 'query': json.dumps(query) }
    url = service_url + '?' + urllib.urlencode(params) + '&cursor'
    response = json.loads(urllib.urlopen(url).read())
    cursor = response['cursor']
    
    # write results
    for result in response['result']:
        name = normalizer(result['name'])
        f.write(name + '\n')
        if aliases:
            for name in result['/common/topic/alias']:
                name = normalizer(name)
                f.write(name + '\n')
    
    while cursor:
        params = { 'query': json.dumps(query), 'cursor': cursor }
        url = service_url + '?' + urllib.urlencode(params)
        response = json.loads(urllib.urlopen(url).read())
        cursor = response['cursor']
        
        # write results
        for result in response['result']:
            if result['name']:
                name = normalizer(result['name'])
                f.write(name + '\n')
            if aliases:
                for name in result['/common/topic/alias']:
                    name = normalizer(name)
                    f.write(name + '\n')
{'fraction_non_alphanumeric': 0.12335356774710801, 'fraction_numerical': 0.0038941701981445426, 'mean_word_length': 1.7747060692723229, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 4, 'https://': 4, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '464635', 'n_tokens_mistral': 2615, 'n_tokens_neox': 2534, 'n_words': 601}