File size: 4,723 Bytes
31f25cb | 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | import json
import os
import glob
from tqdm import tqdm
from huggingface_hub import snapshot_download
# ==========================================
# 1. 配置
# ==========================================
DATASET_ID = "stepfun-ai/Step-3.5-Flash-SFT"
OUTPUT_DIR = "shell/playground/data/hf_data_0330"
os.makedirs(OUTPUT_DIR, exist_ok=True)
OUTPUT_PATH = os.path.join(OUTPUT_DIR, "Step-3.5-Flash-SFT_train.jsonl")
ROLE_MAP = {
"user": "human",
"human": "human",
"assistant": "gpt",
"gpt": "gpt",
"system": "system",
}
def main():
print(f"📦 正在同步 '{DATASET_ID}' 的所有原始文件 (绕过 Arrow 解析) ...")
print(f"⏳ 会自动下载所有分片到缓存，请耐心等待...")
try:
# 1. 直接下载仓库中的所有 JSON 数据文件，不走 Schema 校验
local_dir = snapshot_download(
repo_id=DATASET_ID,
repo_type="dataset",
allow_patterns=["*.json", "*.jsonl"]
)
except Exception as e:
print(f"❌ 下载失败: {e}")
return
# 2. 找到所有下载下来的 json/jsonl 文件
data_files = glob.glob(os.path.join(local_dir, "**", "*.json"), recursive=True) + \
glob.glob(os.path.join(local_dir, "**", "*.jsonl"), recursive=True)
# 过滤掉明显的非数据文件（比如 huggingface 自动生成的元数据）
data_files = [f for f in data_files if "dataset_info" not in f.lower()]
print(f"✅ 成功定位到 {len(data_files)} 个数据文件。开始逐个文件解析...")
success_count = 0
total_records = 0
with open(OUTPUT_PATH, 'w', encoding='utf-8') as f_out:
# 以文件为单位进行进度条显示
for file_path in tqdm(data_files, desc="Processing Files"):
try:
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read().strip()
if not content:
continue
# 兼容处理：文件内容可能是完整的 JSON Array，也可能是 JSONL
if content.startswith('['):
records = json.loads(content)
else:
records = [json.loads(line) for line in content.split('\n') if line.strip()]
except Exception as e:
print(f"\n⚠️ 跳过文件 {os.path.basename(file_path)}: 无法解析 JSON ({e})")
continue
# 3. 处理单文件内的所有记录
for entry in records:
total_records += 1
msgs = entry.get("conversations") or entry.get("messages")
# 如果这个脏数据的 conversations 不是 list (正是由于这个导致 pyarrow 崩溃)，直接跳过
if not isinstance(msgs, list):
continue
new_convs = []
is_agent = False
for m in msgs:
if not isinstance(m, dict): continue
if m.get("tool_calls") or m.get("name") or m.get("tool_call_id"):
is_agent = True
break
role = m.get("role", "")
if role:
role = role.lower()
content = m.get("content", "")
if isinstance(content, list):
content = " ".join([c.get("text", "") for c in content if isinstance(c, dict) and "text" in c])
if role in ROLE_MAP and isinstance(content, str) and content.strip():
new_convs.append({
"from": ROLE_MAP[role],
"value": content
})
if is_agent or not new_convs:
continue
# 4. 写入输出文件
output_item = {
"id": f"step_{total_records}",
"conversations": new_convs
}
if 'source' in entry:
output_item['source'] = entry['source']
f_out.write(json.dumps(output_item, ensure_ascii=False) + '\n')
success_count += 1
print(f"\n✨ 处理完成！")
print(f"📊 共扫描原始记录: {total_records} 条")
print(f"🎯 成功提取可用记录: {success_count} 条")
print(f"📂 结果保存至: {OUTPUT_PATH}")
if __name__ == "__main__":
main() |