简化基于固定ngram的字符串压缩循环（关联BPE算法）

给定字符列表（例如list('Hello▁world▁')）和字符元组列表：

[('l', 'l'), ('ell', 'o▁'), ('Hell', 'o▁'), ('w', 'or'), ('o', 'r'), ('e', 'l'), ('el', 'l'), ('H', 'ell'), ('H', 'e'), ('He', 'll'), ('worl', 'd▁'), ('wor', 'l'), ('l', 'd▁'), ('d', '▁'), ('wor', 'ld▁'), ('H', 'el'), ('o', '▁'), ('w', 'o'), ('l', 'o▁'), ('l', 'o')]

目标是遍历这些元组，若字符列表中存在连续匹配的元组对，则将其合并。现有实现代码可正常运行，但嵌套while循环可进一步简化，希望用正则迭代替换或更优的算法方案（该问题与Byte Pair Encoding（BPE）相关）优化实现。

import copy

def matcher(s, ngram):
  while s:
    window = tuple(s[:2]) # 因为元组是成对的
    if window == ngram:
      yield "".join(window)
      s = s[2:]
    else:
      yield s[0]
      s = s[1:]

def combine_ngrams(s, vocab):
  prev = copy.copy(s)
  while True:
    for v in vocab:
      s = list(matcher(s, v))
    if s == prev:
      break
    else:
      prev = s
  return s


vocab = [('l', 'l'), ('ell', 'o▁'), ('Hell', 'o▁'), ('w', 'or'), ('o', 'r'), ('e', 'l'), ('el', 'l'), ('H', 'ell'), ('H', 'e'), 
         ('He', 'll'), ('worl', 'd▁'), ('wor', 'l'), ('l', 'd▁'), ('d', '▁'), ('wor', 'ld▁'), ('H', 'el'), ('o', '▁'), ('w', 'o'), ('l', 'o▁'), ('l', 'o')]

s = list('Hello▁world▁')

print(combine_ngrams(s, vocab))

输出结果：

['Hello▁', 'world▁']

输入 #

# 示例1
s = list('abcde'); vocab = [('a', 'b'), ('b', 'c'), ('a', 'bc'), ('abc', 'd'), ('abcd', 'e')]

# 示例2
s = list('abcde'); vocab = [('a', 'b'), ('ab', 'c'), ('b', 'c'),  ('a', 'bc'), ('abc', 'd'), ('abcd', 'e')]

# 示例3
s = list('aaab'); vocab = [('a', 'a'), ('a', 'aa'), ('aaa', 'b')]

# 示例4
s = list('Hello▁ポケモンセンター▁world▁'); vocab = [('l', 'l'), ('ell', 'o▁'), ('Hell', 'o▁'), ('w', 'or'), ('o', 'r'), ('e', 'l'), ('el', 'l'), ('H', 'ell'), ('H', 'e'), ('He', 'll'), ('worl', 'd▁'), ('wor', 'l'), ('l', 'd▁'), ('d', '▁'), ('wor', 'ld▁'), ('H', 'el'), ('o', '▁'), ('w', 'o'), ('l', 'o▁'), ('l', 'o')]

输出 #

['ab', 'c', 'd', 'e']

['abcde']

['aa', 'a', 'b']

['Hello▁', 'ポ', 'ケ', 'モ', 'ン', 'セ', 'ン', 'タ', 'ー', '▁', 'world▁']

方案1：正则迭代替换 #

利用正则表达式批量替换连续匹配的组合，核心思路是用特殊分隔符将字符列表拼接成字符串，将vocab中的元组转换为正则匹配模式，迭代替换直到无变化，最后分割回列表。这种方式可避免嵌套while循环，代码更简洁：

import re

def combine_ngrams_regex(s, vocab):
    # 用不会出现在目标字符中的分隔符连接列表
    sep = '\x00'
    s_str = sep.join(s)
    # 按组合总长度倒序排序，优先匹配长组合，避免短组合提前替换导致长组合无法匹配
    sorted_vocab = sorted(vocab, key=lambda x: len(x[0]) + len(x[1]), reverse=True)
    # 生成正则替换规则
    patterns = []
    for a, b in sorted_vocab:
        pattern = re.compile(re.escape(a) + re.escape(sep) + re.escape(b))
        replacement = a + b
        patterns.append((pattern, replacement))
    
    while True:
        prev_str = s_str
        for pattern, repl in patterns:
            s_str = pattern.sub(repl, s_str)
        if s_str == prev_str:
            break
    # 分割回列表
    return s_str.split(sep)

方案2：BPE风格的优化遍历 #

参考BPE算法的核心逻辑，先对vocab按组合长度倒序排序，单次遍历完成一轮替换，直到没有变化，减少无效循环次数：

def combine_ngrams_bpe(s, vocab):
    # 按组合总长度倒序排序，优先处理长组合
    sorted_vocab = sorted(vocab, key=lambda x: len(x[0]) + len(x[1]), reverse=True)
    while True:
        changed = False
        new_s = []
        i = 0
        n = len(s)
        while i < n:
            matched = False
            # 遍历排序后的vocab，尝试匹配当前位置的组合
            for a, b in sorted_vocab:
                len_a = len(a)
                len_b = len(b)
                if i + len_a + len_b <= n and ''.join(s[i:i+len_a]) == a and ''.join(s[i+len_a:i+len_a+len_b]) == b:
                    new_s.append(a + b)
                    i += len_a + len_b
                    matched = True
                    changed = True
                    break
            if not matched:
                new_s.append(s[i])
                i += 1
        if not changed:
            break
        s = new_s
    return s

测试验证 #

用示例输入测试两种优化方案，结果与原实现一致：

# 测试示例1
s1 = list('abcde')
vocab1 = [('a', 'b'), ('b', 'c'), ('a', 'bc'), ('abc', 'd'), ('abcd', 'e')]
print(combine_ngrams_regex(s1, vocab1))  # ['ab', 'c', 'd', 'e']
print(combine_ngrams_bpe(s1, vocab1))    # ['ab', 'c', 'd', 'e']

# 测试示例2
s2 = list('abcde')
vocab2 = [('a', 'b'), ('ab', 'c'), ('b', 'c'),  ('a', 'bc'), ('abc', 'd'), ('abcd', 'e')]
print(combine_ngrams_regex(s2, vocab2))  # ['abcde']
print(combine_ngrams_bpe(s2, vocab2))    # ['abcde']

内容的提问来源于stack exchange，提问作者alvas