import unittest import numpy as np from op_test import OpTest class Segment(object): def __init__(self, chunk_type, start_idx, end_idx): self.chunk_type = chunk_type self.start_idx = start_idx self.end_idx = end_idx def __str__(self): return '(Segment: %s, %s, %s)' % (self.chunk_type, self.start_idx, self.end_idx) __repr__ = __str__ class TestChunkEvalOp(OpTest): num_sequences = 5 batch_size = 50 def parse_scheme(self): if self.scheme == 'IOB': self.num_tag_types = 2 elif self.scheme == 'IOE': self.num_tag_types = 2 def fill_with_chunks(self, data, chunks): for chunk in chunks: if self.scheme == 'IOB': data[chunk.start_idx] = chunk.chunk_type * self.num_tag_types data[chunk.start_idx + 1: chunk.end_idx] = chunk.chunk_type * self.num_tag_types + ( self.num_tag_types - 1) data[chunk.end_idx] = chunk.chunk_type * self.num_tag_types + ( self.num_tag_types - 1 ) if chunk.start_idx < chunk.end_idx else data[chunk.start_idx] elif self.scheme == 'IOE': data[chunk.start_idx: chunk.end_idx] = chunk.chunk_type * self.num_tag_types data[chunk.end_idx] = chunk.chunk_type * self.num_tag_types + ( self.num_tag_types - 1) def rand_chunks(self, starts, num_chunks): if num_chunks < 0: num_chunks = np.random.randint(starts[-1]) chunks = [] # generate chunk beginnings chunk_begins = sorted( np.random.choice( range(starts[-1]), num_chunks, replace=False)) seq_chunk_begins = [] begin_idx = 0 # divide chunks into sequences for i in range(len(starts) - 1): tmp_chunk_begins = [] while begin_idx < len(chunk_begins) and chunk_begins[ begin_idx] < starts[i + 1]: tmp_chunk_begins.append(chunk_begins[begin_idx]) begin_idx += 1 seq_chunk_begins.append(tmp_chunk_begins) # generate chunk ends chunk_ends = [] for i in range(len(seq_chunk_begins)): for j in range(len(seq_chunk_begins[i])): low = seq_chunk_begins[i][j] high = seq_chunk_begins[i][j + 1] if j < len(seq_chunk_begins[ i]) - 1 else starts[i + 1] chunk_ends.append(np.random.randint(low, high)) # generate chunks for chunk_pos in zip(chunk_begins, chunk_ends): chunk_type = np.random.randint(self.num_chunk_types) chunks.append(Segment(chunk_type, *chunk_pos)) return chunks def gen_chunks(self, infer, label, starts): chunks = self.rand_chunks(starts, self.num_infer_chunks + self.num_label_chunks - self.num_correct_chunks) correct_chunks = np.random.choice( range(len(chunks)), self.num_correct_chunks, replace=False) infer_chunks = np.random.choice( [x for x in range(len(chunks)) if x not in correct_chunks], self.num_infer_chunks - self.num_correct_chunks, replace=False) infer_chunks = sorted(correct_chunks.tolist() + infer_chunks.tolist()) label_chunks = np.random.choice( [x for x in range(len(chunks)) if x not in infer_chunks], self.num_label_chunks - self.num_correct_chunks, replace=False) label_chunks = sorted(correct_chunks.tolist() + label_chunks.tolist()) self.fill_with_chunks(infer, [chunks[idx] for idx in infer_chunks]) self.fill_with_chunks(label, [chunks[idx] for idx in label_chunks]) # exclude types in excluded_chunk_types if len(self.excluded_chunk_types) > 0: for idx in correct_chunks: if chunks[idx].chunk_type in self.excluded_chunk_types: self.num_correct_chunks -= 1 for idx in infer_chunks: if chunks[idx].chunk_type in self.excluded_chunk_types: self.num_infer_chunks -= 1 for idx in label_chunks: if chunks[idx].chunk_type in self.excluded_chunk_types: self.num_label_chunks -= 1 return self.num_correct_chunks, self.num_infer_chunks, self.num_label_chunks def set_confs(self): # Use the IOB scheme and labels with 2 chunk types self.scheme = 'IOB' self.num_chunk_types = 2 self.excluded_chunk_types = [] self.other_chunk_type = self.num_chunk_types self.attrs = { 'num_chunk_types': self.num_chunk_types, 'chunk_scheme': self.scheme, 'excluded_chunk_types': self.excluded_chunk_types } self.parse_scheme() self.num_correct_chunks, self.num_infer_chunks, self.num_label_chunks = 4, 5, 9 def set_data(self): infer = np.zeros((self.batch_size, )).astype('int32') infer.fill(self.num_chunk_types * self.num_tag_types) label = np.copy(infer) starts = np.random.choice( range(1, self.batch_size), self.num_sequences - 1, replace=False).tolist() starts.extend([0, self.batch_size]) starts = sorted(starts) self.num_correct_chunks, self.num_infer_chunks, self.num_label_chunks = self.gen_chunks( infer, label, starts) self.inputs = { 'Inference': (infer, [starts]), 'Label': (label, [starts]) } precision = float( self.num_correct_chunks ) / self.num_infer_chunks if self.num_infer_chunks else 0 recall = float(self.num_correct_chunks ) / self.num_label_chunks if self.num_label_chunks else 0 f1 = float(2 * precision * recall) / ( precision + recall) if self.num_correct_chunks else 0 self.outputs = { 'Precision': np.asarray( [precision], dtype='float32'), 'Recall': np.asarray( [recall], dtype='float32'), 'F1-Score': np.asarray( [f1], dtype='float32') } def setUp(self): self.op_type = 'chunk_eval' self.set_confs() self.set_data() def test_check_output(self): self.check_output() class TestChunkEvalOpWithExclude(TestChunkEvalOp): def set_confs(self): # Use the IOE scheme and labels with 3 chunk types self.scheme = 'IOE' self.num_chunk_types = 3 self.excluded_chunk_types = [1] self.other_chunk_type = self.num_chunk_types self.attrs = { 'num_chunk_types': self.num_chunk_types, 'chunk_scheme': self.scheme, 'excluded_chunk_types': self.excluded_chunk_types } self.parse_scheme() self.num_correct_chunks, self.num_infer_chunks, self.num_label_chunks = 15, 18, 20 if __name__ == '__main__': unittest.main()