サブワード分割手法の比較 — Unigram・BPE・WordPieceの学習アルゴリズムと性能を徹底比較

import numpy as np
import matplotlib.pyplot as plt
from collections import Counter, defaultdict

# 比較用の共通コーパス
corpus = [
    "the cat sat on the mat",
    "the cat ate the fish",
    "the dog sat on the mat",
    "the dog chased the cat",
    "a new cat sat on a new mat",
    "the newest cat chased the oldest dog",
    "lower the newest rating",
    "the lower rating is newer",
    "playing with the newest toy",
    "the newest player is playing",
    "unplayable games are lower rated",
    "the newest unplayable game",
]

# 単語頻度の集計（末尾マーカー付き）
def get_word_freqs(corpus):
    word_freqs = Counter()
    for sentence in corpus:
        for word in sentence.strip().split():
            word_freqs[word] += 1
    return word_freqs

word_freqs = get_word_freqs(corpus)
print("=== コーパスの単語頻度（上位15件） ===")
for word, freq in word_freqs.most_common(15):
    print(f"  {word}: {freq}")
print(f"  異なり語数: {len(word_freqs)}")
print(f"  総単語数: {sum(word_freqs.values())}")

class BPETokenizer:
    """BPEトークナイザ"""

    def __init__(self, num_merges=30):
        self.num_merges = num_merges
        self.merges = []
        self.vocab = set()
        self.merge_history = []  # マージ履歴を記録

    def _get_split_words(self, corpus):
        """コーパスを文字に分解（語末マーカー付き）"""
        word_freqs = Counter()
        for sentence in corpus:
            for word in sentence.strip().split():
                word_freqs[word] += 1
        split_words = {}
        for word, freq in word_freqs.items():
            chars = tuple(list(word) + ["</w>"])
            split_words[chars] = freq
        return split_words

    def _get_pair_freqs(self, split_words):
        """隣接ペアの頻度を計算"""
        pair_freqs = Counter()
        for chars, freq in split_words.items():
            for i in range(len(chars) - 1):
                pair_freqs[(chars[i], chars[i+1])] += freq
        return pair_freqs

    def _merge_pair(self, pair, split_words):
        """指定ペアをマージ"""
        new_split = {}
        merged = pair[0] + pair[1]
        for chars, freq in split_words.items():
            new_chars = []
            i = 0
            while i < len(chars):
                if (i < len(chars) - 1
                    and chars[i] == pair[0]
                    and chars[i+1] == pair[1]):
                    new_chars.append(merged)
                    i += 2
                else:
                    new_chars.append(chars[i])
                    i += 1
            new_split[tuple(new_chars)] = freq
        return new_split

    def train(self, corpus):
        """BPEを学習"""
        split_words = self._get_split_words(corpus)
        self.vocab = set()
        for chars in split_words:
            for c in chars:
                self.vocab.add(c)

        for step in range(self.num_merges):
            pair_freqs = self._get_pair_freqs(split_words)
            if not pair_freqs:
                break
            best_pair = max(pair_freqs, key=pair_freqs.get)
            best_freq = pair_freqs[best_pair]
            split_words = self._merge_pair(best_pair, split_words)
            merged = best_pair[0] + best_pair[1]
            self.merges.append(best_pair)
            self.vocab.add(merged)
            self.merge_history.append({
                'step': step + 1,
                'pair': best_pair,
                'merged': merged,
                'freq': best_freq
            })

    def encode(self, word):
        """単語をBPEでエンコード"""
        tokens = list(word) + ["</w>"]
        for pair in self.merges:
            i = 0
            while i < len(tokens) - 1:
                if tokens[i] == pair[0] and tokens[i+1] == pair[1]:
                    tokens = tokens[:i] + [pair[0]+pair[1]] + tokens[i+2:]
                else:
                    i += 1
        return tokens

class WordPieceTokenizer:
    """WordPieceトークナイザ（PMIベース）"""

    def __init__(self, num_merges=30):
        self.num_merges = num_merges
        self.vocab = set()
        self.merges = []
        self.merge_history = []

    def _get_split_words(self, corpus):
        """コーパスを文字に分解"""
        word_freqs = Counter()
        for sentence in corpus:
            for word in sentence.strip().split():
                word_freqs[word] += 1
        split_words = {}
        for word, freq in word_freqs.items():
            chars = tuple(list(word) + ["</w>"])
            split_words[chars] = freq
        return split_words

    def _get_pair_scores(self, split_words):
        """PMIベースのスコアを計算"""
        pair_freqs = Counter()
        token_freqs = Counter()

        for chars, freq in split_words.items():
            for i in range(len(chars)):
                token_freqs[chars[i]] += freq
                if i < len(chars) - 1:
                    pair_freqs[(chars[i], chars[i+1])] += freq

        # PMIスコア: count(ab) / (count(a) * count(b))
        pair_scores = {}
        for pair, freq in pair_freqs.items():
            a, b = pair
            score = freq / (token_freqs[a] * token_freqs[b])
            pair_scores[pair] = score
        return pair_scores

    def _merge_pair(self, pair, split_words):
        """指定ペアをマージ"""
        new_split = {}
        merged = pair[0] + pair[1]
        for chars, freq in split_words.items():
            new_chars = []
            i = 0
            while i < len(chars):
                if (i < len(chars) - 1
                    and chars[i] == pair[0]
                    and chars[i+1] == pair[1]):
                    new_chars.append(merged)
                    i += 2
                else:
                    new_chars.append(chars[i])
                    i += 1
            new_split[tuple(new_chars)] = freq
        return new_split

    def train(self, corpus):
        """WordPieceを学習"""
        split_words = self._get_split_words(corpus)
        self.vocab = set()
        for chars in split_words:
            for c in chars:
                self.vocab.add(c)

        for step in range(self.num_merges):
            pair_scores = self._get_pair_scores(split_words)
            if not pair_scores:
                break
            best_pair = max(pair_scores, key=pair_scores.get)
            best_score = pair_scores[best_pair]
            split_words = self._merge_pair(best_pair, split_words)
            merged = best_pair[0] + best_pair[1]
            self.merges.append(best_pair)
            self.vocab.add(merged)
            self.merge_history.append({
                'step': step + 1,
                'pair': best_pair,
                'merged': merged,
                'score': best_score
            })

    def encode(self, word):
        """最長一致法でエンコード"""
        tokens = []
        chars = word + "</w>"
        start = 0
        while start < len(chars):
            end = len(chars)
            found = False
            while start < end:
                substr = chars[start:end]
                if substr in self.vocab:
                    tokens.append(substr)
                    found = True
                    break
                end -= 1
            if not found:
                tokens.append(chars[start])
                start += 1
            else:
                start = end
        return tokens

class UnigramTokenizer:
    """Unigramトークナイザ（EM + Viterbi）"""

    def __init__(self, target_vocab_size=50, max_piece_length=8):
        self.target_vocab_size = target_vocab_size
        self.max_piece_length = max_piece_length
        self.vocab = {}  # token -> log_prob
        self.char_set = set()

    def _initialize_vocab(self, corpus):
        """初期語彙を構築（全サブストリング）"""
        substring_freqs = Counter()
        char_set = set()

        for sentence in corpus:
            for word in sentence.strip().split():
                word_with_end = word + "</w>"
                for i in range(len(word_with_end)):
                    for j in range(i+1,
                                   min(i+self.max_piece_length+1,
                                       len(word_with_end)+1)):
                        substring_freqs[word_with_end[i:j]] += 1
                for c in word_with_end:
                    char_set.add(c)

        total = sum(substring_freqs.values())
        self.vocab = {}
        for substr, freq in substring_freqs.most_common(
                self.target_vocab_size * 5):
            self.vocab[substr] = np.log(freq / total)

        for c in char_set:
            if c not in self.vocab:
                self.vocab[c] = np.log(1e-10)
        self.char_set = char_set

    def _viterbi_decode(self, text):
        """Viterbiアルゴリズムで最適分割"""
        n = len(text)
        dp = [-float('inf')] * (n + 1)
        dp[0] = 0.0
        backpointer = [0] * (n + 1)

        for j in range(1, n + 1):
            for i in range(max(0, j - self.max_piece_length), j):
                piece = text[i:j]
                if piece in self.vocab:
                    score = dp[i] + self.vocab[piece]
                    if score > dp[j]:
                        dp[j] = score
                        backpointer[j] = i

        # バックトレース
        tokens = []
        j = n
        while j > 0:
            i = backpointer[j]
            tokens.append(text[i:j])
            j = i
        tokens.reverse()
        return tokens, dp[n]

    def train(self, corpus, n_iterations=10):
        """Unigramモデルを学習"""
        self._initialize_vocab(corpus)

        word_freqs = Counter()
        for sentence in corpus:
            for word in sentence.strip().split():
                word_freqs[word + "</w>"] += 1

        self.log_likelihoods = []

        for iteration in range(n_iterations):
            # Eステップ: Viterbi近似で期待出現回数を計算
            token_counts = Counter()
            total_log_prob = 0.0

            for word, freq in word_freqs.items():
                tokens, log_prob = self._viterbi_decode(word)
                total_log_prob += log_prob * freq
                for token in tokens:
                    token_counts[token] += freq

            self.log_likelihoods.append(total_log_prob)

            # Mステップ: 確率を更新
            total_count = sum(token_counts.values())
            for token in self.vocab:
                if token in token_counts:
                    self.vocab[token] = np.log(
                        token_counts[token] / total_count)
                else:
                    self.vocab[token] = np.log(1e-10)

            # 語彙の削減
            if len(self.vocab) > self.target_vocab_size:
                scored = []
                for token in self.vocab:
                    if len(token) == 1 or token in self.char_set:
                        continue
                    count = token_counts.get(token, 0)
                    scored.append((token, count))

                scored.sort(key=lambda x: x[1])
                n_remove = max(1, int(len(self.vocab) * 0.1))
                n_remove = min(n_remove,
                               len(self.vocab) - self.target_vocab_size)

                for token, _ in scored[:n_remove]:
                    if len(self.vocab) <= self.target_vocab_size:
                        break
                    del self.vocab[token]

    def encode(self, word):
        """Viterbi分割でエンコード"""
        tokens, _ = self._viterbi_decode(word + "</w>")
        return tokens

# 3手法を同じコーパスで学習
bpe = BPETokenizer(num_merges=25)
bpe.train(corpus)

wp = WordPieceTokenizer(num_merges=25)
wp.train(corpus)

unigram = UnigramTokenizer(target_vocab_size=40, max_piece_length=8)
unigram.train(corpus, n_iterations=15)

# テスト単語でエンコード結果を比較
test_words = ["newest", "playing", "unplayable", "lower",
              "chased", "rating", "cat", "the"]

print("=" * 70)
print(f"{'単語':<14} {'BPE':<22} {'WordPiece':<22} {'Unigram':<22}")
print("=" * 70)
for word in test_words:
    bpe_tokens = bpe.encode(word)
    wp_tokens = wp.encode(word)
    uni_tokens = unigram.encode(word)
    # 表示用に </w> を除去して見やすくする
    bpe_str = " + ".join(t.replace("</w>", "_") for t in bpe_tokens)
    wp_str = " + ".join(t.replace("</w>", "_") for t in wp_tokens)
    uni_str = " + ".join(t.replace("</w>", "_") for t in uni_tokens)
    print(f"{word:<14} {bpe_str:<22} {wp_str:<22} {uni_str:<22}")

fig, axes = plt.subplots(1, 2, figsize=(14, 6))

# BPEのマージ順序
ax = axes[0]
bpe_steps = [h['step'] for h in bpe.merge_history[:15]]
bpe_labels = [f"{h['pair'][0]}+{h['pair'][1]}" for h in bpe.merge_history[:15]]
bpe_freqs = [h['freq'] for h in bpe.merge_history[:15]]
bars = ax.barh(range(len(bpe_steps)), bpe_freqs, color='#4FC3F7', edgecolor='white')
ax.set_yticks(range(len(bpe_steps)))
ax.set_yticklabels(bpe_labels, fontsize=9)
ax.set_xlabel('Pair Frequency', fontsize=11)
ax.set_title('BPE: Merge Order (by Frequency)', fontsize=13)
ax.invert_yaxis()
ax.grid(axis='x', alpha=0.3)

# WordPieceのマージ順序
ax = axes[1]
wp_steps = [h['step'] for h in wp.merge_history[:15]]
wp_labels = [f"{h['pair'][0]}+{h['pair'][1]}" for h in wp.merge_history[:15]]
wp_scores = [h['score'] for h in wp.merge_history[:15]]
bars = ax.barh(range(len(wp_steps)), wp_scores, color='#FF8A65', edgecolor='white')
ax.set_yticks(range(len(wp_steps)))
ax.set_yticklabels(wp_labels, fontsize=9)
ax.set_xlabel('PMI Score', fontsize=11)
ax.set_title('WordPiece: Merge Order (by PMI)', fontsize=13)
ax.invert_yaxis()
ax.grid(axis='x', alpha=0.3)

plt.tight_layout()
plt.show()

# 各手法でコーパス全体のトークン数を計算
def count_total_tokens(tokenizer, corpus):
    total = 0
    for sentence in corpus:
        for word in sentence.strip().split():
            tokens = tokenizer.encode(word)
            total += len(tokens)
    return total

bpe_total = count_total_tokens(bpe, corpus)
wp_total = count_total_tokens(wp, corpus)
uni_total = count_total_tokens(unigram, corpus)

# 元の文字数
total_chars = sum(len(word) for sentence in corpus
                  for word in sentence.strip().split())

methods = ['BPE', 'WordPiece', 'Unigram']
token_counts = [bpe_total, wp_total, uni_total]
vocab_sizes = [len(bpe.vocab), len(wp.vocab), len(unigram.vocab)]
compression = [total_chars / tc for tc in token_counts]

fig, axes = plt.subplots(1, 3, figsize=(15, 5))

# トークン数の比較
colors = ['#4FC3F7', '#FF8A65', '#81C784']
axes[0].bar(methods, token_counts, color=colors, edgecolor='white', width=0.6)
axes[0].set_ylabel('Total Token Count', fontsize=11)
axes[0].set_title('Total Tokens (Entire Corpus)', fontsize=13)
axes[0].grid(axis='y', alpha=0.3)
for i, v in enumerate(token_counts):
    axes[0].text(i, v + 1, str(v), ha='center', fontsize=11, fontweight='bold')

# 語彙サイズの比較
axes[1].bar(methods, vocab_sizes, color=colors, edgecolor='white', width=0.6)
axes[1].set_ylabel('Vocabulary Size', fontsize=11)
axes[1].set_title('Final Vocabulary Size', fontsize=13)
axes[1].grid(axis='y', alpha=0.3)
for i, v in enumerate(vocab_sizes):
    axes[1].text(i, v + 0.5, str(v), ha='center', fontsize=11, fontweight='bold')

# 圧縮率の比較
axes[2].bar(methods, compression, color=colors, edgecolor='white', width=0.6)
axes[2].set_ylabel('Compression Ratio (chars/token)', fontsize=11)
axes[2].set_title('Compression Ratio', fontsize=13)
axes[2].grid(axis='y', alpha=0.3)
for i, v in enumerate(compression):
    axes[2].text(i, v + 0.05, f'{v:.2f}', ha='center', fontsize=11,
                 fontweight='bold')

plt.tight_layout()
plt.show()

# 語彙サイズを変えてBPEのトークン数を測定
merge_counts = [5, 10, 15, 20, 25, 30, 35, 40]
bpe_token_series = []
wp_token_series = []

for n_merge in merge_counts:
    # BPE
    bpe_tmp = BPETokenizer(num_merges=n_merge)
    bpe_tmp.train(corpus)
    bpe_token_series.append(count_total_tokens(bpe_tmp, corpus))

    # WordPiece
    wp_tmp = WordPieceTokenizer(num_merges=n_merge)
    wp_tmp.train(corpus)
    wp_token_series.append(count_total_tokens(wp_tmp, corpus))

# Unigramは語彙サイズで制御
uni_vocab_sizes = [20, 25, 30, 35, 40, 50, 60, 80]
uni_token_series = []

for vs in uni_vocab_sizes:
    uni_tmp = UnigramTokenizer(target_vocab_size=vs, max_piece_length=8)
    uni_tmp.train(corpus, n_iterations=15)
    uni_token_series.append(count_total_tokens(uni_tmp, corpus))

fig, ax = plt.subplots(figsize=(10, 6))

# BPEとWordPieceは初期語彙 + マージ数 = 語彙サイズ
# 簡易的にマージ数を横軸に
ax.plot(merge_counts, bpe_token_series, 'o-',
        color='#4FC3F7', linewidth=2, markersize=8, label='BPE')
ax.plot(merge_counts, wp_token_series, 's-',
        color='#FF8A65', linewidth=2, markersize=8, label='WordPiece')
ax.plot(uni_vocab_sizes, uni_token_series, '^-',
        color='#81C784', linewidth=2, markersize=8, label='Unigram')

ax.set_xlabel('Vocabulary Size Parameter\n'
              '(num_merges for BPE/WP, target_vocab_size for Unigram)',
              fontsize=11)
ax.set_ylabel('Total Token Count', fontsize=11)
ax.set_title('Vocabulary Size vs Token Count', fontsize=14)
ax.legend(fontsize=12)
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

# 対数尤度の推移を可視化
fig, ax = plt.subplots(figsize=(8, 5))

iterations = list(range(1, len(unigram.log_likelihoods) + 1))
ax.plot(iterations, unigram.log_likelihoods, 'o-',
        color='#81C784', linewidth=2, markersize=8)
ax.set_xlabel('EM Iteration', fontsize=12)
ax.set_ylabel('Log-Likelihood', fontsize=12)
ax.set_title('Unigram Model: EM Convergence', fontsize=14)
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

# 分割結果の詳細比較
detailed_words = ["newest", "playing", "unplayable",
                  "lower", "chased", "rated"]

fig, axes = plt.subplots(len(detailed_words), 1,
                          figsize=(12, len(detailed_words) * 1.2))

colors_map = {'BPE': '#4FC3F7', 'WordPiece': '#FF8A65',
              'Unigram': '#81C784'}

for idx, word in enumerate(detailed_words):
    ax = axes[idx]

    bpe_t = bpe.encode(word)
    wp_t = wp.encode(word)
    uni_t = unigram.encode(word)

    results = {
        'BPE': [t.replace('</w>', '') for t in bpe_t if t != '</w>'],
        'WordPiece': [t.replace('</w>', '') for t in wp_t if t != '</w>'],
        'Unigram': [t.replace('</w>', '') for t in uni_t if t != '</w>'],
    }

    y_positions = {'BPE': 0.7, 'WordPiece': 0.4, 'Unigram': 0.1}

    for method_name, tokens in results.items():
        y = y_positions[method_name]
        x = 0.15
        ax.text(0.0, y, f'{method_name}:', fontsize=9,
                va='center', fontweight='bold',
                color=colors_map[method_name],
                transform=ax.transAxes)
        for token in tokens:
            width = len(token) * 0.04 + 0.02
            rect = plt.Rectangle((x, y - 0.12), width, 0.24,
                                  facecolor=colors_map[method_name],
                                  alpha=0.3, edgecolor=colors_map[method_name],
                                  transform=ax.transAxes, clip_on=False)
            ax.add_patch(rect)
            ax.text(x + width/2, y, token, fontsize=9, ha='center',
                    va='center', transform=ax.transAxes,
                    fontweight='bold')
            x += width + 0.01

    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_title(f'"{word}"', fontsize=11, fontweight='bold', loc='left')
    ax.axis('off')

plt.suptitle('Subword Segmentation Comparison', fontsize=14,
             fontweight='bold', y=1.02)
plt.tight_layout()
plt.show()