Building a Japanese BPE Tokenizer: From Characters to Subwords

# Download and read Japanese news corpus
import tarfile
import urllib.request
import os

filename = "ldcc-20140209.tar.gz"
if not os.path.exists(filename):
    print("Downloading...")
    urllib.request.urlretrieve("https://www.rondhuit.com/download/ldcc-20140209.tar.gz", filename)

text = ""
articles = []
article_count_limit = 1000

with tarfile.open(filename, "r") as tar:
    for member in tar.getmembers():
        if member.isfile() and member.name.endswith(".txt"):
            file_content = tar.extractfile(member).read().decode("utf-8")
            lines = file_content.split("\n")
            if len(lines) >= 4:  # Skip URL, date, title
                content = "\n".join(lines[3:]).strip()
                if content:
                    articles.append(content)
                    text += content + "\n\n"

print(f"Total text length: {len(text):,}")
print(f"Total articles: {len(articles):,}")
print(f"Initial vocab size: {len(set(text)):,}")

def get_pairs(tokens):
    """Find all adjacent character pairs and their frequencies"""
    pairs = defaultdict(int)
    for i in range(len(tokens) - 1):
        pairs[(tokens[i], tokens[i + 1])] += 1
    return pairs

def merge_tokens(tokens, pair):
    """Merge the most frequent pair into a single token"""
    new_tokens = []
    i = 0
    while i < len(tokens):
        if (i < len(tokens) - 1 and 
            tokens[i] == pair[0] and tokens[i + 1] == pair[1]):
            new_tokens.append(pair[0] + pair[1])  # Merge pair
            i += 2
        else:
            new_tokens.append(tokens[i])
            i += 1
    return new_tokens

# Initialize with character-level tokens
tokens = list(text)  # Split into individual characters
vocab = set(tokens)
target_vocab_size = 5000

while len(vocab) < target_vocab_size:
    # Find most frequent adjacent pair
    pairs = get_pairs(tokens)
    if not pairs:
        break
    
    best_pair = max(pairs, key=pairs.get)
    
    # Merge all instances of this pair
    tokens = merge_tokens(tokens, best_pair)
    vocab.add("".join(best_pair))
    
    if merge_count % 100 == 0:
        print(f"Merge {merge_count}: vocab {len(vocab)}, tokens {len(tokens)}")

def tokenize_with_bpe(text):
    """Apply learned BPE merges to new text"""
    tokens = list(text)  # Start with characters
    changed = True
    
    while changed:
        changed = False
        new_tokens = []
        i = 0
        while i < len(tokens):
            # Check if adjacent pair exists in our vocabulary
            if (i < len(tokens) - 1 and 
                tokens[i] + tokens[i + 1] in vocab):
                new_tokens.append(tokens[i] + tokens[i + 1])
                i += 2
                changed = True
            else:
                new_tokens.append(tokens[i])
                i += 1
        tokens = new_tokens
    return tokens

# Training Results:
Total text length: 659,423
Total articles: 1,000
Initial vocab size: 2,913
Post-BPE vocab size: 5,000

# Compression Results:
Original characters: 659,423
BPE tokens: 379,352
Compression: 42.5%