174 lines
7.1 KiB
Python
Executable File
174 lines
7.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import argparse
|
|
import gzip
|
|
import json
|
|
import os
|
|
import re
|
|
import time
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
|
|
from huggingface_hub import HfApi, hf_hub_download
|
|
from transformers import AutoTokenizer
|
|
|
|
|
|
def clean_text(text):
|
|
if text is None:
|
|
return ""
|
|
text = str(text).replace("\x00", " ")
|
|
text = re.sub(r"[ \t\r\f\v]+", " ", text)
|
|
text = re.sub(r"\n{4,}", "\n\n\n", text)
|
|
return text.strip()
|
|
|
|
|
|
def iter_jsonl(path):
|
|
with open(path, "r", encoding="utf-8", errors="replace") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
yield json.loads(line)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
|
|
def download_with_retry(filename, args):
|
|
last = None
|
|
for attempt in range(1, args.retries + 1):
|
|
try:
|
|
return hf_hub_download(
|
|
repo_id=args.repo,
|
|
repo_type="dataset",
|
|
filename=filename,
|
|
endpoint=args.endpoint,
|
|
token=os.environ.get("HF_TOKEN"),
|
|
local_dir=args.raw_dir,
|
|
)
|
|
except Exception as exc:
|
|
last = exc
|
|
print(json.dumps({"event": "download_retry", "file": filename, "attempt": attempt, "error": repr(exc)[:800]}, ensure_ascii=False), flush=True)
|
|
time.sleep(min(120, 5 * attempt))
|
|
raise RuntimeError(f"download failed for {filename}: {last!r}")
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--repo", default="BAAI/CCI3-HQ")
|
|
ap.add_argument("--endpoint", default=os.environ.get("HF_ENDPOINT", "https://hf-mirror.com"))
|
|
ap.add_argument("--tokenizer", default="model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2")
|
|
ap.add_argument("--raw-dir", default="data/raw_jsonl/cci3_hq_probe")
|
|
ap.add_argument("--out-dir", default="data/cpt_docmix_5b_sources_8192_20260614")
|
|
ap.add_argument("--target-tokens", type=int, default=1_250_000_000)
|
|
ap.add_argument("--min-tokens", type=int, default=128)
|
|
ap.add_argument("--max-doc-tokens", type=int, default=32768)
|
|
ap.add_argument("--log-every", type=int, default=5000)
|
|
ap.add_argument("--retries", type=int, default=16)
|
|
ap.add_argument("--start-index", type=int, default=0)
|
|
ap.add_argument("--keep-raw", action="store_true")
|
|
args = ap.parse_args()
|
|
|
|
base = Path.cwd()
|
|
raw_dir = base / args.raw_dir
|
|
out_dir = base / args.out_dir
|
|
doc_dir = out_dir / "documents"
|
|
doc_dir.mkdir(parents=True, exist_ok=True)
|
|
raw_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
token = os.environ.get("HF_TOKEN")
|
|
if not token:
|
|
raise SystemExit("HF_TOKEN is required for gated CCI3-HQ")
|
|
|
|
api = HfApi(endpoint=args.endpoint, token=token)
|
|
files = [f for f in api.list_repo_files(args.repo, repo_type="dataset") if f.startswith("data/part_") and f.endswith(".jsonl")]
|
|
files = sorted(files)
|
|
if args.start_index:
|
|
files = [f for f in files if int(Path(f).stem.split("_")[-1]) >= args.start_index]
|
|
|
|
tok = AutoTokenizer.from_pretrained(base / args.tokenizer, trust_remote_code=True)
|
|
tmp_out = doc_dir / "chinese_clean.jsonl.gz.tmp"
|
|
final_out = doc_dir / "chinese_clean.jsonl.gz"
|
|
if tmp_out.exists():
|
|
tmp_out.unlink()
|
|
|
|
stats = {
|
|
"target_tokens": args.target_tokens,
|
|
"tokens": 0,
|
|
"docs": 0,
|
|
"rows_seen": 0,
|
|
"files_done": [],
|
|
"tokens_by_file": Counter(),
|
|
"rejected": Counter(),
|
|
"started_at": time.time(),
|
|
"repo": args.repo,
|
|
"endpoint": args.endpoint,
|
|
}
|
|
|
|
with gzip.open(tmp_out, "wt", encoding="utf-8") as w:
|
|
for filename in files:
|
|
if stats["tokens"] >= args.target_tokens:
|
|
break
|
|
local_path = Path(download_with_retry(filename, args))
|
|
file_tokens = 0
|
|
file_docs = 0
|
|
for row in iter_jsonl(local_path):
|
|
stats["rows_seen"] += 1
|
|
if stats["tokens"] >= args.target_tokens:
|
|
break
|
|
text = clean_text(row.get("text") or row.get("content"))
|
|
if not text:
|
|
stats["rejected"]["empty"] += 1
|
|
continue
|
|
ntok = len(tok.encode(text, add_special_tokens=False))
|
|
if ntok < args.min_tokens:
|
|
stats["rejected"]["too_short"] += 1
|
|
continue
|
|
if ntok > args.max_doc_tokens:
|
|
stats["rejected"]["too_long"] += 1
|
|
continue
|
|
idx = stats["docs"]
|
|
rec = {
|
|
"id": f"chinese_clean_cci3_hq_{idx:09d}",
|
|
"category": "chinese_clean",
|
|
"source": f"{args.repo}:{filename}",
|
|
"text": text,
|
|
"token_count": ntok,
|
|
"metadata": {"cci3_id": row.get("id"), "cci3_score": row.get("score")},
|
|
}
|
|
w.write(json.dumps(rec, ensure_ascii=False) + "\n")
|
|
stats["docs"] += 1
|
|
stats["tokens"] += ntok
|
|
file_docs += 1
|
|
file_tokens += ntok
|
|
if stats["docs"] % args.log_every == 0:
|
|
print(json.dumps({"event": "progress", "docs": stats["docs"], "tokens": stats["tokens"], "target": args.target_tokens, "file": filename, "elapsed_sec": time.time() - stats["started_at"]}, ensure_ascii=False), flush=True)
|
|
stats["files_done"].append({"file": filename, "docs": file_docs, "tokens": file_tokens, "path": str(local_path)})
|
|
stats["tokens_by_file"][filename] += file_tokens
|
|
print(json.dumps({"event": "file_done", "file": filename, "docs": file_docs, "tokens": file_tokens, "total_tokens": stats["tokens"]}, ensure_ascii=False), flush=True)
|
|
if not args.keep_raw:
|
|
try:
|
|
local_path.unlink()
|
|
except Exception:
|
|
pass
|
|
|
|
if stats["tokens"] < args.target_tokens:
|
|
raise SystemExit(f"only collected {stats['tokens']} / {args.target_tokens} tokens")
|
|
|
|
backup = final_out.with_suffix(".jsonl.gz.failed_empty_20260614")
|
|
if final_out.exists() and final_out.stat().st_size < 1024:
|
|
final_out.replace(backup)
|
|
elif final_out.exists():
|
|
final_out.replace(final_out.with_suffix(".jsonl.gz.backup_20260614"))
|
|
tmp_out.replace(final_out)
|
|
|
|
stats["elapsed_sec"] = time.time() - stats["started_at"]
|
|
stats["tokens_by_file"] = dict(stats["tokens_by_file"])
|
|
stats["rejected"] = dict(stats["rejected"])
|
|
(out_dir / "chinese_clean_5b_fix_stats.json").write_text(json.dumps(stats, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
(out_dir / ".chinese_clean_5b_ready").write_text(json.dumps({"tokens": stats["tokens"], "docs": stats["docs"], "elapsed_sec": stats["elapsed_sec"]}, ensure_ascii=False), encoding="utf-8")
|
|
print(json.dumps({"event": "done", "tokens": stats["tokens"], "docs": stats["docs"], "elapsed_sec": stats["elapsed_sec"], "output": str(final_out)}, ensure_ascii=False), flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|