cwenzi/neuroflow-cpp

1import json
2import os
3import argparse
4import logging
5
6logger = logging.getLogger(__name__)
7
8SPECIAL_TOKENS = {"<pad>": 0, "<s>": 1, "</s>": 2, "<unk>": 3}
9
10
11def build_chinese_chars(start_id=4, count=3755):
12    vocab = {}
13    idx = start_id
14    for code in range(0x4E00, 0x4E00 + count):
15        if idx >= start_id + count:
16            break
17        char = chr(code)
18        vocab[char] = idx
19        idx += 1
20    return vocab
21
22
23def build_english_subwords(start_id, count=500):
24    vocab = {}
25    idx = start_id
26    prefixes = ["un", "re", "pre", "dis", "mis", "over", "under", "out", "sub", "inter", "anti", "non", "semi", "multi", "bi", "co", "ex", "in", "im", "il", "ir"]
27    suffixes = ["ing", "ed", "er", "est", "ly", "tion", "sion", "ment", "ness", "ity", "ous", "ive", "able", "ible", "ful", "less", "al", "ial", "ic", "ical"]
28    roots = ["the", "be", "to", "of", "and", "a", "in", "that", "have", "it", "for", "not", "on", "with", "he", "as", "you", "do", "at", "this", "but", "his", "by", "from", "they", "we", "say", "her", "she", "or", "an", "will", "my", "one", "all", "would", "there", "their", "what", "so", "up", "out", "if", "about", "who", "get", "which", "go", "me", "when", "make", "can", "like", "time", "no", "just", "him", "know", "take", "people", "into", "year", "your", "good", "some", "could", "them", "see", "other", "than", "then", "now", "look", "only", "come", "its", "over", "think", "also", "back", "after", "use", "two", "how", "our", "work", "first", "well", "way", "even", "new", "want", "because", "any", "these", "give", "day", "most", "us"]
29    for r in roots:
30        if idx >= start_id + count:
31            break
32        vocab[r] = idx
33        idx += 1
34    for p in prefixes:
35        for r in roots[:50]:
36            if idx >= start_id + count:
37                break
38            token = p + r
39            if token not in vocab:
40                vocab[token] = idx
41                idx += 1
42    for s in suffixes:
43        for r in roots[:50]:
44            if idx >= start_id + count:
45                break
46            token = r + s
47            if token not in vocab:
48                vocab[token] = idx
49                idx += 1
50    return vocab
51
52
53def build_punctuation(start_id):
54    vocab = {}
55    idx = start_id
56    puncts = list("。，、；：！？…—""''（）【】《》·～、.,!?;:\"'()[]{}<>-_/\\|@#$%^&*+=~`")
57    for p in puncts:
58        vocab[p] = idx
59        idx += 1
60    return vocab
61
62
63def build_digits(start_id):
64    vocab = {}
65    idx = start_id
66    for d in "0123456789":
67        vocab[d] = idx
68        idx += 1
69    for combo in ["10", "00", "01", "20", "30", "50", "100", "200", "500", "1000"]:
70        vocab[combo] = idx
71        idx += 1
72    return vocab
73
74
75def build_merges(vocab, max_merges=2000):
76    merges = []
77    tokens = sorted(vocab.keys(), key=lambda t: len(t), reverse=True)
78    for i, t1 in enumerate(tokens):
79        if len(merges) >= max_merges:
80            break
81        for t2 in tokens[i:i+10]:
82            if len(merges) >= max_merges:
83                break
84            combined = t1 + t2
85            if combined in vocab:
86                merges.append(f"{t1} {t2}")
87    return merges
88
89
90def generate_tokenizer_json(vocab_size=5000):
91    vocab = dict(SPECIAL_TOKENS)
92    next_id = len(vocab)
93
94    cn = build_chinese_chars(next_id, 3755)
95    vocab.update(cn)
96    next_id = max(vocab.values()) + 1
97
98    en = build_english_subwords(next_id, 500)
99    vocab.update(en)
100    next_id = max(vocab.values()) + 1
101
102    punct = build_punctuation(next_id)
103    vocab.update(punct)
104    next_id = max(vocab.values()) + 1
105
106    digits = build_digits(next_id)
107    vocab.update(digits)
108
109    while len(vocab) > vocab_size:
110        max_id = max(vocab.values())
111        for k, v in list(vocab.items()):
112            if v == max_id:
113                del vocab[k]
114                break
115
116    merges = build_merges(vocab)
117
118    return {
119        "model_type": "bpe",
120        "vocab_size": vocab_size,
121        "special_tokens": SPECIAL_TOKENS,
122        "vocab": vocab,
123        "merges": merges,
124    }
125
126
127if __name__ == "__main__":
128    logging.basicConfig(level=logging.INFO)
129    parser = argparse.ArgumentParser(description="NeuroFlow BPE词表适配器")
130    parser.add_argument("--vocab-size", type=int, default=5000, help="词表大小")
131    parser.add_argument("--output", type=str, default="configs/tokenizer_cn_013.json", help="输出路径")
132    args = parser.parse_args()
133
134    tokenizer = generate_tokenizer_json(args.vocab_size)
135    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
136    with open(args.output, "w", encoding="utf-8") as f:
137        json.dump(tokenizer, f, indent=2, ensure_ascii=False)
138    logger.info(f"词表已保存到 {args.output}, vocab_size={len(tokenizer['vocab'])}, merges={len(tokenizer['merges'])}")