"""Classes and functions for working with entity labels."""
class Label(object):
"""Base class for a label.
A label contains four attributes of primary interest:
* entity_type - the type of entity which is labeled
* start - the start offset of the label in the source text
* length - the length of the label in the source text
* entity - the actual text of the entity
"""
def __init__(self, entity_type, start, length, entity):
"""Initialize a data processor with the location of the data."""
self.entity_type = entity_type
self.start = start
self.length = length
self.entity = entity
def __repr__(self):
return f'Label({self.entity_type}, {self.start}, {self.length}, {self.entity})'
def map_entity_type(self, mapping, force_upper=True):
if force_upper:
self.entity_type = mapping[self.entity_type.upper()]
else:
self.entity_type = mapping[self.entity_type]
def shift(self, s):
self.start += s
return self
def contains(self, i):
"""Returns true if any label contains the offset."""
return (self.start >= i) & ((self.start + self.length) < i)
def overlaps(self, start, stop):
"""Returns true if any label contains the start/stop offset."""
contains_start = (self.start >= start) & (self.start < stop)
contains_stop = ((self.start + self.length) >=
start) & ((self.start + self.length) < stop)
return contains_start | contains_stop
def within(self, start, stop):
"""Returns true if the label is within a start/stop offset."""
after_start = (self.start >= start) or ((self.start + self.length) >= start)
before_stop = self.start < stop
return after_start & before_stop
def convert_to_bio_scheme(tokens: list) -> list:
def b_or_i(w, w_prev):
if w == 'O':
return 'O'
elif w == w_prev:
return f'I-{w}'
else:
return f'B-{w}'
return [
[b_or_i(w, None if i == 0 else sequence[i-1]) for i, w in enumerate(sequence)]
for sequence in tokens
]