tokenize.py

def find_spans(sample, common_words):
    
    spans = []
    for word in common_words:
        count = 0
        s = re.finditer(f'(^|\s){re.escape(word)}[^a-zA-Z]', sample, flags=re.IGNORECASE)
        for idx in s:
            start, end = idx.span()
            count += 1
            spans.append([start, end])
    spans = np.array(spans)
    spans = spans[spans[:,0].argsort()]

    new_span = []
    i = 0
    num_span = spans.shape[0]
    while i < num_span:
        if i < (num_span-1) and spans[i, 1] > spans[i+1, 0]:
            start_i = i
            i += 1
            while i < (num_span-1) and spans[i, 1] > spans[i+1, 0]:
                i += 1
            new_span.append([spans[start_i, 0], spans[i-1, 1]])
        else:
            new_span.append(spans[i, :].tolist())
            i += 1
    spans = np.array(new_span)
    
    return spans

def extract_spans(sample, spans):
    index  = 0
    chunks = []
    
    for span in spans:
        
        if index != span[0]:
            chunks += [(sample[index:span[0]], False)]
            
        chunks += [(sample[span[0]:span[1]], True)]
        index = span[1]
    
    if (index < (len(sample) - 1)):
        chunks += [(sample[index:], False)]
    
    return chunks

def tokenize(sample, tokenizer, common_words):
    
    spans  = find_spans(sample, common_words)
    chunks = extract_spans(sample, spans)
    
    prev_num_tokens = 0
    
    current_sample = ''
    current_tokens = []
    current_labels = []
    
    truncated_sample = ''
    
    for chunk_index, (chunk_str, chunk_label) in enumerate(chunks):
        
        current_sample    += chunk_str
        current_tokens     = tokenizer.encode(current_sample, add_special_tokens=False)
        num_current_tokens = len(current_tokens)
    
        if chunk_index < (len(chunks) - 2):
            next_sample     = current_sample + chunks[chunk_index + 1][0]
            next_tokens     = tokenizer.encode(next_sample, add_special_tokens=False)
            num_next_tokens = len(next_tokens)
            
            if current_tokens[num_current_tokens-1] != next_tokens[num_current_tokens-1]:
                continue
        
        new_tokens      = current_tokens[prev_num_tokens:]
        current_labels += [chunk_label] * len(new_tokens)
        prev_num_tokens = num_current_tokens
        
    assert len(current_tokens) == len(current_labels)    
    
    return current_tokens, current_labels 
      

usage.py

sample = 'The quick brown fox jumps over the lazy dog. Here is a dog mSv/a. test. ' * 2

tokens, mask = tokenize(sample, tokenizer, ['the', 'a'])    

print(len(tokens), tokens)
print(len(mask  ), mask)