<commit_before>#!/usr/bin/env python
"""
Container for CoNLL03 NEL annotation
"""

class Mention(object):
    """Named entity mention."""
    def __init__(self, start, end, link=None):
        """
        start - begin token offset (slice semantics)
        end - end token offset (slice semantics)
        link - e.g., Wikipedia title (None == NIL)
        """
        self.start = start
        self.end = end
        self.link = link

    def __str__(self):
        return '<Mention [{}:{}] -> {}>'.format(self.start, self.end, self.link)

    def __cmp__(self, other):
        return cmp(self.start, other.start) or cmp(self.end, other.end)

    def span_match(self, other):
        return cmp(self, other) == 0

    def fuzzy_span_match(self, other):
        raise NotImplementedError

    def link_match(self, other):
        return self.link == other.link

class Document(object):
    def __init__(self, id, mentions=None):
        self.id, self.split = self._parse_id(id)
        self.mentions = mentions or []

    def _parse_id(self, id):
        split = 'train'
        if 'testa' in id:
            split = 'testa'
        elif 'testb' in id:
            split = 'testb'
        id = id.replace(split, '')
        return id, split

    @classmethod
    def from_string(cls, s):
        raise NotImplementedError
    
class Data(object):

    @classmethod
    def from_string(cls, s):
        raise NotImplementedError<commit_msg>first cut of data object with mention information for basic accuracy measures.<commit_after>#!/usr/bin/env python
"""
Container for CoNLL03 NEL annotation
"""
import gzip
from collections import OrderedDict

class Mention(object):
    """Named entity mention."""
    def __init__(self, start, end, link=None):
        """
        start - begin token offset (slice semantics)
        end - end token offset (slice semantics)
        link - e.g., Wikipedia url (None == NIL)
        """
        self.start = start
        self.end = end
        self.link = link

    @property
    def title(self):
        if self.link is None:
            return None
        else:
            return self.link[self.link.rfind('/')+1:].replace('_', ' ')

    def __str__(self):
        return '<Mention [{}:{}] -> {}>'.format(self.start, self.end, self.link)

    def __cmp__(self, other):
        return cmp(self.start, other.start) or cmp(self.end, other.end)

    def span_match(self, other):
        return cmp(self, other) == 0

    def fuzzy_span_match(self, other):
        raise NotImplementedError

    def link_match(self, other):
        return self.link == other.link

class Document(object):
    def __init__(self, id, mentions=None):
        self.id, self.split = self._parse_id(id)
        self.mentions = mentions or []

    def _parse_id(self, id):
        split = 'train'
        if 'testa' in id:
            split = 'testa'
        elif 'testb' in id:
            split = 'testb'
        id = id.replace(split, '')
        return id, split

    @classmethod
    def from_lines(cls, lines):
        """
        Return document object.
        s - CoNLL-formatted lines
        """
        doc = cls._init_document(lines)
        if doc is not None:
            for m in cls._iter_mentions(lines):
                doc.mentions.append(m)
        return doc

    @classmethod
    def _init_document(cls, lines):
        for line in lines:
            if line.startswith('-DOCSTART-'):
                id = line.strip()[12:-1].split()[0]
                return cls(id)

    @classmethod
    def _iter_mentions(cls, lines):
        tok_id = -1
        start = None
        for line in lines:
            if line.strip() == '' or line.startswith('-DOCSTART-'):
                continue # blank line
            tok_id += 1
            tok, bi, name, link = cls._parse_line(line)
            if bi is None:
                continue # not an entity mention
            if bi == 'B':
                start = tok_id
            if tok == name.split()[-1]:
                yield Mention(start, tok_id, link)
                start = None

    @classmethod
    def _parse_line(cls, line):
        cols = line.rstrip('\n').split('\t')
        tok, bi, name, link = None, None, None, None
        if len(cols) >= 1:
            tok = cols[0] # current token
        if len(cols) >= 3:
            bi = cols[1] # begin/inside tag
            name = cols[2] # full mention text
        if len(cols) >= 5:
            link = cols[4] # Wikipedia url
        return tok, bi, name, link
    
class Data(object):
    def __init__(self, documents=None):
        self.documents = documents or OrderedDict()

    @classmethod
    def from_file(cls, f):
        """
        Return data object.
        f - CoNLL-formatted file
        """
        fh = gzip.open(f) if f.endswith('.gz') else open(f)
        data = cls()
        for d in cls._iter_documents(fh.readlines()):
            data.documents[d.id] = d
        return data

    @classmethod
    def _iter_documents(cls, lines):
        doc = []
        for i, line in enumerate(lines):
            if line.startswith('-DOCSTART-'):
                if len(doc) > 0:
                    yield Document.from_lines(doc)
                doc = []
            if line.strip() != '':
                doc.append(line)
        yield Document.from_lines(doc)
        
        
{'fraction_non_alphanumeric': 0.11987263532496722, 'fraction_numerical': 0.004120621839295748, 'mean_word_length': 1.8374070138150904, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 11, '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': '4405844', 'n_tokens_mistral': 1654, 'n_tokens_neox': 1583, 'n_words': 450}