Splitting each word using longest substring that matches in a vocab/dictionary

Question

This function takes in a string and a fixed list of vocab:

• splits a string into words by spaces
• for each word, check in the dictionary/vocab and find the longest matching substrings, if none matching use the [UNK] word
• if a word is longer than a predefined length, it'll also become [UNK]

Code:

def tokenize(text, vocab, max_input_chars_per_word=10, unk_token="[UNK]"):
    output_tokens = []
    for token in text.split():
        chars = list(token)
        if len(chars) > max_input_chars_per_word:
            output_tokens.append(unk_token)
            continue

        is_bad = False
        start = 0
        sub_tokens = []

        while start < len(chars):
            end = len(chars)
            cur_substr = None
            while start < end:
                substr = "".join(chars[start:end])
                if start > 0:
                    substr = "##" + substr
                if substr in vocab:
                    cur_substr = substr
                    break
                end -= 1
            if cur_substr is None:
                is_bad = True
                break
            sub_tokens.append(cur_substr)
            start = end

        if is_bad:
            output_tokens.append(unk_token)
        else:
            output_tokens.extend(sub_tokens)
    return output_tokens

And example input and expected output:

vocab = ["the", "go", "##es", "to", "eat", "pum", "##pkins", "of", "##gos", "##stein"
         "#400", "1", "boredom", "##folk", "man", "##go", "out", "folks", "##0", 
        "un", "##aff", "##able"]


s = "the unaffable folks goes to eat 1400 folkspkinsgosgo pumpkins and 10 mangos out of boredom"

tokenize(s, vocab)

[out]:

['the',
 'un',
 '##aff',
 '##able',
 'folks',
 'go',
 '##es',
 'to',
 'eat',
 '[UNK]',
 '[UNK]',
 'pum',
 '##pkins',
 '[UNK]',
 '1',
 '##0',
 'man',
 '##gos',
 'out',
 'of',
 'boredom']

The nested while loop looks a complicated and checking each token isn't parallelized since it looks like it's independent of other tokens. How can this function be improved?

Simpler loops, or checking in parallel or maybe an easier to find longest matching substring from a list when iterating through the tokens? Or regexes?

Answer 1

The following is a basic syntax and usage pass without looking into your algorithm in too much depth. I'll first show the suggested code and then highlight significant differences where it can offer improvements on the original.

#!/usr/bin/env python3

from typing import Iterable, Set


def tokenize(
    text: str,
    vocab: Set[str],
    max_input_chars_per_word=10,
    unk_token="[UNK]"
) -> Iterable[str]:
    for token in text.split():
        n = len(token)
        if n > max_input_chars_per_word:
            yield unk_token
            continue

        start = 0
        sub_tokens = []

        while start < n:
            end = n
            cur_substr = None
            while start < end:
                substr = token[start:end]
                if start > 0:
                    substr = "##" + substr
                if substr in vocab:
                    cur_substr = substr
                    break
                end -= 1
            if cur_substr is None:
                yield unk_token
                break
            sub_tokens.append(cur_substr)
            start = end
        else:
            yield from sub_tokens


def test():
    vocab = {"the", "go", "##es", "to", "eat", "pum", "##pkins", "of", "##gos", "##stein"
             "#400", "1", "boredom", "##folk", "man",
             "##go", "out", "folks", "##0",
             "un", "##aff", "##able"}

    s = "the unaffable folks goes to eat 1400 folkspkinsgosgo pumpkins and 10 mangos out of boredom"

    result = tokenize(s, vocab)
    assert list(result) == [
        'the',
        'un',
        '##aff',
        '##able',
        'folks',
        'go',
        '##es',
        'to',
        'eat',
        '[UNK]',
        '[UNK]',
        'pum',
        '##pkins',
        '[UNK]',
        '1',
        '##0',
        'man',
        '##gos',
        'out',
        'of',
        'boredom']


test()

Rationale

• Add PEP484 type hints such as Iterable for the return of your function.
• Require that you accept a set for vocab. All you do on it is membership testing, i.e. if substr in vocab:, so a set will be much faster.
• Make this function a generator and yield instead of building up an output_tokens list.
• Store the token length in a variable such as n, considering you need it so often.
• Do not split out token into chars, and do not call join. All you need is the token string.
• Do not track an is_bad flag. Use a for/else to detect whether a break occurred.
• Use your test data in actual assert tests.

Answer 2

Here's my take on it. I split the vocab into 2 set. The first set, heads if for substrings that start at the beginning of the string, and the second set tails for other substrings.

The else clause on a loop gets executed when the loop terminates "normally", but is skipped when break is used. Some people don't like these else clauses, but I find them handy.

def tokenize(text, heads, tails, max_token_length=10, unknown_token="[UNK]"):
    output_tokens = []

    for token in text.split():

        if len(token) > max_token_length:
            output_tokens.append(unknown_token)
            continue

        sub_tokens = []

        substrs, flag = heads, ''
        while token:
            for end in range(len(token),0,-1):
                if token[:end] in substrs:
                    sub_tokens.append(f"{flag}{token[:end]}")
                    token = token[end:]
                    substrs, flag = tails, '##'
                    break

            else:
                output_tokens.append(unknown_token)
                break

        else:
            output_tokens.extend(sub_tokens)

    return output_tokens

Used like so:

heads = set(v for v in vocab if v[:2]!='##')
tails = set(v[2:] for v in vocab if v[:2]=='##')

tokenize(s, heads, tails)