conll2003.py

conll2003.py

import os
import collections

path = os.path.expanduser('~/data/conll2003/eng.train')

# Read raw data.
with open(path) as f:
    dataset = []
    example = dict()
    for line in f:
        line = line.strip()
        if not line or len(line) == 0:
            # TODO: Add example.
            if len(example) > 0:
                dataset.append(example)
            example = dict()
            continue
        token, b, c, bio_entity = line.split()
        example.setdefault('bio', []).append(bio_entity)
        example.setdefault('token', []).append(token)

    if len(example) > 0:
        dataset.append(example)

del example

# Resolve BIO tags.
all_spans = []
for ex in dataset:
    span = []
    recent_entity_type = None
    for tok, bio in zip(ex['token'], ex['bio']):
        if bio.startswith('B'):
            # DONE
            if recent_entity_type is None:
                if len(span) > 0:
                    all_spans.append(span)
                    span = []

            entity_type = bio.split('-')[1]

            recent_entity_type = entity_type

            span.append(tok)
        elif bio.startswith('I'):
            entity_type = bio.split('-')[1]

            # DONE
            if recent_entity_type is None or recent_entity_type != entity_type:
                if len(span) > 0:
                    all_spans.append(span)
                    span = []

            recent_entity_type = entity_type

            span.append(tok)
        if bio.startswith('O'):
            # DONE
            if len(span) > 0:
                all_spans.append(span)
                span = []
            recent_entity_type = None


# Count lengths.
for span in all_spans:
    print(span)

lengths = collections.Counter([len(sp) for sp in all_spans])
print(lengths)