72 lines
2.9 KiB
Python
72 lines
2.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Build a deterministic punctuation-balanced fine-tuning corpus."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import heapq
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
CATEGORIES = {
|
|
"question": lambda text: "?" in text,
|
|
"exclamation": lambda text: "!" in text,
|
|
"quotes": lambda text: "«" in text or "»" in text,
|
|
"ellipsis": lambda text: "…" in text or "..." in text,
|
|
"neutral": lambda text: not any(mark in text for mark in ("?", "!", "«", "»", "…", "...")),
|
|
}
|
|
|
|
|
|
def build(args: argparse.Namespace) -> None:
|
|
if args.output.exists() or args.metadata.exists():
|
|
raise FileExistsError("Refusing to overwrite output or metadata")
|
|
excluded = set()
|
|
if args.exclude:
|
|
excluded = {line.strip() for line in args.exclude.read_text(encoding="utf-8-sig").splitlines() if line.strip()}
|
|
heaps: dict[str, list[tuple[int, str]]] = {name: [] for name in CATEGORIES}
|
|
scanned = 0
|
|
for text in args.source.read_text(encoding="utf-8").splitlines():
|
|
text = text.strip()
|
|
if not text or text in excluded:
|
|
continue
|
|
scanned += 1
|
|
for name, predicate in CATEGORIES.items():
|
|
if not predicate(text):
|
|
continue
|
|
score = int.from_bytes(hashlib.blake2b(f"{name}\0{text}".encode("utf-8"), digest_size=8).digest(), "big")
|
|
heap = heaps[name]
|
|
if len(heap) < args.per_category:
|
|
heapq.heappush(heap, (-score, text))
|
|
elif score < -heap[0][0]:
|
|
heapq.heapreplace(heap, (-score, text))
|
|
selected: dict[str, set[str]] = {
|
|
name: {text for _, text in heap} for name, heap in heaps.items()
|
|
}
|
|
combined = sorted(set().union(*selected.values()), key=lambda text: hashlib.sha256(text.encode("utf-8")).digest())
|
|
args.output.write_text("\n".join(combined) + "\n", encoding="utf-8")
|
|
metadata = {
|
|
"source": str(args.source), "exclude": str(args.exclude) if args.exclude else None,
|
|
"scanned": scanned, "per_category_limit": args.per_category,
|
|
"category_counts": {name: len(values) for name, values in selected.items()},
|
|
"unique_sentences": len(combined), "output_bytes": args.output.stat().st_size,
|
|
}
|
|
args.metadata.write_text(json.dumps(metadata, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
print(json.dumps(metadata, ensure_ascii=False, indent=2))
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--source", type=Path, required=True)
|
|
parser.add_argument("--exclude", type=Path)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--metadata", type=Path, required=True)
|
|
parser.add_argument("--per-category", type=int, default=4000)
|
|
return parser.parse_args()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(build(parse_args()) or 0)
|