1 line
37 KiB
JSON
1 line
37 KiB
JSON
{"repo": "albertyang514/albertlm", "files": {"tools/deepseek_distill.py": {"dd3bf58f5dc3579788cfd043d42e59c78ceba7d2": {"t": 1785745440.8197756, "n": 16983, "h": "c5c6e3b0d5383e20", "c": "#!/usr/bin/env python3\nfrom __future__ import annotations\n\nimport argparse\nimport asyncio\nimport hashlib\nimport json\nimport math\nimport os\nimport random\nimport re\nimport sys\nimport time\nimport unicodedata\nfrom pathlib import Path\nfrom typing import Any\n\nROOT = Path('/data/AlbertLM')\nWORK = ROOT / 'data' / 'synthetic' / 'deepseek-v4-pro'\nJOBS = WORK / 'jobs.jsonl'\nRAW = WORK / 'raw.jsonl'\nERRORS = WORK / 'errors.jsonl'\nCLEAN = WORK / 'clean.jsonl'\nCORPUS = WORK / 'corpus.txt'\nBIN = WORK / 'train-deepseek-v4-pro.bin'\nSTATS = WORK / 'stats.json'\n\nSYSTEM_PROMPT = '''你正在为通用语言模型生成高质量合成预训练语料。只输出最终正文，不要提到提示词、数据集、蒸馏、AI身份或生成过程。正文必须自洽、准确、信息密度高；不确定的具体事实不要编造。需要推导时给出可公开阅读的解释和关键步骤，但不要输出隐藏思维过程。代码必须完整且有解释，数学公式必须定义符号。避免空泛套话、营销文案、重复段落和机械总结。'''\n\nCATEGORIES: dict[str, list[str]] = {\n 'science': ['物理学', '化学', '生物学', '天文学', '地球科学', '环境科学', '材料科学', '神经科学'],\n 'mathematics': ['代数', '概率统计', '微积分', '离散数学', '线性代数', '数论', '优化方法', '信息论'],\n 'software': ['Python', '操作系统', '计算机网络', '数据库', '编译原理', '分布式系统', '软件工程', '网络安全', '机器学习系统', '高性能计算'],\n 'engineering': ['电子工程', '控制理论', '通信工程', '机械工程', '能源系统', '土木工程', '工业设计', '可靠性工程'],\n 'humanities': ['世界史', '中国史', '哲学', '语言学', '文学理论', '艺术史', '宗教学研究', '考古学'],\n 'social_science': ['经济学', '社会学', '政治学', '法学基础', '心理学', '公共管理', '传播学', '国际关系'],\n 'practical': ['个人财务基础', '项目管理', '科学写作', '信息检索', '批判性思维', '实验设计', '数据分析', '职业沟通'],\n 'creative': ['叙事结构', '人物塑造', '场景描写', '科幻设定', '悬疑结构', '散文写作', '诗歌鉴赏', '跨文化写作'],\n}\n\nANGLES = [\n '核心概念与常见误区', '从基本原理到实际应用', '一个具体问题的系统分析', '历史演变与现代方法',\n '关键机制、限制与权衡', '面向初学者但保持技术严谨', '面向工程实践的完整说明', '用例子解释抽象概念',\n '比较两种主流方法并说明适用边界', '给出可验证的推导或实验思路', '从失败案例反推正确方法', '整理为结构清晰的参考文章',\n]\n\nLANGUAGES = ['zh-CN'] * 12 + ['en'] * 5 + ['ja'] * 2\n\n\ndef write_jsonl(path: Path, obj: dict[str, Any]) -> None:\n path.parent.mkdir(parents=True, exist_ok=True)\n with path.open('a', encoding='utf-8') as f:\n f.write(json.dumps(obj, ensure_ascii=False) + '\\n')\n\n\ndef iter_jsonl(path: Path):\n if not path.exists():\n return\n with path.open('r', encoding='utf-8', errors='replace') as f:\n for line_no, line in enumerate(f, 1):\n line = line.strip()\n if not line:\n continue\n try:\n yield json.loads(line)\n except json.JSONDecodeError:\n print(f'跳过损坏 JSONL 行: {path}:{line_no}', file=sys.stderr)\n\n\ndef make_jobs(args: argparse.Namespace) -> None:\n WORK.mkdir(parents=True, exist_ok=True)\n rng = random.Random(args.seed)\n categories = list(CATEGORIES)\n with JOBS.open('w', encoding='utf-8') as f:\n for i in range(args.count):\n category = categories[i % len(categories)]\n topic = rng.choice(CATEGORIES[category])\n angle = rng.choice(ANGLES)\n language = rng.choice(LANGUAGES)\n diversity_key = hashlib.sha256(f'{args.seed}:{i}'.encode()).hexdigest()[:12]\n if language == 'zh-CN':\n request = (\n f'写一篇约1200至2200字的高质量参考正文。领域：{category}；主题范围：{topic}；'\n f'切入角度：{angle}；多样性键：{diversity_key}。请选择一个具体且不过度常见的子题，'\n '正文应自成一体，包含必要定义、机制、例子和边界条件。不要写问答格式，不要写“下面将”。'\n )\n elif language == 'en':\n request = (\n f'Write a self-contained high-quality reference article of roughly 900-1600 words. '\n f'Domain: {category}; topic area: {topic}; angle: {angle}; diversity key: {diversity_key}. '\n 'Choose a specific, non-obvious subtopic. Include definitions, mechanisms, examples, limitations, '\n 'and practical implications. Do not use a Q&A format or mention these instructions.'\n )\n else:\n request = (\n f'約1500〜2500字の、独立して読める高品質な日本語の解説文を書いてください。'\n f'分野：{category}、主題：{topic}、観点：{angle}、多様性キー：{diversity_key}。'\n '具体的で一般的すぎない小テーマを選び、定義、仕組み、例、限界を含めてください。'\n '問答形式や指示への言及は避けてください。'\n )\n job = {\n 'id': f'ds4p-{i:08d}',\n 'category': category,\n 'topic': topic,\n 'angle': angle,\n 'language': language,\n 'diversity_key': diversity_key,\n 'messages': [\n {'role': 'system', 'content': SYSTEM_PROMPT},\n {'role': 'user', 'content': request},\n ],\n }\n f.write(json.dumps(job, ensure_ascii=False) + '\\n')\n print(f'jobs={args.count} path={JOBS}')\n\n\nasync def generate(args: argparse.Namespace) -> None:\n try:\n import httpx\n except ImportError as exc:\n raise SystemExit('缺少 httpx：pip install httpx') from exc\n\n api_key = os.environ.get('DEEPSEEK_API_KEY')\n if not api_key:\n raise SystemExit('未设置 DEEPSEEK_API_KEY')\n if not JOBS.exists():\n raise SystemExit(f'不存在 {JOBS}，先运行 make-jobs')\n\n completed = {str(x.get('id')) for x in iter_jsonl(RAW) or [] if x.get('id')}\n jobs = [x for x in iter_jsonl(JOBS) or [] if str(x.get('id')) not in completed]\n if args.limit > 0:\n jobs = jobs[:args.limit]\n print(f'total_pending={len(jobs)} completed={len(completed)} concurrency={args.concurrency}')\n if not jobs:\n return\n\n queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue()\n for job in jobs:\n queue.put_nowait(job)\n for _ in range(args.concurrency):\n queue.put_nowait(None)\n\n write_lock = asyncio.Lock()\n counter_lock = asyncio.Lock()\n done = 0\n failed = 0\n started = time.time()\n\n timeout = httpx.Timeout(args.timeout, connect=30.0, read=args.timeout, write=60.0, pool=60.0)\n limits = httpx.Limits(max_connections=args.concurrency, max_keepalive_connections=args.concurrency)\n\n async with httpx.AsyncClient(\n base_url='https://api.deepseek.com',\n headers={'Authorization': f'Bearer {api_key}', 'Content-Type': 'application/json'},\n timeout=timeout,\n limits=limits,\n http2=True,\n ) as client:\n async def worker(worker_id: int) -> None:\n nonlocal done, failed\n while True:\n job = await queue.get()\n if job is None:\n queue.task_done()\n return\n last_error: Exception | None = None\n for attempt in range(args.retries):\n try:\n payload = {\n 'model': args.model,\n 'messages': job['messages'],\n 'thinking': {'type': 'disabled'},\n 'temperature': args.temperature,\n 'max_tokens': args.max_tokens,\n 'stream': False,\n 'user_id': 'albertlm-distill',\n }\n response = await client.post('/chat/completions', json=payload)\n if response.status_code in {429, 500, 502, 503, 504}:\n raise RuntimeError(f'HTTP {response.status_code}: {response.text[:300]}')\n response.raise_for_status()\n data = response.json()\n choice = data['choices'][0]\n content = choice['message'].get('content') or ''\n if not content.strip():\n raise RuntimeError('empty output')\n result = {\n 'id': job['id'],\n 'category': job['category'],\n 'topic': job['topic'],\n 'angle': job['angle'],\n 'language': job['language'],\n 'prompt': job['messages'][-1]['content'],\n 'output': content,\n 'model': data.get('model', args.model),\n 'finish_reason': choice.get('finish_reason'),\n 'usage': data.get('usage', {}),\n 'created': data.get('created'),\n 'request_id': response.headers.get('x-request-id'),\n 'timestamp': time.time(),\n }\n async with write_lock:\n write_jsonl(RAW, result)\n async with counter_lock:\n done += 1\n if done % 10 == 0 or done == len(jobs):\n elapsed = max(time.time() - started, 1e-6)\n print(f'done={done}/{len(jobs)} failed={failed} rate={done/elapsed:.2f} req/s', flush=True)\n last_error = None\n break\n except Exception as exc: # noqa: BLE001\n last_error = exc\n if attempt + 1 < args.retries:\n delay = min(60.0, (2 ** attempt) + random.random() * 2)\n await asyncio.sleep(delay)\n if last_error is not None:\n async with write_lock:\n write_jsonl(ERRORS, {\n 'id': job.get('id'),\n 'error_type': type(last_error).__name__,\n 'error': str(last_error),\n 'timestamp': time.time(),\n })\n async with counter_lock:\n failed += 1\n queue.task_done()\n\n workers = [asyncio.create_task(worker(i)) for i in range(args.concurrency)]\n await queue.join()\n await asyncio.gather(*workers)\n print(f'finished done={done} failed={failed} raw={RAW}')\n\n\ndef normalize_text(text: str) -> str:\n text = unicodedata.normalize('NFC', text)\n text = text.replace('\\x00', '')\n text = re.sub(r'\\r\\n?', '\\n', text)\n text = text.replace('\\t', ' ')\n text = re.sub(r'[ \\t]+\\n', '\\n', text)\n text = re.sub(r'\\n{4,}', '\\n\\n\\n', text)\n return text.strip()\n\n\ndef repeated_line_ratio(text: str) -> float:\n lines = [x.strip() for x in text.splitlines() if x.strip()]\n if len(lines) < 6:\n return 0.0\n return 1.0 - (len(set(lines)) / len(lines))\n\n\ndef clean(args: argparse.Namespace) -> None:\n if not RAW.exists():\n raise SystemExit(f'不存在 {RAW}')\n seen: set[str] = set()\n kept = 0\n rejected = 0\n chars = 0\n usage_in = 0\n usage_out = 0\n refusal_markers = ('作为一个ai', '作为 ai', 'i cannot comply', 'i can\\'t comply', '申し訳ありませんが')\n with CLEAN.open('w', encoding='utf-8') as clean_file, CORPUS.open('w', encoding='utf-8') as corpus_file:\n for item in iter_jsonl(RAW) or []:\n text = normalize_text(str(item.get('output') or ''))\n reason = None\n if item.get('finish_reason') != 'stop':\n reason = 'incomplete'\n elif len(text) < args.min_chars:\n reason = 'too_short'\n elif len(text) > args.max_chars:\n reason = 'too_long'\n elif any(marker in text[:300].lower() for marker in refusal_markers):\n reason = 'refusal'\n elif repeated_line_ratio(text) > args.max_repeated_line_ratio:\n reason = 'repetition'\n digest = hashlib.sha256(re.sub(r'\\s+', ' ', text).encode('utf-8')).hexdigest()\n if digest in seen:\n reason = 'duplicate'\n if reason:\n rejected += 1\n continue\n seen.add(digest)\n output = {\n 'id': item.get('id'),\n 'category': item.get('category'),\n 'topic': item.get('topic'),\n 'language': item.get('language'),\n 'prompt': item.get('prompt'),\n 'output': text,\n 'sha256': digest,\n 'model': item.get('model'),\n 'usage': item.get('usage', {}),\n }\n clean_file.write(json.dumps(output, ensure_ascii=False) + '\\n')\n corpus_file.write(text + '\\n\\n')\n kept += 1\n chars += len(text)\n usage = item.get('usage') or {}\n usage_in += int(usage.get('prompt_tokens') or 0)\n usage_out += int(usage.get('completion_tokens') or 0)\n stats = {\n 'kept_documents': kept,\n 'rejected_documents': rejected,\n 'characters': chars,\n 'api_prompt_tokens': usage_in,\n 'api_completion_tokens': usage_out,\n 'clean_jsonl': str(CLEAN),\n 'corpus_txt': str(CORPUS),\n }\n STATS.write_text(json.dumps(stats, ensure_ascii=False, indent=2) + '\\n', encoding='utf-8')\n print(json.dumps(stats, ensure_ascii=False, indent=2))\n\n\ndef find_special_id(tokenizer, names: list[str]) -> int | None:\n for name in names:\n token_id = tokenizer.token_to_id(name)\n if token_id is not None:\n return int(token_id)\n return None\n\n\ndef tokenize(args: argparse.Namespace) -> None:\n try:\n import numpy as np\n from tokenizers import Tokenizer\n except ImportError as exc:\n raise SystemExit('缺少 numpy/tokenizers') from exc\n tokenizer_path = Path(args.tokenizer)\n if not tokenizer_path.exists():\n raise SystemExit(f'不存在 tokenizer: {tokenizer_path}')\n if not CLEAN.exists():\n raise SystemExit(f'不存在 {CLEAN}，先运行 clean')\n tokenizer = Tokenizer.from_file(str(tokenizer_path))\n bos_id = find_special_id(tokenizer, ['<bos>', '<s>', '[BOS]'])\n eos_id = find_special_id(tokenizer, ['<eos>', '</s>', '[EOS]'])\n unk_id = tokenizer.token_to_id('<unk>')\n total_tokens = 0\n documents = 0\n with BIN.open('wb') as out:\n for item in iter_jsonl(CLEAN) or []:\n text = str(item['output'])\n ids = tokenizer.encode(text, add_special_tokens=False).ids\n if unk_id is not None and unk_id in ids:\n raise RuntimeError(f\"document {item.get('id')} contains <unk>\")\n if bos_id is not None:\n ids.insert(0, bos_id)\n if eos_id is not None:\n ids.append(eos_id)\n if not ids:\n continue\n np.asarray(ids, dtype=np.uint32).tofile(out)\n total_tokens += len(ids)\n documents += 1\n stats = {}\n if STATS.exists():\n try:\n stats = json.loads(STATS.read_text(encoding='utf-8'))\n except json.JSONDecodeError:\n stats = {}\n stats.update({\n 'tokenized_documents': documents,\n 'albertlm_tokens': total_tokens,\n 'binary_path': str(BIN),\n 'binary_bytes': BIN.stat().st_size,\n 'tokenizer_path': str(tokenizer_path),\n 'bos_id': bos_id,\n 'eos_id': eos_id,\n })\n STATS.write_text(json.dumps(stats, ensure_ascii=False, indent=2) + '\\n', encoding='utf-8')\n print(json.dumps(stats, ensure_ascii=False, indent=2))\n\n\ndef status(_: argparse.Namespace) -> None:\n def count(path: Path) -> int:\n return sum(1 for _ in iter_jsonl(path) or []) if path.exists() else 0\n print(f'jobs={count(JOBS)} raw={count(RAW)} errors={count(ERRORS)} clean={count(CLEAN)}')\n for path in [JOBS, RAW, ERRORS, CLEAN, CORPUS, BIN, STATS]:\n if path.exists():\n print(f'{path}: {path.stat().st_size / 1024 / 1024:.2f} MiB')\n if STATS.exists():\n print(STATS.read_text(encoding='utf-8'))\n\n\ndef main() -> int:\n parser = argparse.ArgumentParser()\n sub = parser.add_subparsers(dest='command', required=True)\n\n p = sub.add_parser('make-jobs')\n p.add_argument('--count', type=int, default=10000)\n p.add_argument('--seed', type=int, default=20260715)\n p.set_defaults(func=make_jobs)\n\n p = sub.add_parser('generate')\n p.add_argument('--model', default='deepseek-v4-pro')\n p.add_argument('--concurrency', type=int, default=32)\n p.add_argument('--max-tokens', type=int, default=1800)\n p.add_argument('--temperature', type=float, default=0.7)\n p.add_argument('--timeout', type=float, default=900.0)\n p.add_argument('--retries', type=int, default=8)\n p.add_argument('--limit', type=int, default=0, help='0 means all pending jobs')\n p.set_defaults(func=lambda args: asyncio.run(generate(args)))\n\n p = sub.add_parser('clean')\n p.add_argument('--min-chars', type=int, default=500)\n p.add_argument('--max-chars', type=int, default=30000)\n p.add_argument('--max-repeated-line-ratio', type=float, default=0.35)\n p.set_defaults(func=clean)\n\n p = sub.add_parser('tokenize')\n p.add_argument('--tokenizer', default=str(ROOT / 'tokenizer' / 'tokenizer.json'))\n p.set_defaults(func=tokenize)\n\n p = sub.add_parser('status')\n p.set_defaults(func=status)\n\n args = parser.parse_args()\n args.func(args)\n return 0\n\n\nif __name__ == '__main__':\n raise SystemExit(main())\n"}, "afb9a3c73a1a559dd7c621dfc2fef6a9f7cce9b4": {"t": 1785753252.0837462, "n": 16983, "h": "c5c6e3b0d5383e20", "c": "#!/usr/bin/env python3\nfrom __future__ import annotations\n\nimport argparse\nimport asyncio\nimport hashlib\nimport json\nimport math\nimport os\nimport random\nimport re\nimport sys\nimport time\nimport unicodedata\nfrom pathlib import Path\nfrom typing import Any\n\nROOT = Path('/data/AlbertLM')\nWORK = ROOT / 'data' / 'synthetic' / 'deepseek-v4-pro'\nJOBS = WORK / 'jobs.jsonl'\nRAW = WORK / 'raw.jsonl'\nERRORS = WORK / 'errors.jsonl'\nCLEAN = WORK / 'clean.jsonl'\nCORPUS = WORK / 'corpus.txt'\nBIN = WORK / 'train-deepseek-v4-pro.bin'\nSTATS = WORK / 'stats.json'\n\nSYSTEM_PROMPT = '''你正在为通用语言模型生成高质量合成预训练语料。只输出最终正文，不要提到提示词、数据集、蒸馏、AI身份或生成过程。正文必须自洽、准确、信息密度高；不确定的具体事实不要编造。需要推导时给出可公开阅读的解释和关键步骤，但不要输出隐藏思维过程。代码必须完整且有解释，数学公式必须定义符号。避免空泛套话、营销文案、重复段落和机械总结。'''\n\nCATEGORIES: dict[str, list[str]] = {\n 'science': ['物理学', '化学', '生物学', '天文学', '地球科学', '环境科学', '材料科学', '神经科学'],\n 'mathematics': ['代数', '概率统计', '微积分', '离散数学', '线性代数', '数论', '优化方法', '信息论'],\n 'software': ['Python', '操作系统', '计算机网络', '数据库', '编译原理', '分布式系统', '软件工程', '网络安全', '机器学习系统', '高性能计算'],\n 'engineering': ['电子工程', '控制理论', '通信工程', '机械工程', '能源系统', '土木工程', '工业设计', '可靠性工程'],\n 'humanities': ['世界史', '中国史', '哲学', '语言学', '文学理论', '艺术史', '宗教学研究', '考古学'],\n 'social_science': ['经济学', '社会学', '政治学', '法学基础', '心理学', '公共管理', '传播学', '国际关系'],\n 'practical': ['个人财务基础', '项目管理', '科学写作', '信息检索', '批判性思维', '实验设计', '数据分析', '职业沟通'],\n 'creative': ['叙事结构', '人物塑造', '场景描写', '科幻设定', '悬疑结构', '散文写作', '诗歌鉴赏', '跨文化写作'],\n}\n\nANGLES = [\n '核心概念与常见误区', '从基本原理到实际应用', '一个具体问题的系统分析', '历史演变与现代方法',\n '关键机制、限制与权衡', '面向初学者但保持技术严谨', '面向工程实践的完整说明', '用例子解释抽象概念',\n '比较两种主流方法并说明适用边界', '给出可验证的推导或实验思路', '从失败案例反推正确方法', '整理为结构清晰的参考文章',\n]\n\nLANGUAGES = ['zh-CN'] * 12 + ['en'] * 5 + ['ja'] * 2\n\n\ndef write_jsonl(path: Path, obj: dict[str, Any]) -> None:\n path.parent.mkdir(parents=True, exist_ok=True)\n with path.open('a', encoding='utf-8') as f:\n f.write(json.dumps(obj, ensure_ascii=False) + '\\n')\n\n\ndef iter_jsonl(path: Path):\n if not path.exists():\n return\n with path.open('r', encoding='utf-8', errors='replace') as f:\n for line_no, line in enumerate(f, 1):\n line = line.strip()\n if not line:\n continue\n try:\n yield json.loads(line)\n except json.JSONDecodeError:\n print(f'跳过损坏 JSONL 行: {path}:{line_no}', file=sys.stderr)\n\n\ndef make_jobs(args: argparse.Namespace) -> None:\n WORK.mkdir(parents=True, exist_ok=True)\n rng = random.Random(args.seed)\n categories = list(CATEGORIES)\n with JOBS.open('w', encoding='utf-8') as f:\n for i in range(args.count):\n category = categories[i % len(categories)]\n topic = rng.choice(CATEGORIES[category])\n angle = rng.choice(ANGLES)\n language = rng.choice(LANGUAGES)\n diversity_key = hashlib.sha256(f'{args.seed}:{i}'.encode()).hexdigest()[:12]\n if language == 'zh-CN':\n request = (\n f'写一篇约1200至2200字的高质量参考正文。领域：{category}；主题范围：{topic}；'\n f'切入角度：{angle}；多样性键：{diversity_key}。请选择一个具体且不过度常见的子题，'\n '正文应自成一体，包含必要定义、机制、例子和边界条件。不要写问答格式，不要写“下面将”。'\n )\n elif language == 'en':\n request = (\n f'Write a self-contained high-quality reference article of roughly 900-1600 words. '\n f'Domain: {category}; topic area: {topic}; angle: {angle}; diversity key: {diversity_key}. '\n 'Choose a specific, non-obvious subtopic. Include definitions, mechanisms, examples, limitations, '\n 'and practical implications. Do not use a Q&A format or mention these instructions.'\n )\n else:\n request = (\n f'約1500〜2500字の、独立して読める高品質な日本語の解説文を書いてください。'\n f'分野：{category}、主題：{topic}、観点：{angle}、多様性キー：{diversity_key}。'\n '具体的で一般的すぎない小テーマを選び、定義、仕組み、例、限界を含めてください。'\n '問答形式や指示への言及は避けてください。'\n )\n job = {\n 'id': f'ds4p-{i:08d}',\n 'category': category,\n 'topic': topic,\n 'angle': angle,\n 'language': language,\n 'diversity_key': diversity_key,\n 'messages': [\n {'role': 'system', 'content': SYSTEM_PROMPT},\n {'role': 'user', 'content': request},\n ],\n }\n f.write(json.dumps(job, ensure_ascii=False) + '\\n')\n print(f'jobs={args.count} path={JOBS}')\n\n\nasync def generate(args: argparse.Namespace) -> None:\n try:\n import httpx\n except ImportError as exc:\n raise SystemExit('缺少 httpx：pip install httpx') from exc\n\n api_key = os.environ.get('DEEPSEEK_API_KEY')\n if not api_key:\n raise SystemExit('未设置 DEEPSEEK_API_KEY')\n if not JOBS.exists():\n raise SystemExit(f'不存在 {JOBS}，先运行 make-jobs')\n\n completed = {str(x.get('id')) for x in iter_jsonl(RAW) or [] if x.get('id')}\n jobs = [x for x in iter_jsonl(JOBS) or [] if str(x.get('id')) not in completed]\n if args.limit > 0:\n jobs = jobs[:args.limit]\n print(f'total_pending={len(jobs)} completed={len(completed)} concurrency={args.concurrency}')\n if not jobs:\n return\n\n queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue()\n for job in jobs:\n queue.put_nowait(job)\n for _ in range(args.concurrency):\n queue.put_nowait(None)\n\n write_lock = asyncio.Lock()\n counter_lock = asyncio.Lock()\n done = 0\n failed = 0\n started = time.time()\n\n timeout = httpx.Timeout(args.timeout, connect=30.0, read=args.timeout, write=60.0, pool=60.0)\n limits = httpx.Limits(max_connections=args.concurrency, max_keepalive_connections=args.concurrency)\n\n async with httpx.AsyncClient(\n base_url='https://api.deepseek.com',\n headers={'Authorization': f'Bearer {api_key}', 'Content-Type': 'application/json'},\n timeout=timeout,\n limits=limits,\n http2=True,\n ) as client:\n async def worker(worker_id: int) -> None:\n nonlocal done, failed\n while True:\n job = await queue.get()\n if job is None:\n queue.task_done()\n return\n last_error: Exception | None = None\n for attempt in range(args.retries):\n try:\n payload = {\n 'model': args.model,\n 'messages': job['messages'],\n 'thinking': {'type': 'disabled'},\n 'temperature': args.temperature,\n 'max_tokens': args.max_tokens,\n 'stream': False,\n 'user_id': 'albertlm-distill',\n }\n response = await client.post('/chat/completions', json=payload)\n if response.status_code in {429, 500, 502, 503, 504}:\n raise RuntimeError(f'HTTP {response.status_code}: {response.text[:300]}')\n response.raise_for_status()\n data = response.json()\n choice = data['choices'][0]\n content = choice['message'].get('content') or ''\n if not content.strip():\n raise RuntimeError('empty output')\n result = {\n 'id': job['id'],\n 'category': job['category'],\n 'topic': job['topic'],\n 'angle': job['angle'],\n 'language': job['language'],\n 'prompt': job['messages'][-1]['content'],\n 'output': content,\n 'model': data.get('model', args.model),\n 'finish_reason': choice.get('finish_reason'),\n 'usage': data.get('usage', {}),\n 'created': data.get('created'),\n 'request_id': response.headers.get('x-request-id'),\n 'timestamp': time.time(),\n }\n async with write_lock:\n write_jsonl(RAW, result)\n async with counter_lock:\n done += 1\n if done % 10 == 0 or done == len(jobs):\n elapsed = max(time.time() - started, 1e-6)\n print(f'done={done}/{len(jobs)} failed={failed} rate={done/elapsed:.2f} req/s', flush=True)\n last_error = None\n break\n except Exception as exc: # noqa: BLE001\n last_error = exc\n if attempt + 1 < args.retries:\n delay = min(60.0, (2 ** attempt) + random.random() * 2)\n await asyncio.sleep(delay)\n if last_error is not None:\n async with write_lock:\n write_jsonl(ERRORS, {\n 'id': job.get('id'),\n 'error_type': type(last_error).__name__,\n 'error': str(last_error),\n 'timestamp': time.time(),\n })\n async with counter_lock:\n failed += 1\n queue.task_done()\n\n workers = [asyncio.create_task(worker(i)) for i in range(args.concurrency)]\n await queue.join()\n await asyncio.gather(*workers)\n print(f'finished done={done} failed={failed} raw={RAW}')\n\n\ndef normalize_text(text: str) -> str:\n text = unicodedata.normalize('NFC', text)\n text = text.replace('\\x00', '')\n text = re.sub(r'\\r\\n?', '\\n', text)\n text = text.replace('\\t', ' ')\n text = re.sub(r'[ \\t]+\\n', '\\n', text)\n text = re.sub(r'\\n{4,}', '\\n\\n\\n', text)\n return text.strip()\n\n\ndef repeated_line_ratio(text: str) -> float:\n lines = [x.strip() for x in text.splitlines() if x.strip()]\n if len(lines) < 6:\n return 0.0\n return 1.0 - (len(set(lines)) / len(lines))\n\n\ndef clean(args: argparse.Namespace) -> None:\n if not RAW.exists():\n raise SystemExit(f'不存在 {RAW}')\n seen: set[str] = set()\n kept = 0\n rejected = 0\n chars = 0\n usage_in = 0\n usage_out = 0\n refusal_markers = ('作为一个ai', '作为 ai', 'i cannot comply', 'i can\\'t comply', '申し訳ありませんが')\n with CLEAN.open('w', encoding='utf-8') as clean_file, CORPUS.open('w', encoding='utf-8') as corpus_file:\n for item in iter_jsonl(RAW) or []:\n text = normalize_text(str(item.get('output') or ''))\n reason = None\n if item.get('finish_reason') != 'stop':\n reason = 'incomplete'\n elif len(text) < args.min_chars:\n reason = 'too_short'\n elif len(text) > args.max_chars:\n reason = 'too_long'\n elif any(marker in text[:300].lower() for marker in refusal_markers):\n reason = 'refusal'\n elif repeated_line_ratio(text) > args.max_repeated_line_ratio:\n reason = 'repetition'\n digest = hashlib.sha256(re.sub(r'\\s+', ' ', text).encode('utf-8')).hexdigest()\n if digest in seen:\n reason = 'duplicate'\n if reason:\n rejected += 1\n continue\n seen.add(digest)\n output = {\n 'id': item.get('id'),\n 'category': item.get('category'),\n 'topic': item.get('topic'),\n 'language': item.get('language'),\n 'prompt': item.get('prompt'),\n 'output': text,\n 'sha256': digest,\n 'model': item.get('model'),\n 'usage': item.get('usage', {}),\n }\n clean_file.write(json.dumps(output, ensure_ascii=False) + '\\n')\n corpus_file.write(text + '\\n\\n')\n kept += 1\n chars += len(text)\n usage = item.get('usage') or {}\n usage_in += int(usage.get('prompt_tokens') or 0)\n usage_out += int(usage.get('completion_tokens') or 0)\n stats = {\n 'kept_documents': kept,\n 'rejected_documents': rejected,\n 'characters': chars,\n 'api_prompt_tokens': usage_in,\n 'api_completion_tokens': usage_out,\n 'clean_jsonl': str(CLEAN),\n 'corpus_txt': str(CORPUS),\n }\n STATS.write_text(json.dumps(stats, ensure_ascii=False, indent=2) + '\\n', encoding='utf-8')\n print(json.dumps(stats, ensure_ascii=False, indent=2))\n\n\ndef find_special_id(tokenizer, names: list[str]) -> int | None:\n for name in names:\n token_id = tokenizer.token_to_id(name)\n if token_id is not None:\n return int(token_id)\n return None\n\n\ndef tokenize(args: argparse.Namespace) -> None:\n try:\n import numpy as np\n from tokenizers import Tokenizer\n except ImportError as exc:\n raise SystemExit('缺少 numpy/tokenizers') from exc\n tokenizer_path = Path(args.tokenizer)\n if not tokenizer_path.exists():\n raise SystemExit(f'不存在 tokenizer: {tokenizer_path}')\n if not CLEAN.exists():\n raise SystemExit(f'不存在 {CLEAN}，先运行 clean')\n tokenizer = Tokenizer.from_file(str(tokenizer_path))\n bos_id = find_special_id(tokenizer, ['<bos>', '<s>', '[BOS]'])\n eos_id = find_special_id(tokenizer, ['<eos>', '</s>', '[EOS]'])\n unk_id = tokenizer.token_to_id('<unk>')\n total_tokens = 0\n documents = 0\n with BIN.open('wb') as out:\n for item in iter_jsonl(CLEAN) or []:\n text = str(item['output'])\n ids = tokenizer.encode(text, add_special_tokens=False).ids\n if unk_id is not None and unk_id in ids:\n raise RuntimeError(f\"document {item.get('id')} contains <unk>\")\n if bos_id is not None:\n ids.insert(0, bos_id)\n if eos_id is not None:\n ids.append(eos_id)\n if not ids:\n continue\n np.asarray(ids, dtype=np.uint32).tofile(out)\n total_tokens += len(ids)\n documents += 1\n stats = {}\n if STATS.exists():\n try:\n stats = json.loads(STATS.read_text(encoding='utf-8'))\n except json.JSONDecodeError:\n stats = {}\n stats.update({\n 'tokenized_documents': documents,\n 'albertlm_tokens': total_tokens,\n 'binary_path': str(BIN),\n 'binary_bytes': BIN.stat().st_size,\n 'tokenizer_path': str(tokenizer_path),\n 'bos_id': bos_id,\n 'eos_id': eos_id,\n })\n STATS.write_text(json.dumps(stats, ensure_ascii=False, indent=2) + '\\n', encoding='utf-8')\n print(json.dumps(stats, ensure_ascii=False, indent=2))\n\n\ndef status(_: argparse.Namespace) -> None:\n def count(path: Path) -> int:\n return sum(1 for _ in iter_jsonl(path) or []) if path.exists() else 0\n print(f'jobs={count(JOBS)} raw={count(RAW)} errors={count(ERRORS)} clean={count(CLEAN)}')\n for path in [JOBS, RAW, ERRORS, CLEAN, CORPUS, BIN, STATS]:\n if path.exists():\n print(f'{path}: {path.stat().st_size / 1024 / 1024:.2f} MiB')\n if STATS.exists():\n print(STATS.read_text(encoding='utf-8'))\n\n\ndef main() -> int:\n parser = argparse.ArgumentParser()\n sub = parser.add_subparsers(dest='command', required=True)\n\n p = sub.add_parser('make-jobs')\n p.add_argument('--count', type=int, default=10000)\n p.add_argument('--seed', type=int, default=20260715)\n p.set_defaults(func=make_jobs)\n\n p = sub.add_parser('generate')\n p.add_argument('--model', default='deepseek-v4-pro')\n p.add_argument('--concurrency', type=int, default=32)\n p.add_argument('--max-tokens', type=int, default=1800)\n p.add_argument('--temperature', type=float, default=0.7)\n p.add_argument('--timeout', type=float, default=900.0)\n p.add_argument('--retries', type=int, default=8)\n p.add_argument('--limit', type=int, default=0, help='0 means all pending jobs')\n p.set_defaults(func=lambda args: asyncio.run(generate(args)))\n\n p = sub.add_parser('clean')\n p.add_argument('--min-chars', type=int, default=500)\n p.add_argument('--max-chars', type=int, default=30000)\n p.add_argument('--max-repeated-line-ratio', type=float, default=0.35)\n p.set_defaults(func=clean)\n\n p = sub.add_parser('tokenize')\n p.add_argument('--tokenizer', default=str(ROOT / 'tokenizer' / 'tokenizer.json'))\n p.set_defaults(func=tokenize)\n\n p = sub.add_parser('status')\n p.set_defaults(func=status)\n\n args = parser.parse_args()\n args.func(args)\n return 0\n\n\nif __name__ == '__main__':\n raise SystemExit(main())\n"}}}} |