本篇内容主要介绍如何根据提供的文本内容创建 WordPiece vocabulary，代码来自谷歌；
一、简单介绍
wordpiece的目的是：通过考虑单词内部构造，充分利用subwords的优势，在把长word转化为短word提高文字的灵活性以及提高word转化的效率这两处之间取得一个良好的平衡；
前者会增加词表大小，后者会减少词表大小
二、步骤流程
2.1 预处理
在读取所有的文本内容后，第一步便是对文本内容预处理；
- 对英文来说，我们可以把字符都转化为小写形式，去掉accents，á 变成 a，然后利用
whitespace和punctuation进行分割；也就是空格和标点符号； - 对中文来说，我们可以把繁体转化为简体，分割的方式就只有单个单个字词进行分割了，而优化的方式只有从外部引入
tokenizer对文本内容做分词，然后进行后续步骤，不然单个中文字词无法进行分解，有人可能想通过偏旁部首来，但偏旁部首如何区分顺序呢？后续的内容将围绕英文展开；
在将文本切割成以word为单位的小块后，我们进行下一步骤；
2.2 计数
在预处理这一阶段，我们得到了以word为单位的小块，为了了解总体words的情况，我们需要对word进行计数处理，并按照数量从大到小排列；如果文本内容很大，我们可以在此做一个优化，过滤掉数量太大或者太小的word以及过滤掉长度太长的word；
由于word piece的本质是subwords，为了合理的把word转化为subwords，我们必须考虑word的基本单元；因为如果词表中缺少组成word的基本单元，那么该词的表示就无法实现或者不完整和其他词照成混淆；
所以在这里我们统计所有的词其单个字符出现的次数；同理，在这里我们可以优化一下，删除出现次数较少的字符，由于出现次数较少的字符删除了，哪包含这些字符的words也就无法表示，所以我们同时要删除包含这些字符的words；
2.3 分割
在这一环节，我们对计数字典的word进行分割，其处理方式如下：
首先对word设置一个首指针和一个尾指针，以指针之间的内容求匹配计数字典和字符字典的合集，若成功，则将首指针指向尾指针，然后尾指针重新指向最后的位置，若失败，则将尾指针向首指针移动一步；直到停止首尾位置一致，若首尾在尾部则返回output_tokens，若在其他地方则说明不能分词，返回None
实现过程如图所示：
实现代码如下：
def get_split_indices(word, curr_tokens, include_joiner_token, joiner): indices = [] start = 0 while start < len(word): end = len(word) while end > start: subtoken = word[start:end] # Subtoken includes the joiner token. if include_joiner_token and start > 0: subtoken = joiner + subtoken # If subtoken is part of vocab, 'end' is a valid start index. if subtoken in curr_tokens: indices.append(end) break end -= 1 if end == start: return 1 start = end return indices if __name__ == '__main__': res = get_split_indices('hello', ['h', '##e', '##llo', '##o'], True, '##') # print(res) res: [1, 2, 5]
2.4 添加subword
上一步分割的作用实际上是在找最大的分词块，但是其采用的是一种贪婪算法，并不是最优解；在对word进行分割找最大的分割块的indice之后，我们可以更快的找到常常出现在一起的字符串；处理方式如下：获取每一个以indice位置开始，长度依次增加的subword，构建subword字典并计数，其每次增加的数目应该是word.count，
这种遍历方式产生的subword的数目过于庞大，因此如果有需要，我们需要对其进行一些优化，比如删掉一些长度较长的subword，删除一些次数比较小的subword，这样添加subword的这一步骤就算完成了；
但是要注意的是，这里的subword出现了重复计数，我们考虑了长的字符串，那么短的字符串一定会被考虑，这里我们从长字符串到短字符开始遍历，当确定长字符串有一定数目确定为vocabulary中的元素时，我们把所有有相同前缀的短字符串减去长字符串的数目避免影响；
与此同时，vocabulary并不一定包含了字符字典，所以我们需要将其进行合并，最后得到的vocabulary就是wordpiece vocabulary
三、代码实现
首先我们对word进行预处理，这里代码省略；
在这里我们传入一个iterable迭代器，然后用collections库中的Counter，对每个词进行计数；
def count_words(iterable) -> collections.Counter: """Converts a iterable of arrays of words into a `Counter` of word counts.""" counts = collections.Counter() for words in iterable: # Convert a RaggedTensor to a flat/dense Tensor. words = getattr(words, 'flat_values', words) # Flatten any dense tensor words = np.reshape(words, [-1]) counts.update(words) # Decode the words if necessary. example_word = next(iter(counts.keys())) if isinstance(example_word, bytes): counts = collections.Counter( {word.decode('utf-8'): count for word, count in counts.items()}) return counts
根据当前词频以及upper_thresh和lower_thresh确定词频的界限；
def get_search_threshs(word_counts, upper_thresh, lower_thresh): """Clips the thresholds for binary search based on current word counts. The upper threshold parameter typically has a large default value that can result in many iterations of unnecessary search. Thus we clip the upper and lower bounds of search to the maximum and the minimum wordcount values. Args: word_counts: list of (string, int) tuples upper_thresh: int, upper threshold for binary search lower_thresh: int, lower threshold for binary search Returns: upper_search: int, clipped upper threshold for binary search lower_search: int, clipped lower threshold for binary search """ counts = [count for _, count in word_counts] max_count = max(counts) min_count = min(counts) if upper_thresh is None: upper_search = max_count else: upper_search = max_count if max_count < upper_thresh else upper_thresh if lower_thresh is None: lower_search = min_count else: lower_search = min_count if min_count > lower_thresh else lower_thresh return upper_search, lower_search
对单个的char的数量设置一个上限；
def get_allowed_chars(all_counts, max_unique_chars): """Get the top max_unique_chars characters within our wordcounts. We want each character to be in the vocabulary so that we can keep splitting down to the character level if necessary. However, in order not to inflate our vocabulary with rare characters, we only keep the top max_unique_chars characters. Args: all_counts: list of (string, int) tuples max_unique_chars: int, maximum number of unique single-character tokens Returns: set of strings containing top max_unique_chars characters in all_counts """ char_counts = collections.defaultdict(int) for word, count in all_counts: for char in word: char_counts[char] += count # Sort by count, then alphabetically. sorted_counts = sorted(sorted(char_counts.items(), key=lambda x: x[0]), key=lambda x: x[1], reverse=True) allowed_chars = set() for i in range(min(len(sorted_counts), max_unique_chars)): allowed_chars.add(sorted_counts[i][0]) return allowed_chars
结合all_counts和allowed_chars，删掉包含allowed_char的字符，控制结果为max_input_tokens个出现次数最大的word
def filter_input_words(all_counts, allowed_chars, max_input_tokens): """Filters out words with unallowed chars and limits words to max_input_tokens. Args: all_counts: list of (string, int) tuples allowed_chars: list of single-character strings max_input_tokens: int, maximum number of tokens accepted as input Returns: list of (string, int) tuples of filtered wordcounts """ # Ensure that the input is sorted so that if `max_input_tokens` is reached # the least common tokens are dropped. all_counts = sorted( all_counts, key=lambda word_and_count: word_and_count[1], reverse=True) filtered_counts = [] for word, count in all_counts: if (max_input_tokens != -1 and len(filtered_counts) >= max_input_tokens): break has_unallowed_chars = False for char in word: if char not in allowed_chars: has_unallowed_chars = True break if has_unallowed_chars: continue filtered_counts.append((word, count)) return filtered_counts
获得split的index，要保证curr_tokens可以对word进行分割；
def get_split_indices(word, curr_tokens, include_joiner_token, joiner): """Gets indices for valid substrings of word, for iterations > 0. For iterations > 0, rather than considering every possible substring, we only want to consider starting points corresponding to the start of wordpieces in the current vocabulary. Args: word: string we want to split into substrings curr_tokens: string to int dict of tokens in vocab (from previous iteration) include_joiner_token: bool whether to include joiner token joiner: string used to indicate suffixes Returns: list of ints containing valid starting indices for word """ indices = [] start = 0 while start < len(word): end = len(word) while end > start: subtoken = word[start:end] # Subtoken includes the joiner token. if include_joiner_token and start > 0: subtoken = joiner + subtoken # If subtoken is part of vocab, 'end' is a valid start index. if subtoken in curr_tokens: indices.append(end) break end -= 1 if end == start: return None start = end return indices
进行最后的步骤；
import collections from typing import List, Optional Params = collections.namedtuple('Params', [ 'upper_thresh', 'lower_thresh', 'num_iterations', 'max_input_tokens', 'max_token_length', 'max_unique_chars', 'vocab_size', 'slack_ratio', 'include_joiner_token', 'joiner', 'reserved_tokens' ]) def extract_char_tokens(word_counts): """Extracts all single-character tokens from word_counts. Args: word_counts: list of (string, int) tuples Returns: set of single-character strings contained within word_counts """ seen_chars = set() for word, _ in word_counts: for char in word: seen_chars.add(char) return seen_chars def ensure_all_tokens_exist(input_tokens, output_tokens, include_joiner_token, joiner): """Adds all tokens in input_tokens to output_tokens if not already present. Args: input_tokens: set of strings (tokens) we want to include output_tokens: string to int dictionary mapping token to count include_joiner_token: bool whether to include joiner token joiner: string used to indicate suffixes Returns: string to int dictionary with all tokens in input_tokens included """ for token in input_tokens: if token not in output_tokens: output_tokens[token] = 1 if include_joiner_token: joined_token = joiner + token if joined_token not in output_tokens: output_tokens[joined_token] = 1 return output_tokens def get_search_threshs(word_counts, upper_thresh, lower_thresh): """Clips the thresholds for binary search based on current word counts. The upper threshold parameter typically has a large default value that can result in many iterations of unnecessary search. Thus we clip the upper and lower bounds of search to the maximum and the minimum wordcount values. Args: word_counts: list of (string, int) tuples upper_thresh: int, upper threshold for binary search lower_thresh: int, lower threshold for binary search Returns: upper_search: int, clipped upper threshold for binary search lower_search: int, clipped lower threshold for binary search """ counts = [count for _, count in word_counts] max_count = max(counts) min_count = min(counts) if upper_thresh is None: upper_search = max_count else: upper_search = max_count if max_count < upper_thresh else upper_thresh if lower_thresh is None: lower_search = min_count else: lower_search = min_count if min_count > lower_thresh else lower_thresh return upper_search, lower_search def get_input_words(word_counts, reserved_tokens, max_token_length): """Filters out words that are longer than max_token_length or are reserved. Args: word_counts: list of (string, int) tuples reserved_tokens: list of strings max_token_length: int, maximum length of a token Returns: list of (string, int) tuples of filtered wordcounts """ all_counts = [] for word, count in word_counts: if len(word) > max_token_length or word in reserved_tokens: continue all_counts.append((word, count)) return all_counts def generate_final_vocabulary(reserved_tokens, char_tokens, curr_tokens): """Generates final vocab given reserved, single-character, and current tokens. Args: reserved_tokens: list of strings (tokens) that must be included in vocab char_tokens: set of single-character strings curr_tokens: string to int dict mapping token to count Returns: list of strings representing final vocabulary """ sorted_char_tokens = sorted(list(char_tokens)) vocab_char_arrays = [] vocab_char_arrays.extend(reserved_tokens) vocab_char_arrays.extend(sorted_char_tokens) # Sort by count, then alphabetically. sorted_tokens = sorted(sorted(curr_tokens.items(), key=lambda x: x[0]), key=lambda x: x[1], reverse=True) for token, _ in sorted_tokens: vocab_char_arrays.append(token) seen_tokens = set() # Adding unique tokens to list to maintain sorted order. vocab_words = [] for word in vocab_char_arrays: if word in seen_tokens: continue seen_tokens.add(word) vocab_words.append(word) return vocab_words def learn_with_thresh(word_counts, thresh, params): """Wordpiece learning algorithm to produce a vocab given frequency threshold. Args: word_counts: list of (string, int) tuples thresh: int, frequency threshold for a token to be included in the vocab params: Params namedtuple, parameters for learning Returns: list of strings, vocabulary generated for the given thresh """ # Set of single-character tokens. char_tokens = extract_char_tokens(word_counts) curr_tokens = ensure_all_tokens_exist(char_tokens, {}, params.include_joiner_token, params.joiner) for iteration in range(params.num_iterations): subtokens = [dict() for _ in range(params.max_token_length + 1)] # Populate array with counts of each subtoken. for word, count in word_counts: if iteration == 0: split_indices = range(1, len(word) + 1) else: split_indices = get_split_indices(word, curr_tokens, params.include_joiner_token, params.joiner) if not split_indices: continue start = 0 for index in split_indices: for end in range(start + 1, len(word) + 1): subtoken = word[start:end] length = len(subtoken) if params.include_joiner_token and start > 0: subtoken = params.joiner + subtoken if subtoken in subtokens[length]: # Subtoken exists, increment count. subtokens[length][subtoken] += count else: # New subtoken, add to dict. subtokens[length][subtoken] = count start = index next_tokens = {} # Get all tokens that have a count above the threshold. for length in range(params.max_token_length, 0, -1): for token, count in subtokens[length].items(): if count >= thresh: next_tokens[token] = count # Decrement the count of all prefixes. if len(token) > length: # This token includes the joiner. joiner_len = len(params.joiner) for i in range(1 + joiner_len, length + joiner_len): prefix = token[0:i] if prefix in subtokens[i - joiner_len]: subtokens[i - joiner_len][prefix] -= count else: for i in range(1, length): prefix = token[0:i] if prefix in subtokens[i]: subtokens[i][prefix] -= count # Add back single-character tokens. curr_tokens = ensure_all_tokens_exist(char_tokens, next_tokens, params.include_joiner_token, params.joiner) vocab_words = generate_final_vocabulary(params.reserved_tokens, char_tokens, curr_tokens) return vocab_words def learn_binary_search(word_counts, lower, upper, params): """Performs binary search to find wordcount frequency threshold. Given upper and lower bounds and a list of (word, count) tuples, performs binary search to find the threshold closest to producing a vocabulary of size vocab_size. Args: word_counts: list of (string, int) tuples lower: int, lower bound for binary search upper: int, upper bound for binary search params: Params namedtuple, parameters for learning Returns: list of strings, vocab that is closest to target vocab_size """ thresh = (upper + lower) // 2 current_vocab = learn_with_thresh(word_counts, thresh, params) current_vocab_size = len(current_vocab) # Allow count to be within k% of the target count, where k is slack ratio. slack_count = params.slack_ratio * params.vocab_size if slack_count < 0: slack_count = 0 is_within_slack = (current_vocab_size <= params.vocab_size) and ( params.vocab_size - current_vocab_size <= slack_count) # We've created a vocab within our goal range (or, ran out of search space). if is_within_slack or lower >= upper or thresh <= 1: return current_vocab current_vocab = None if current_vocab_size > params.vocab_size: return learn_binary_search(word_counts, thresh + 1, upper, params) else: return learn_binary_search(word_counts, lower, thresh - 1, params)
整合：
def learn(word_counts, vocab_size: int, reserved_tokens: List[str], upper_thresh: Optional[int] = int(1e7), lower_thresh: Optional[int] = 10, num_iterations: int = 4, max_input_tokens: Optional[int] = int(5e6), max_token_length: int = 50, max_unique_chars: int = 1000, slack_ratio: float = 0.05, include_joiner_token: bool = True, joiner: str = '##') -> List[str]: """Takes in wordcounts and returns wordpiece vocabulary. Args: word_counts: (word, count) pairs as a dictionary, or list of tuples. vocab_size: The target vocabulary size. This is the maximum size. reserved_tokens: A list of tokens that must be included in the vocabulary. upper_thresh: Initial upper bound on the token frequency threshold. lower_thresh: Initial lower bound on the token frequency threchold. num_iterations: Number of iterations to run. max_input_tokens: The maximum number of words in the initial vocabulary. The words with the lowest counts are discarded. Use `None` or `-1` for "no maximum". max_token_length: The maximum token length. Counts for longer words are discarded. max_unique_chars: The maximum alphabet size. This prevents rare characters from inflating the vocabulary. Counts for words containing characters ouside of the selected alphabet are discarded. slack_ratio: The maximum deviation acceptable from `vocab_size` for an acceptable vocabulary. The acceptable range of vocabulary sizes is from `vocab_size*(1-slack_ratio)` to `vocab_size`. include_joiner_token: If true, include the `joiner` token in the output vocabulary. joiner: The prefix to include on suffix tokens in the output vocabulary. Usually "##". For example 'places' could be tokenized as `['place', '##s']`. Returns: string, final vocabulary with each word separated by newline """ if isinstance(word_counts, dict): word_counts = word_counts.items() params = Params(upper_thresh, lower_thresh, num_iterations, max_input_tokens, max_token_length, max_unique_chars, vocab_size, slack_ratio, include_joiner_token, joiner, reserved_tokens) upper_search, lower_search = get_search_threshs(word_counts, params.upper_thresh, params.lower_thresh) all_counts = get_input_words(word_counts, params.reserved_tokens, params.max_token_length) allowed_chars = get_allowed_chars(all_counts, params.max_unique_chars) filtered_counts = filter_input_words(all_counts, allowed_chars, params.max_input_tokens) vocab = learn_binary_search(filtered_counts, lower_search, upper_search, params) return vocab
