cwenzi/neuroflow-cpp
1
1"""
2preprocess_corpus.py — 语料预处理：文本 → TOK1 二进制格式
3
4TOK1 格式规范：
5 Magic: 4 bytes "TOK1"
6 Version: 2 bytes (uint16 LE) = 1
7 Header: 4 bytes (uint32 LE) = tokenizer_vocab_size
8 4 bytes (uint32 LE) = max_seq_len
9 4 bytes (uint32 LE) = total_sample_count (占位，写完回填)
10 Samples: [2 bytes seq_len (uint16 LE)] [seq_len * 4 bytes token_ids (uint32 LE each)] ...
11 Footer: 4 bytes "END\0"
12
13支持输入格式：txt, json, jsonl, csv, 目录递归
14"""
15
16import os
17import sys
18import json
19import struct
20import argparse
21import logging
22from pathlib import Path
23from collections.abc import Iterator
24
25logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s')
26logger = logging.getLogger(__name__)
27
28MAGIC = b'TOK1'
29VERSION = 1
30FOOTER = b'END\x00'
31
32SPECIAL_TOKENS = {"<pad>": 0, "<s>": 1, "</s>": 2, "<unk>": 3}
33
34
35class SimpleBPETokenizer:
36 """轻量级BPE分词器，兼容tokenizer_cn_013.json格式"""
37
38 def __init__(self, vocab_path: str):
39 with open(vocab_path, 'r', encoding='utf-8') as f:
40 data = json.load(f)
41
42 if isinstance(data, dict) and 'model' in data:
43 self.vocab = data['model'].get('vocab', {})
44 self.merges = data['model'].get('merges', [])
45 elif isinstance(data, dict) and 'vocab' in data:
46 self.vocab = data['vocab']
47 self.merges = data.get('merges', [])
48 else:
49 self.vocab = data if isinstance(data, dict) else {}
50 self.merges = []
51
52 self.id_to_token = {v: k for k, v in self.vocab.items()}
53 self.vocab_size = len(self.vocab)
54
55 self.merge_ranks = {}
56 for i, merge in enumerate(self.merges):
57 if isinstance(merge, list) and len(merge) == 2:
58 pair = (merge[0], merge[1])
59 elif isinstance(merge, str) and ' ' in merge:
60 parts = merge.split(' ', 1)
61 pair = (parts[0], parts[1])
62 else:
63 continue
64 self.merge_ranks[pair] = i
65
66 logger.info(f"词表加载: vocab_size={self.vocab_size}, merges={len(self.merge_ranks)}")
67
68 def _tokenize_word(self, word: str) -> list:
69 tokens = list(word) if len(word) > 0 else []
70 if not tokens:
71 return tokens
72
73 while len(tokens) >= 2:
74 best_pair = None
75 best_rank = float('inf')
76 for i in range(len(tokens) - 1):
77 pair = (tokens[i], tokens[i + 1])
78 rank = self.merge_ranks.get(pair, float('inf'))
79 if rank < best_rank:
80 best_rank = rank
81 best_pair = pair
82 best_pos = i
83
84 if best_pair is None or best_rank == float('inf'):
85 break
86
87 tokens[best_pos] = best_pair[0] + best_pair[1]
88 tokens.pop(best_pos + 1)
89
90 return tokens
91
92 def encode(self, text: str, max_seq_len: int = 128) -> list:
93 if not text or not text.strip():
94 return []
95
96 text = text.strip()
97 words = []
98 current = []
99 for ch in text:
100 if '\u4e00' <= ch <= '\u9fff':
101 if current:
102 words.append(''.join(current))
103 current = []
104 words.append(ch)
105 elif ch.isspace():
106 if current:
107 words.append(''.join(current))
108 current = []
109 else:
110 current.append(ch)
111 if current:
112 words.append(''.join(current))
113
114 token_ids = [SPECIAL_TOKENS.get("<s>", 1)]
115 for word in words:
116 sub_tokens = self._tokenize_word(word)
117 for st in sub_tokens:
118 tid = self.vocab.get(st, SPECIAL_TOKENS.get("<unk>", 3))
119 token_ids.append(tid)
120 if len(token_ids) >= max_seq_len - 1:
121 break
122 if len(token_ids) >= max_seq_len - 1:
123 break
124
125 token_ids.append(SPECIAL_TOKENS.get("</s>", 2))
126 return token_ids[:max_seq_len]
127
128
129def read_texts_from_json(file_path: str) -> Iterator[str]:
130 """流式读取JSON文件中的文本字段"""
131 with open(file_path, 'r', encoding='utf-8') as f:
132 brace_depth = 0
133 record = ''
134 for line in f:
135 for ch in line:
136 if ch == '{':
137 if brace_depth == 0:
138 record = ''
139 brace_depth += 1
140 elif ch == '}':
141 brace_depth -= 1
142 if brace_depth == 0 and record:
143 record += '}'
144 for text in _extract_texts_from_record(record):
145 yield text
146 record = ''
147 if brace_depth > 0:
148 record += ch
149
150
151def _extract_texts_from_record(record: str) -> Iterator[str]:
152 """从JSON record中提取文本字段"""
153 for field in ['"text"', '"content"', '"title"', '"question"', '"answer"']:
154 pos = record.find(field)
155 if pos == -1:
156 continue
157 colon = record.find(':', pos + len(field))
158 if colon == -1:
159 continue
160 colon += 1
161 while colon < len(record) and record[colon] != '"':
162 colon += 1
163 if colon >= len(record):
164 continue
165 colon += 1
166 end = colon
167 while end < len(record) and record[end] != '"':
168 if record[end] == '\\':
169 end += 1
170 end += 1
171 text = record[colon:end]
172 text = _unescape_json(text)
173 if len(text) >= 10:
174 yield text
175
176
177def _unescape_json(s: str) -> str:
178 return s.replace('\\n', '\n').replace('\\t', '\t').replace('\\"', '"').replace('\\\\', '\\')
179
180
181def read_texts_from_jsonl(file_path: str) -> Iterator[str]:
182 """逐行读取JSONL文件"""
183 with open(file_path, 'r', encoding='utf-8') as f:
184 for line in f:
185 line = line.strip()
186 if not line:
187 continue
188 try:
189 obj = json.loads(line)
190 except json.JSONDecodeError:
191 for text in _extract_texts_from_record(line):
192 yield text
193 continue
194 for field in ['text', 'content', 'title', 'question', 'answer']:
195 if field in obj and isinstance(obj[field], str) and len(obj[field]) >= 10:
196 yield obj[field]
197
198
199def read_texts_from_txt(file_path: str) -> Iterator[str]:
200 """按段落读取纯文本文件"""
201 with open(file_path, 'r', encoding='utf-8', errors='replace') as f:
202 paragraph = ''
203 for line in f:
204 if line.strip() == '':
205 if len(paragraph) >= 10:
206 yield paragraph.strip()
207 paragraph = ''
208 else:
209 paragraph += line
210 if len(paragraph) > 10000:
211 yield paragraph.strip()
212 paragraph = ''
213 if len(paragraph) >= 10:
214 yield paragraph.strip()
215
216
217def read_texts_from_csv(file_path: str) -> Iterator[str]:
218 """读取CSV，第一列作为文本"""
219 with open(file_path, 'r', encoding='utf-8', errors='replace') as f:
220 for line in f:
221 line = line.strip()
222 if not line or line.startswith('#'):
223 continue
224 parts = line.split(',', 1)
225 text = parts[0].strip().strip('"')
226 if len(text) >= 10:
227 yield text
228
229
230def read_texts(file_path: str) -> Iterator[str]:
231 """根据扩展名自动分派读取器"""
232 ext = Path(file_path).suffix.lower()
233 if ext == '.json':
234 return read_texts_from_json(file_path)
235 elif ext == '.jsonl':
236 return read_texts_from_jsonl(file_path)
237 elif ext == '.txt':
238 return read_texts_from_txt(file_path)
239 elif ext in ('.csv', '.tsv'):
240 return read_texts_from_csv(file_path)
241 else:
242 logger.warning(f"跳过不支持的格式: {file_path}")
243 return iter([])
244
245
246def read_texts_recursive(path: str) -> Iterator[tuple]:
247 """递归遍历目录，返回 (文件路径, 文本) 对"""
248 p = Path(path)
249 if p.is_file():
250 for text in read_texts(str(p)):
251 yield (str(p), text)
252 elif p.is_dir():
253 files = sorted(p.rglob('*'))
254 for f in files:
255 if not f.is_file():
256 continue
257 ext = f.suffix.lower()
258 if ext in ('.json', '.jsonl', '.txt', '.csv', '.tsv'):
259 logger.info(f"读取: {f}")
260 try:
261 for text in read_texts(str(f)):
262 yield (str(f), text)
263 except Exception as e:
264 logger.warning(f"跳过 {f}: {e}")
265 else:
266 raise ValueError(f"路径不存在: {path}")
267
268
269def write_tok1(output_path: str, tokenizer: SimpleBPETokenizer,
270 text_iter: Iterator[tuple], max_seq_len: int = 128,
271 max_samples: int = 0, min_tokens: int = 4):
272 """将文本流写入TOK1二进制格式"""
273 total_written = 0
274 total_skipped = 0
275
276 with open(output_path, 'wb') as f:
277 f.write(MAGIC)
278 f.write(struct.pack('<H', VERSION))
279 f.write(struct.pack('<I', tokenizer.vocab_size))
280 f.write(struct.pack('<I', max_seq_len))
281 count_offset = f.tell()
282 f.write(struct.pack('<I', 0))
283
284 for file_path, text in text_iter:
285 token_ids = tokenizer.encode(text, max_seq_len)
286 if len(token_ids) < min_tokens:
287 total_skipped += 1
288 continue
289
290 seq_len = len(token_ids)
291 f.write(struct.pack('<H', seq_len))
292 for tid in token_ids:
293 f.write(struct.pack('<I', tid))
294
295 total_written += 1
296 if total_written % 10000 == 0:
297 logger.info(f"已写入 {total_written} 样本, 跳过 {total_skipped}")
298
299 if max_samples > 0 and total_written >= max_samples:
300 logger.info(f"达到采样上限 {max_samples}，停止")
301 break
302
303 f.write(FOOTER)
304
305 f.seek(count_offset)
306 f.write(struct.pack('<I', total_written))
307
308 file_size = os.path.getsize(output_path)
309 logger.info(f"写入完成: {total_written} 样本, 跳过 {total_skipped}")
310 logger.info(f"输出文件: {output_path} ({file_size / 1024 / 1024:.1f} MB)")
311
312
313def split_large_json(input_path: str, output_dir: str, chunk_size: int = 50000):
314 """将大型JSON文件拆分为JSONL分片"""
315 os.makedirs(output_dir, exist_ok=True)
316 chunk_idx = 0
317 count = 0
318 out = None
319
320 for text in read_texts_from_json(input_path):
321 if out is None:
322 chunk_path = os.path.join(output_dir, f'chunk_{chunk_idx:04d}.jsonl')
323 out = open(chunk_path, 'w', encoding='utf-8')
324 logger.info(f"开始分片: {chunk_path}")
325
326 obj = json.dumps({"text": text}, ensure_ascii=False)
327 out.write(obj + '\n')
328 count += 1
329
330 if count >= chunk_size:
331 out.close()
332 logger.info(f"分片 {chunk_idx} 完成: {count} 条")
333 chunk_idx += 1
334 count = 0
335 out = None
336
337 if out:
338 out.close()
339 logger.info(f"分片 {chunk_idx} 完成: {count} 条")
340
341 logger.info(f"拆分完成: {chunk_idx + 1} 个分片")
342
343
344def verify_tok1(file_path: str, max_show: int = 5):
345 """验证TOK1文件格式"""
346 with open(file_path, 'rb') as f:
347 magic = f.read(4)
348 assert magic == MAGIC, f"Magic不匹配: {magic}"
349 version = struct.unpack('<H', f.read(2))[0]
350 vocab_size = struct.unpack('<I', f.read(4))[0]
351 max_seq_len = struct.unpack('<I', f.read(4))[0]
352 total = struct.unpack('<I', f.read(4))[0]
353
354 logger.info(f"TOK1文件: version={version}, vocab_size={vocab_size}, "
355 f"max_seq_len={max_seq_len}, total_samples={total}")
356
357 shown = 0
358 for i in range(total):
359 seq_len = struct.unpack('<H', f.read(2))[0]
360 ids = []
361 for _ in range(seq_len):
362 ids.append(struct.unpack('<I', f.read(4))[0])
363 if shown < max_show:
364 logger.info(f" 样本{i}: len={seq_len}, ids[:10]={ids[:10]}")
365 shown += 1
366
367 footer = f.read(4)
368 assert footer == FOOTER, f"Footer不匹配: {footer}"
369
370 logger.info("验证通过")
371
372
373def main():
374 parser = argparse.ArgumentParser(description='语料预处理: 文本 → TOK1 二进制格式')
375 sub = parser.add_subparsers(dest='command')
376
377 p_process = sub.add_parser('process', help='预处理语料为TOK1格式')
378 p_process.add_argument('--input', required=True, help='输入路径(文件或目录)')
379 p_process.add_argument('--tokenizer', required=True, help='tokenizer JSON路径')
380 p_process.add_argument('--output', required=True, help='输出TOK1文件路径')
381 p_process.add_argument('--max-seq-len', type=int, default=128, help='最大序列长度')
382 p_process.add_argument('--max-samples', type=int, default=0, help='最大样本数(0=不限)')
383 p_process.add_argument('--min-tokens', type=int, default=4, help='最小token数')
384
385 p_split = sub.add_parser('split', help='拆分大型JSON为JSONL分片')
386 p_split.add_argument('--input', required=True, help='输入JSON文件')
387 p_split.add_argument('--output-dir', required=True, help='输出目录')
388 p_split.add_argument('--chunk-size', type=int, default=50000, help='每片条数')
389
390 p_verify = sub.add_parser('verify', help='验证TOK1文件')
391 p_verify.add_argument('--input', required=True, help='TOK1文件路径')
392
393 args = parser.parse_args()
394
395 if args.command == 'process':
396 tokenizer = SimpleBPETokenizer(args.tokenizer)
397 text_iter = read_texts_recursive(args.input)
398 write_tok1(args.output, tokenizer, text_iter,
399 max_seq_len=args.max_seq_len,
400 max_samples=args.max_samples,
401 min_tokens=args.min_tokens)
402 elif args.command == 'split':
403 split_large_json(args.input, args.output_dir, args.chunk_size)
404 elif args.command == 'verify':
405 verify_tok1(args.input)
406 else:
407 parser.print_help()
408
409
410if __name__ == '__main__':
411 main()