basic-tokenizer/basic_tokenizer.py at main · vision9527/basic-tokenizer · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
class BaiscTokenizer:

    def __init__(self):
        self.vocab = {} # int -> bytes
        self.merge_rule = {} # (int, int) ->  int

    def train(self, text, vocab_size):
        """
        接受训练token的数据集，使用bpe算法生成词汇表vocab和合并规则merge_rule
        """
        self.vocab = {i:bytes([i]) for i in range(256)} # 初始化基本字符
        token_ids = text.encode("utf-8")
        for new_id in range(256, vocab_size):
            # 遍历数据，找到频率最高的字节对
            stats = get_stat(token_ids)
            if not stats:
                break
            max_pair = max(stats, key=stats.get)
            # 将原始token_ids中的最高频率字节对，替换成new_id
            token_ids = merge_rule(token_ids, max_pair, new_id)
            # 将new_id加入vocab和merge_rule
            self.vocab[new_id] = self.vocab[max_pair[0]] + self.vocab[max_pair[1]]
            self.merge_rule[max_pair] = new_id

    def encode(self, text):
        """
        将输入的文本，编码成token id
        """
        res_ids = text.encode("utf-8")
        # res_ids: [12, 453, 12, 3243, 545, 2, 1, 12, 3243]
        while True:
            stats = get_stat(res_ids)
            target_pair = None
            temp_pair_id = 999999999999
            for pair, frenq in stats.items():
                if frenq < temp_pair_id:
                    target_pair = pair
                    temp_pair_id = frenq
            id = self.merge_rule.get(target_pair)
            if not id:
                break
            res_ids = merge_rule(res_ids, target_pair, id)
        return res_ids

    def decode(self, token_ids):
        """
        将输入的token id解码成文本字符串
        """
        return "".join([self.vocab[i].decode("utf-8") for i in token_ids])

def merge_rule(ids, pair, new_id):
    """
    将ids中出现的pair替换成new_id
    """
    res_ids = []
    start = 0
    while start < len(ids):
        if start < len(ids)-1 and ids[start] == pair[0] and ids[start+1] == pair[1]:
            res_ids.append(new_id)
            start = start + 2
        else:
            res_ids.append(ids[start])
            start = start + 1
    return res_ids


def get_stat(ids):
    """
    计算字节对频率
    """
    res = {}
    for i in range(len(ids)-1):
        pair = (ids[i], ids[i+1])
        res[pair] = res.get(pair, 0) + 1
    return res

if __name__ == "__main__":
    tokenizer = BaiscTokenizer()
    text = "hello world"
    tokenizer.train(text, 300)
    for k, v in tokenizer.vocab.items():
        print(f'vocab id:{k}, bytes: {v}')
    for k, v in tokenizer.merge_rule.items():
        print(f'merge_rule pair:{k}, id: {v}')

    text = "hello"
    encode_ids = tokenizer.encode(text)
    print(f'encode_ids:{encode_ids}')
    decode_text = tokenizer.decode(encode_ids)
    print(f'decode_text: {decode_text}')
    print(text==decode_text)