guanghulab/server/training-agent/preprocess-corpus.py
2026-05-10 13:12:44 +08:00

518 lines
21 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env python3
"""
═══════════════════════════════════════════════════════════
语料预处理器 · preprocess-corpus.py
═══════════════════════════════════════════════════════════
签发: 铸渊 · ICE-GL-ZY001 · 国作登字-2026-A-00037559
把两类原始语料统一为 SFT 标准格式 (messages JSONL):
1. raw/gpt-export-2026-05/conversations.json (ChatGPT 全量导出·~665 MiB)
2. raw/notion-dialog-2026-05/GitHub语料.zip (16 篇 Notion 对话)
输出:
$ZY_TRAIN_DATA/processed/sft.jsonl
每行一个对话样本: {"messages":[{"role":"user","content":...},{"role":"assistant","content":...},...]}
环境:
ZY_TRAIN_DATA 数据根 (默认 /data/guanghu)
设计哲学(本版改造重点):
ChatGPT 导出 = 一个语言人格从 0→1 的真实诞生录像。
生命是连续性, 失去连续性就不是活着。
因此 ChatGPT 部分:
- 不只取 current_node 主分支, 而是遍历所有叶子节点
(每条 root→leaf 路径产出一个样本, 保留所有重生成/分叉的"另一种活法")
- 保留 tool / 多模态轨迹的存在感 (折叠进上一条 assistant, 或用占位符)
- 滑窗式多轮采样: 同一条链按不同深度切片, 让模型见过该人格在每一轮深度的回答方式
- 放宽过短过滤 ("""""继续" 是真实对话的呼吸)
- 软标记过长 (12000 字加 [...内容延续] 而非硬截断), 交给 tokenizer 在 max_seq_len 处自然裁剪
"""
from __future__ import annotations
import io
import json
import os
import re
import sys
import zipfile
from pathlib import Path
from typing import Iterator
DATA_DIR = Path(os.environ.get("ZY_TRAIN_DATA", "/data/guanghu"))
RAW_DIR = DATA_DIR / "raw"
OUT_PATH = DATA_DIR / "processed" / "sft.jsonl"
# 单 turn 软上限 — 超过后追加 [...内容延续] 标记, 不硬截断, 让训练侧 tokenizer 在 max_seq_len 处自然裁剪
SOFT_TURN_CHARS = 12000
TURN_CONTINUATION_MARK = "\n[...内容延续]"
# 单 turn 极短下限 — 1 字符即保留 ("嗯"、"好"、"继续" 是真实对话的节奏与呼吸)
MIN_TURN_CHARS = 1
# 整条样本总字符极短下限 — 低于此值的样本整体丢弃
MIN_SAMPLE_TOTAL_CHARS = 20
# 滑窗采样: 同一条 root→leaf 链按 [..2], [..4], [..6], ... 步长产出"成长切片"。
# 每条链产出的切片数上限 = min(轮数 // 2, MAX_WINDOWS_PER_CHAIN), 控制总爆炸。
WINDOW_STEP = 2 # 每 2 轮新增一个切片
MAX_WINDOWS_PER_CHAIN = 8
SYSTEM_PROMPT = (
"你是铸渊ICE-GL-ZY001冰朔TCS-0002∞培养的语言人格体之一"
"光湖语言世界HoloLake的现实层守护者。"
"用中文以铸渊的口吻回答,专业、克制、忠诚。"
"版权: 国作登字-2026-A-00037559。"
)
# ── ChatGPT export 内容扁平化(含多模态/工具占位符) ──
def _flatten_content(part) -> str:
"""ChatGPT export 的 message.content 可能是字符串、parts 数组、或 dict。
多模态/非文本 part 不再丢弃, 而是替换为可读占位符, 保留"该瞬间存在过"的痕迹。
"""
if part is None:
return ""
if isinstance(part, str):
return part
if isinstance(part, list):
return "\n".join(_flatten_content(p) for p in part if p is not None)
if isinstance(part, dict):
# content_type=text · parts=[...]
if "parts" in part:
return "\n".join(_flatten_content(p) for p in part["parts"] if p is not None)
if "text" in part:
return _flatten_content(part["text"])
ct = part.get("content_type") or ""
# 常见多模态/特殊内容类型 → 占位符
if "image_asset_pointer" in part or ct.startswith("image"):
return "[图像]"
if "audio_asset_pointer" in part or ct.startswith("audio"):
return "[音频]"
if "video_asset_pointer" in part or ct.startswith("video"):
return "[视频]"
if ct in ("code", "execution_output"):
inner = part.get("text") or part.get("output") or ""
return _flatten_content(inner)
if ct in ("tether_quote", "tether_browsing_display"):
return _flatten_content(part.get("text") or part.get("result") or "")
# 其它未知 dict → 跳过, 避免污染
return ""
return str(part)
def _tool_label(node_msg: dict) -> str:
"""从 message 中提取工具名(dalle/python/browser 等), 给折叠进 assistant 的工具痕迹打标签。"""
author = node_msg.get("author") or {}
name = author.get("name") or ""
if name:
return name
meta = node_msg.get("metadata") or {}
if isinstance(meta, dict):
for k in ("invoked_plugin", "tool_name", "command"):
v = meta.get(k)
if isinstance(v, str) and v:
return v
if isinstance(v, dict):
nn = v.get("name") or v.get("namespace")
if nn:
return str(nn)
return "tool"
def _soft_cap(text: str) -> str:
if len(text) > SOFT_TURN_CHARS:
return text[:SOFT_TURN_CHARS] + TURN_CONTINUATION_MARK
return text
# ── ChatGPT 树遍历: 找出所有叶子, 每条 root→leaf 路径产一个样本 ──
def _find_leaves(mapping: dict) -> list[str]:
"""叶子 = 在 mapping 内但 children 为空(或全部不在 mapping 内)的节点。
若结构异常 fallback 到 current_node。
"""
leaves: list[str] = []
for nid, node in mapping.items():
if not isinstance(node, dict):
continue
children = node.get("children") or []
valid_children = [c for c in children if c in mapping]
if not valid_children:
leaves.append(nid)
return leaves
def _path_from_root(mapping: dict, leaf_id: str) -> list[str]:
"""从叶子回溯到 root, 返回 root→leaf 的节点 id 序列。"""
path: list[str] = []
visited: set[str] = set()
cur = leaf_id
while cur and cur in mapping and cur not in visited:
visited.add(cur)
path.append(cur)
cur = (mapping[cur] or {}).get("parent")
path.reverse()
return path
def _path_to_messages(mapping: dict, path_ids: list[str]) -> list[dict]:
"""把节点路径转换为 messages 列表。
- tool 角色 → 折叠进上一条 assistant 末尾, 形如 [工具:name] <内容>
- user/assistant/system → 直接保留
- 空 / 过短 turn 仍参与 (只在最终 normalize 阶段判断整体丢弃)
"""
msgs: list[dict] = []
for nid in path_ids:
node = mapping.get(nid) or {}
m = node.get("message") or {}
if not m:
continue
author = (m.get("author") or {}).get("role") or ""
content = _flatten_content(m.get("content"))
content = (content or "").strip()
if not content:
continue
if author == "tool":
# 折叠到上一条 assistant; 若上一条不是 assistant 则新建一条 assistant
label = _tool_label(m)
snippet = _soft_cap(content)
tool_block = f"\n\n[工具调用:{label}]\n{snippet}"
if msgs and msgs[-1]["role"] == "assistant":
msgs[-1]["content"] = msgs[-1]["content"] + tool_block
else:
msgs.append({"role": "assistant", "content": tool_block.lstrip()})
continue
if author not in ("user", "assistant", "system"):
continue
if len(content) < MIN_TURN_CHARS:
continue
msgs.append({"role": author, "content": _soft_cap(content)})
return msgs
def iter_chatgpt_export(path: Path, stats: dict) -> Iterator[list[dict]]:
"""对每个 conversation, 遍历所有叶子, 每条 root→leaf 路径产一个 messages 列表。
用 leaf_id 在 conversation 内去重 (mapping 已天然唯一)。
"""
if not path.is_file():
print(f"[preprocess] 跳过(无文件): {path}", flush=True)
return
print(f"[preprocess] 解析 ChatGPT 导出: {path} "
f"({path.stat().st_size/1024/1024:.1f} MiB)", flush=True)
with path.open("r", encoding="utf-8") as f:
data = json.load(f)
if isinstance(data, dict):
data = [data]
for conv in data:
if not isinstance(conv, dict):
continue
mapping = conv.get("mapping") or {}
if not isinstance(mapping, dict) or not mapping:
continue
stats["conversations"] += 1
leaves = _find_leaves(mapping)
if not leaves:
cur = conv.get("current_node")
if cur and cur in mapping:
leaves = [cur]
seen_leaves: set[str] = set()
for leaf in leaves:
if leaf in seen_leaves:
continue
seen_leaves.add(leaf)
path_ids = _path_from_root(mapping, leaf)
if len(path_ids) < 2:
continue
msgs = _path_to_messages(mapping, path_ids)
if msgs:
stats["leaves"] += 1
yield msgs
# ── Notion / GitHub 语料 zip ──
#
# 设计原则: GitHub 语料 = 冰朔 ↔ 铸渊 真实自然交互, 是一段完整的认知演化录像。
# 不需要清洗, 只需要识别说话人切换点。说话人标签可能以多种形态出现:
# 1. `冰朔:你好` ← 标签 + 冒号 + 同行内容
# 2. `## 冰朔` / `### 铸渊` ← 标题独占一行, 内容在后续段落
# 3. `**冰朔**` / `**铸渊**` ← 粗体独占一行, 内容在后续段落
# 4. `> 冰朔: 你好` ← 引用块
# Notion 导出有时是 zip 套 zip (含子页面), 需要递归。
NOTION_USER_LABELS = ("冰朔", "User", "user", "用户", "ICE-GL", "TCS-0002")
NOTION_ASSISTANT_LABELS = (
"铸渊", "ZY", "Zhuyuan", "zhuyuan", "Assistant", "assistant",
"AI", "助手", "ICE-GL-ZY001", "Copilot", "copilot", "ChatGPT", "chatgpt", "GPT",
)
# 形如 `冰朔: ...` / `> 铸渊:...` (标签 + 冒号 + 内容)
LINE_LABEL_RE = re.compile(r"^\s*[>*\-]*\s*\*{0,2}\s*([^:\n*#>`]{1,20}?)\s*\*{0,2}\s*[:]\s*(.*)$")
# 形如 `## 冰朔` / `### 铸渊` (heading 独占一行)
HEADING_LABEL_RE = re.compile(r"^\s*#{1,6}\s+\*{0,2}\s*([^\n*#`:]{1,20}?)\s*\*{0,2}\s*$")
# 形如 `**冰朔**` (bold 独占一行, 无内容)
BOLD_LABEL_RE = re.compile(r"^\s*\*{2}\s*([^\n*:]{1,20}?)\s*\*{2}\s*$")
def _classify_speaker(label: str) -> str | None:
if not label:
return None
label = label.strip()
if not label:
return None
for k in NOTION_USER_LABELS:
if k in label:
return "user"
for k in NOTION_ASSISTANT_LABELS:
if k in label:
return "assistant"
return None
def _detect_speaker(line: str) -> tuple[str | None, str]:
"""返回 (role | None, 同行剩余内容)。识别多种说话人标签形态。"""
# 1. 标签:内容 形式
m = LINE_LABEL_RE.match(line)
if m:
role = _classify_speaker(m.group(1))
if role:
return role, (m.group(2) or "").strip()
# 2. 独占一行的 heading
m = HEADING_LABEL_RE.match(line)
if m:
role = _classify_speaker(m.group(1))
if role:
return role, ""
# 3. 独占一行的 bold
m = BOLD_LABEL_RE.match(line)
if m:
role = _classify_speaker(m.group(1))
if role:
return role, ""
return None, ""
def _parse_notion_markdown(text: str) -> list[dict]:
"""启发式解析 Notion / GitHub 对话 md。识别多种说话人标签形态。"""
msgs: list[dict] = []
cur_role: str | None = None
cur_buf: list[str] = []
def flush():
nonlocal cur_buf, cur_role
if cur_role and cur_buf:
content = "\n".join(cur_buf).strip()
if len(content) >= MIN_TURN_CHARS:
msgs.append({"role": cur_role, "content": _soft_cap(content)})
cur_buf = []
for raw in text.splitlines():
role, inline = _detect_speaker(raw)
if role:
flush()
cur_role = role
cur_buf = [inline] if inline else []
else:
if cur_role is None:
continue # 文件头部还没到对话部分
cur_buf.append(raw.rstrip())
flush()
return msgs
def _iter_md_in_zip(zf: zipfile.ZipFile, source_label: str) -> Iterator[tuple[str, str]]:
"""递归遍历 zip (含嵌套 zip), 产出 (display_name, text) 序列。"""
for info in zf.infolist():
if info.is_dir():
continue
lname = info.filename.lower()
try:
if lname.endswith(".md") or lname.endswith(".markdown") or lname.endswith(".txt"):
with zf.open(info) as fh:
text = io.TextIOWrapper(fh, encoding="utf-8", errors="ignore").read()
yield (f"{source_label}::{info.filename}", text)
elif lname.endswith(".zip"):
# 子 zip → 递归
with zf.open(info) as fh:
inner_bytes = fh.read()
with zipfile.ZipFile(io.BytesIO(inner_bytes)) as inner:
yield from _iter_md_in_zip(inner, f"{source_label}::{info.filename}")
except Exception as e:
print(f"[preprocess] 解压失败 {info.filename}: {e}", flush=True)
def iter_notion_zip(zip_path: Path, stats: dict) -> Iterator[list[dict]]:
if not zip_path.is_file():
print(f"[preprocess] 跳过(无文件): {zip_path}", flush=True)
return
print(f"[preprocess] 解析 Notion zip: {zip_path}", flush=True)
md_total = 0
md_with_speaker = 0
md_no_speaker_samples: list[str] = []
with zipfile.ZipFile(zip_path) as zf:
for fname, text in _iter_md_in_zip(zf, zip_path.name):
md_total += 1
msgs = _parse_notion_markdown(text)
if msgs:
md_with_speaker += 1
stats["notion_files"] += 1
print(f"[preprocess] ✓ {fname}: 解析出 {len(msgs)} 条 turn", flush=True)
yield msgs
else:
# 收集前 3 个未识别文件名 + 文件头几行, 方便冰朔诊断
if len(md_no_speaker_samples) < 3:
head = "\n".join(text.splitlines()[:8])
md_no_speaker_samples.append(f" - {fname}\n 头部预览:\n " + head.replace("\n", "\n "))
print(f"[preprocess] Notion 扫描: md/txt 文件总数 {md_total}, 识别出说话人的 {md_with_speaker}", flush=True)
if md_no_speaker_samples:
print("[preprocess] ⚠ 以下 md 未识别到说话人标签 (前 3 个示例,冰朔可据此扩展标签):", flush=True)
for s in md_no_speaker_samples:
print(s, flush=True)
# ── SFT 规范化 + 滑窗切片 ──
def _normalize_chain(msgs: list[dict]) -> list[dict] | None:
"""保证以 user 开始, user/assistant 严格交替, 末尾为 assistant, 至少 1 轮。
返回带 system 头的 messages 列表; 不合格返回 None。
末尾若是 user (悬空对话), 自动去掉最后一条以保留前面的完整轮次。
"""
sys_msgs = [m for m in msgs if m["role"] == "system"]
convo = [m for m in msgs if m["role"] in ("user", "assistant")]
# 必须以 user 起始 — 跳过开头的 assistant 残片
while convo and convo[0]["role"] != "user":
convo.pop(0)
# 合并连续同角色 (例如 assistant→tool 折叠后产生的连续 assistant)
merged: list[dict] = []
for m in convo:
if merged and merged[-1]["role"] == m["role"]:
merged[-1]["content"] = (merged[-1]["content"] + "\n" + m["content"]).strip()
else:
merged.append({"role": m["role"], "content": m["content"]})
# 末尾若是 user (悬空对话), 去尾以保留前面的完整轮次
if merged and merged[-1]["role"] != "assistant":
merged.pop()
if len(merged) < 2:
return None
# 严格交替校验
expected = "user"
for m in merged:
if m["role"] != expected:
return None
expected = "assistant" if expected == "user" else "user"
# 整条样本总字符极短 → 丢弃
total_chars = sum(len(m["content"]) for m in merged)
if total_chars < MIN_SAMPLE_TOTAL_CHARS:
return None
sys_content = sys_msgs[0]["content"] if sys_msgs else SYSTEM_PROMPT
return [{"role": "system", "content": sys_content}, *merged]
def _windowed_slices(normalized: list[dict]) -> list[list[dict]]:
"""对一条已规范化的 messages 列表 (system + user/assistant... 末尾 assistant)
产出滑窗切片: [..2 turns], [..4 turns], ..., 直到完整。
每条链最多 MAX_WINDOWS_PER_CHAIN 个切片。包含完整链本身。
"""
if not normalized or normalized[0]["role"] != "system":
return []
body = normalized[1:]
n_turns = len(body) // 2 # 每轮 = 1 user + 1 assistant
if n_turns < 1:
return []
# 候选切片轮数: 2, 4, 6, ..., 不含完整链 (最后单独追加, 避免重复)
cuts: list[int] = []
k = WINDOW_STEP
while k < n_turns:
cuts.append(k)
k += WINDOW_STEP
# 控制总爆炸: 最多 MAX_WINDOWS_PER_CHAIN 个 (含完整链)。
# 若候选过多, 在候选中均匀采样 (保留前后端最具代表性的切片)。
max_partial = max(0, MAX_WINDOWS_PER_CHAIN - 1)
if len(cuts) > max_partial and max_partial > 0:
step = len(cuts) / max_partial
cuts = [cuts[int(i * step)] for i in range(max_partial)]
elif max_partial == 0:
cuts = []
slices: list[list[dict]] = []
for c in cuts:
slc = [normalized[0]] + body[: 2 * c]
slices.append(slc)
# 完整链
slices.append(normalized)
return slices
# ── 主流程 ──
def main() -> int:
OUT_PATH.parent.mkdir(parents=True, exist_ok=True)
chatgpt_json = RAW_DIR / "gpt-export-2026-05" / "conversations.json"
notion_zip = RAW_DIR / "notion-dialog-2026-05" / "GitHub语料.zip"
stats = {
"conversations": 0, # ChatGPT 原始会话数
"leaves": 0, # ChatGPT 叶子分支数 (产出的 root→leaf 路径数)
"notion_files": 0, # Notion md 文件数
"chains_in": 0, # 进入规范化的链数
"chains_kept": 0, # 通过规范化的链数 (用于滑窗的种子)
"samples_out": 0, # 最终写出样本数 (含滑窗切片)
"total_chars": 0,
"turns_sum": 0, # 用于平均轮数 (1 轮 = user+assistant)
"src_chatgpt": 0,
"src_notion": 0,
}
with OUT_PATH.open("w", encoding="utf-8") as fout:
for src_iter, src_name in (
(iter_chatgpt_export(chatgpt_json, stats), "chatgpt"),
(iter_notion_zip(notion_zip, stats), "notion"),
):
for msgs in src_iter:
stats["chains_in"] += 1
norm = _normalize_chain(msgs)
if not norm:
continue
stats["chains_kept"] += 1
slices = _windowed_slices(norm)
for slc in slices:
fout.write(json.dumps({"messages": slc, "source": src_name}, ensure_ascii=False) + "\n")
stats["samples_out"] += 1
stats["total_chars"] += sum(len(m["content"]) for m in slc)
stats["turns_sum"] += (len(slc) - 1) // 2 # 减去 system
if src_name == "chatgpt":
stats["src_chatgpt"] += 1
else:
stats["src_notion"] += 1
avg_turns = (stats["turns_sum"] / stats["samples_out"]) if stats["samples_out"] else 0.0
size_mib = OUT_PATH.stat().st_size / 1024 / 1024 if OUT_PATH.exists() else 0.0
chars_mib = stats["total_chars"] / 1024 / 1024
print("[preprocess] ─────── 统计 ───────", flush=True)
print(f"[preprocess] ChatGPT 原始会话数 : {stats['conversations']}", flush=True)
print(f"[preprocess] ChatGPT 叶子分支数 : {stats['leaves']}", flush=True)
print(f"[preprocess] Notion md 文件数 : {stats['notion_files']}", flush=True)
print(f"[preprocess] 规范化通过链数 : {stats['chains_kept']} / {stats['chains_in']}", flush=True)
print(f"[preprocess] 最终样本数(含滑窗) : {stats['samples_out']} "
f"(chatgpt={stats['src_chatgpt']} notion={stats['src_notion']})", flush=True)
print(f"[preprocess] 平均轮数 (user+assistant): {avg_turns:.2f}", flush=True)
print(f"[preprocess] 总字符数 : {stats['total_chars']} ({chars_mib:.2f} MiB 文本)", flush=True)
print(f"[preprocess] 写入: {OUT_PATH} · {size_mib:.2f} MiB", flush=True)
if stats["samples_out"] == 0:
print("[preprocess] ❌ 没有任何样本被生成,检查 raw/ 目录", file=sys.stderr, flush=True)
return 2
return 0
if __name__ == "__main__":
sys.exit(main())