bpe-tokenizer

bpe-tokenize.py

for s in data_trn[:nr_trn]:
    cur_tokens += list(s)

init_token_count = len(cur_tokens)
nr_actual_tokens = len(stoi)
nr_desired_tokens = 5000

print(f"Initial token count: {init_token_count}")

while nr_actual_tokens < nr_desired_tokens:
    new_tokens = []
    counts = {}
    candidates = {}

    for i in range(len(cur_tokens)):
        if i == len(cur_tokens) - 1:
            break

        left = cur_tokens[i]
        right = cur_tokens[i+1]

        # Prevent merging on the stop_char so it is easier to deliminate
        # each sample from the dataset
        if left == stop_char or right == stop_char:
            continue

        tok = left + right

        if tok not in counts:
            counts[tok] = 1
            candidates[tok] = {}
            candidates[tok]["left"] = left
            candidates[tok]["right"] = right
        else:
            counts[tok] += 1

    # Need the key which has max count
    new_token = max(counts, key=counts.get)
    left = candidates[new_token]['left']
    right = candidates[new_token]['right']
    cursor = 0

    for i in range(len(cur_tokens)):
        if i == 0:
            continue

        # Check for merge condition
        #
        # We merge if the left and right tokens match, and the cursor is not i.
        # If cursor is i, it means we just merged, and could merge again (two matches
        # in a row), but merges are non-overlapping, so we just skip over, keeping
        # cursor where it is.
        if cur_tokens[i-1] == left and cur_tokens[i] == right and cursor != i:
            if cursor < i - 1:
                # Cursor is behind the left token, so copy [cursor, left token)
                new_tokens += cur_tokens[cursor:i-1] + [new_token]
            else:
                new_tokens += [new_token]

            # anytime we merge, we move the cursor to the right of the right-merge token
            cursor = i + 1

    # Grab the end. this also cleanly covers the degenerate case where no merges happened
    if cursor <= len(cur_tokens) - 1:
        new_tokens += cur_tokens[cursor:]

    cur_tokens = new_tokens
    new_len = len(cur_tokens)

    # Now add the new token to the token dictionary
    stoi[new_token] = nr_actual_tokens
    itos[nr_actual_tokens] = new_token
    nr_actual_tokens += 1
    
    print(f"Merged {new_token}")
    
    if nr_actual_tokens % 1000 == 0:
        print(f"Tokens: {nr_actual_tokens} ({(init_token_count - new_len) / init_token_count:.2f}% reduction)")
        
    if len(new_token) > 40:
        print(f"Max token length is {len(new_token)}: {new_token}   (stopping)")
        print(f"Number of tokens: {nr_actual_tokens}")
        break