<commit_before>import re

from collections import defaultdict
from .prepare import PrepareCorpus
from ..doc import Doc, Candidate, Chain, Mention
from ..model import model
from ..process import tag
from ..process.tokenise import RegexTokeniser

import logging
log = logging.getLogger()

DOCSTART_MARKER = '-DOCSTART-'

@PrepareCorpus.Register
class ConllPrepare(object):
    """ Tokenise and tag a conll document """
    def __init__(self, conll_data_path, doc_type):
        self.in_path = conll_data_path
        self.doc_predicate = {
            'all': lambda _: True,
            'train': self.is_training_doc,
            'dev': self.is_dev_doc,
            'test': self.is_test_doc
        }[doc_type]

    @staticmethod
    def is_training_doc(doc_id):
        return 'test' not in doc_id
    @staticmethod
    def is_test_doc(doc_id):
        return 'testb' in doc_id
    @staticmethod
    def is_dev_doc(doc_id):
        return 'testa' in doc_id

    def __call__(self):
        """ Prepare documents """
        # todo: parameterise model tags
        redirects = model.Redirects('wikipedia')
        candidates = model.Candidates('wikipedia')

        total_mentions = 0
        total_recalled_mentions = 0
        
        log.info('Preparing conll documents...')
        for doc, conll_tags in self.iter_docs(self.in_path, self.doc_predicate):
            mentions = []
            for gold_entity, (start, end) in conll_tags:
                m = Mention(start, doc.text[start:end])
                m.resolution = redirects.map(gold_entity)
                m.surface_form = doc.text[start:end].replace(' ,', ',')
                mentions.append(m)

            chains = tag.Tagger.cluster_mentions(mentions)
            unique_chains = []

            # join chains with the same gold standard resolution, split chains with different gold resolution
            # this requires gold annot information so is only valid if building the train set or isolating disambiguation performance
            for chain in chains:
                mbr = defaultdict(list)
                for m in chain.mentions:
                    mbr[m.resolution].append(m)
                for r, ms in mbr.iteritems():
                    unique_chains.append(Chain(mentions=ms, resolution=Candidate(r)))

            doc.chains = unique_chains
 
            for chain in doc.chains:
                total_mentions += len(chain.mentions)
                # longest mention string
                sf = sorted(chain.mentions, key=lambda m: len(m.surface_form), reverse=True)[0].surface_form

                cs = candidates.search(sf)
                if cs:
                    if chain.resolution.id not in cs:
                        log.warn('Entity (%s) not in candidate set for (%s) in doc (%s).', chain.resolution.id, sf, doc.id)
                    else:
                        total_recalled_mentions += len(chain.mentions)
                    chain.candidates = [Candidate(e) for e in cs]
                else:
                    log.warn('Missing alias (%s): %s' % (doc.id, sf))
            yield doc
        log.info("Candidate Recall = %.1f%", float(total_recalled_mentions*100)/total_mentions)

    @staticmethod
    def iter_docs(path, doc_id_predicate, redirect_model = None, max_docs = None):
        """ Read AIDA-CoNLL formatted documents """
        redirect_model = redirect_model or {}

        with open(path, 'rd') as f:
            doc_id = None
            doc_tokens = None
            doc_mentions = None
            doc_count = 0
            for line in f:<commit_msg>Add raw candidate recall counts to output after conll prepare<commit_after>import re

from collections import defaultdict
from .prepare import PrepareCorpus
from ..doc import Doc, Candidate, Chain, Mention
from ..model import model
from ..process import tag
from ..process.tokenise import RegexTokeniser

import logging
log = logging.getLogger()

DOCSTART_MARKER = '-DOCSTART-'

@PrepareCorpus.Register
class ConllPrepare(object):
    """ Tokenise and tag a conll document """
    def __init__(self, conll_data_path, doc_type):
        self.in_path = conll_data_path
        self.doc_predicate = {
            'all': lambda _: True,
            'train': self.is_training_doc,
            'dev': self.is_dev_doc,
            'test': self.is_test_doc
        }[doc_type]

    @staticmethod
    def is_training_doc(doc_id):
        return 'test' not in doc_id
    @staticmethod
    def is_test_doc(doc_id):
        return 'testb' in doc_id
    @staticmethod
    def is_dev_doc(doc_id):
        return 'testa' in doc_id

    def __call__(self):
        """ Prepare documents """
        # todo: parameterise model tags
        redirects = model.Redirects('wikipedia')
        candidates = model.Candidates('wikipedia')

        total_mentions = 0
        total_recalled_mentions = 0
        
        log.info('Preparing conll documents...')
        for doc, conll_tags in self.iter_docs(self.in_path, self.doc_predicate):
            mentions = []
            for gold_entity, (start, end) in conll_tags:
                m = Mention(start, doc.text[start:end])
                m.resolution = redirects.map(gold_entity)
                m.surface_form = doc.text[start:end].replace(' ,', ',')
                mentions.append(m)

            chains = tag.Tagger.cluster_mentions(mentions)
            unique_chains = []

            # join chains with the same gold standard resolution, split chains with different gold resolution
            # this requires gold annot information so is only valid if building the train set or isolating disambiguation performance
            for chain in chains:
                mbr = defaultdict(list)
                for m in chain.mentions:
                    mbr[m.resolution].append(m)
                for r, ms in mbr.iteritems():
                    unique_chains.append(Chain(mentions=ms, resolution=Candidate(r)))

            doc.chains = unique_chains
 
            for chain in doc.chains:
                total_mentions += len(chain.mentions)
                # longest mention string
                sf = sorted(chain.mentions, key=lambda m: len(m.surface_form), reverse=True)[0].surface_form

                cs = candidates.search(sf)
                if cs:
                    if chain.resolution.id not in cs:
                        log.warn('Entity (%s) not in candidate set for (%s) in doc (%s).', chain.resolution.id, sf, doc.id)
                    else:
                        total_recalled_mentions += len(chain.mentions)
                    chain.candidates = [Candidate(e) for e in cs]
                else:
                    log.warn('Missing alias (%s): %s' % (doc.id, sf))
            yield doc

        log.info("Candidate Recall = %.1f%% (%i/%i)", float(total_recalled_mentions*100)/total_mentions, total_recalled_mentions, total_mentions)

    @staticmethod
    def iter_docs(path, doc_id_predicate, redirect_model = None, max_docs = None):
        """ Read AIDA-CoNLL formatted documents """
        redirect_model = redirect_model or {}

        with open(path, 'rd') as f:
            doc_id = None
            doc_tokens = None
            doc_mentions = None
            doc_count = 0
            for line in f:
{'fraction_non_alphanumeric': 0.09224838221120749, 'fraction_numerical': 0.0022029464408646563, 'mean_word_length': 2.017864561695056, '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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4512852', 'n_tokens_mistral': 2105, 'n_tokens_neox': 1991, 'n_words': 551}