关于LLaMA Tokenizer的一些坑...

使用LLaMA Tokenizer对 jsonl 文件进行分词，并将分词结果保存到 txt 文件中，分词代码如下：

import jsonlines
import sentencepiece as spm
from tqdm import tqdm

jsonl_file = '/path/to/jsonl_file'
txt_file = '/path/to/txt_file'

tokenizer = spm.SentencePieceProcessor('./tokenizer.model')
w = open(txt_file, mode='w', encoding='utf-8')

with jsonlines.open(jsonl_file, mode='r') as r:
    for line in tqdm(r):
        ids = tokenizer.Encode(line['text'])
        tokenized_text = ' '.join(tokenizer.IdToPiece(i) for i in ids)
        w.write(f"{tokenized_text}\n")

w.close()

从以上代码可以看出，txt 文件中的每行内容实际上是 jsonl 文件对应行的文档的分词结果，分词之间以空格分隔。理论上，这意味着 txt 文件的行数应与 jsonl 文件的行数相匹配，均等同于文档数量。然而，实际情况可能出现 txt 文件的行数显著超过 jsonl 文件的行数。

既然行数对不上，那自然就需要找到第一个出现问题的行。为此，我们重新执行分词过程，并在每行的开头添加该行的索引（从 1 1 1 开始）以便追踪。

with jsonlines.open(jsonl_file, mode='r') as r:
    for idx, line in tqdm(enumerate(r, start=1)):
        ids = tokenizer.Encode(line['text'])
        tokenized_text = ' '.join(tokenizer.IdToPiece(i) for i in ids)
        w.write(f"{idx} {tokenized_text}\n")

然后遍历 txt 文件，找到第一个出现问题的行：

with open(txt_file, mode='r') as r:
    for cnt, line in tqdm(enumerate(r, start=1)):
        idx = line[:line.index(' ')]
        if str(cnt) != idx:
            print(cnt)
            break

打开 txt 文件，发现出现问题的行的开头压根没有数字，这说明分词结果被某种特殊字符换行了。要注意LLaMA词表是没有换行符 \n 的（被映射到 <unk> 了），那还能是什么字符导致换行的呢？

将 jsonl 文件中对应的行抽取出来，单独对它进行分词，然后将分词结果打印到终端上，发现并没有换行，然而将这个分词结果单独写入到一个新的文件中时，换行出现了。通过对每一个token单独进行分析，发现其中有一个token含有 \r，而这个字符正是导致换行的罪魁祸首！

? \r 字符是回车符（Carriage Return, CR），在ASCII表中的位置是十进制的13或者十六进制的0x0D。这个字符最初设计用于打字机和早期计算机，指的是将打印头移动到一行的开始位置而不换到下一行，这样新的文本就会覆盖掉同一行上旧的文本。

接下来，我们就可以找出LLaMA词表中所有可能会导致换行的token了：

tokenizer = spm.SentencePieceProcessor('./tokenizer.model')
vocab_size = tokenizer.GetPieceSize()
for i in range(vocab_size):
    piece = tokenizer.IdToPiece(i)
    if '\r' in piece:
        print(f"{i}: {list(piece)}")

一共有 24 24 24 个：

2104: [';', '\r']
3238: ['>', '\r']
3336: ['▁', '{', '\r']
4970: ['▁', '}', '\r']
6075: [')', ';', '\r']
6756: ['▁', '\r']
8117: ['}', '\r']
8443: [')', '\r']
10175: ['"', '>', '\r']
11167: [',', '\r']
14078: ['(', ')', ';', '\r']
14626: ['{', '\r']
15231: ['"', ',', '\r']
16737: ['%', ';', '\r']
17822: ["'", ')', ';', '\r']
18584: ['"', ')', ';', '\r']
19451: ['"', '\r']
22993: ['.', '\r']
23592: ["'", ',', '\r']
24426: ['▁', '}', ')', ';', '\r']
25982: ['"', ';', '\r']
26471: ['(', ')', '\r']
29722: ['▁', '*', '/', '\r']
30004: ['\r']

可以看出，位于索引 30004 处的token刚好就是 \r，而其他token则均是以 \r 结尾。

很多场景下，我们可能需要保证 jsonl 和 txt 的行数一致才能进行下一步，为此有两种选择方案：

无论是哪种方案，都涉及到对是否包含的判断。设24个token构成的列表为 a，某个文档的分词结果为 b，于是问题便归结为判断 a 中是否有元素出现在了 b 中（反之亦然）。根据这篇博客的结论，我们可以用集合法来快速判断。

tokenizer = spm.SentencePieceProcessor('./tokenizer.model')
stop = {2104, 3238, 3336, 4970, 6075, 6756, 8117, 8443, 10175, 11167, 14078, 14626, 15231, 16737, 17822, 18584, 19451, 22993, 23592, 24426, 25982, 26471, 29722, 30004}

def valid(ids):
    return not bool(set(ids).intersection(stop))

with jsonlines.open(jsonl_file, mode='r') as r:
    for idx, line in tqdm(enumerate(r, start=1)):
        ids = tokenizer.Encode(line['text'])
        if not valid(ids):
            print(f"Line {idx} is invalid data!")