vimeml-tiny-zh-v1 / source /scripts /benchmarks /prepare_predict_ime_zh.py

Release vimeml-tiny-zh-v1 (step 844730): FP32 bundle, INT8 Core ML packages, source snapshot

81cffe7 verified 1 day ago

20.5 kB

"""Build the Chinese prediction benchmark ime-predict-zh-v1 (docs/zh-CN/benchmarks/predict-zh-plan.md).
usage:
  python -I -X utf8 scripts/benchmarks/prepare_predict_ime_zh.py scan OUT [--sources ...] [--workers 10]
      Leak filter, before any sampling: every 16-Han-character n-gram of every message of every
      candidate dialogue is searched in export-r1/train (Han-normalized, as in the v1 leak check).
      Candidates are all dialogues of the chosen sources that never entered ime-standard-zh-v1
      (its draft, which is a superset of development and blind). Writes OUT/scan.json (aggregates
      and the hit dialogue keys) and OUT/hits.jsonl (local only).
Dialogue texts stay in ignored directories. Nothing here reads a model score.
"""
import argparse
import collections
import glob
import json
import re
import sys
import unicodedata
import zipfile
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
RAW = ROOT / "datasets"
STANDARD_DRAFT = ROOT / "outputs/benchmarks/ime-standard-zh-v1/draft-r1/items.jsonl"
TRAIN = ROOT / "outputs/zh-corpus/export-r1/train"
SOURCES = ("crosswoz", "risawoz", "kdconv", "naturalconv", "cped")
HAN = re.compile("[〇㐀-䶿一-鿿豈-﫿]+")
N = 16
MULTIPLIER = 0x9E3779B97F4A7C15
SEPARATOR = 0


def norm(s):
    s = unicodedata.normalize("NFC", s)
    s = "".join(c for c in s if unicodedata.category(c)[0] != "C" or c in "\n\t")
    return " ".join(s.split())


def han(s):
    return "".join(HAN.findall(s))


def protect_values(values):
    out = set()
    for v in values:
        if isinstance(v, (list, tuple)):
            out |= protect_values(v)
            continue
        if not isinstance(v, str):
            continue
        v = re.sub(r"（[^）]*）|\([^)]*\)", "", norm(v).replace(" ", ""))
        if len(v) >= 2 and HAN.search(v):
            out.add(v)
    return out


# Every loader yields whole dialogues: messages are (text, human, eligible). `human` marks a person's
# turn; `eligible` marks turns prediction targets may come from. `protected` holds annotated entity and
# slot values, which targets must not cover.
def load_crosswoz(raw):
    for split in ("train", "val", "test"):
        with zipfile.ZipFile(raw / "ime-standard-zh-v1/raw/CrossWOZ/data/crosswoz" / f"{split}.json.zip") as z:
            data = json.loads(z.read(z.namelist()[0]))
        for did, d in data.items():
            messages, protected = [], set()
            for m in d["messages"]:
                protected |= protect_values([a[3] for a in m.get("dialog_act", [])])
                protected |= protect_values([s[3] for s in m.get("user_state", [])])
                user = m["role"] == "usr"
                messages.append((m["content"], user, user))
            yield dict(source="crosswoz", source_split=split, dialogue_id=str(did), messages=messages, protected=protected)


def load_risawoz(raw):
    names = {"all_train10000.json": "train", "all_dev600.json": "dev", "all_test600_new.json": "test"}
    with zipfile.ZipFile(raw / "ime-standard-zh-v1/raw/RiSAWOZ/RiSAWOZ-data/source-data.zip") as z:
        for n, split in names.items():
            for d in json.loads(z.read("source-data/" + n)):
                messages, protected = [], set()
                for t in d["dialogue"]:
                    protected |= protect_values([a[3] for a in t.get("user_actions", []) if len(a) > 3])
                    protected |= protect_values([a[3] for a in t.get("system_actions", []) if len(a) > 3])
                    protected |= protect_values(list(t.get("belief_state", {}).get("inform slot-values", {}).values()))
                    messages.append((t["user_utterance"], True, True))
                    if t.get("system_utterance"):
                        messages.append((t["system_utterance"], False, False))
                yield dict(source="risawoz", source_split=split, dialogue_id=d["dialogue_id"], messages=messages,
                           protected=protected)


def load_kdconv(raw):
    for domain in ("film", "music", "travel"):
        for split in ("train", "dev", "test"):
            data = json.loads((raw / "ime-standard-zh-v1/raw/KdConv/data" / domain / f"{split}.json").read_text(encoding="utf-8"))
            for k, d in enumerate(data):
                protected, messages = protect_values([d.get("name", "")]), []
                for m in d["messages"]:
                    for a in m.get("attrs", []):
                        protected |= protect_values([a.get("name", "")])
                        if len(a.get("attrvalue", "")) <= 15:
                            protected |= protect_values([a["attrvalue"]])
                    messages.append((m["message"], True, not m.get("attrs")))  # targets avoid turns quoting knowledge
                yield dict(source="kdconv", source_split=split, dialogue_id=f"{domain}/{split}/{k}", messages=messages,
                           protected=protected)


def load_naturalconv(raw):
    base = raw / "ime-predict-zh-v1/raw/NaturalConv"
    split_of = {}
    for split in ("train", "dev", "test"):
        for line in (base / f"{split}.txt").read_text(encoding="utf-8").split():
            split_of[line] = split
    for d in json.loads((base / "dialog_release.json").read_text(encoding="utf-8")):
        yield dict(source="naturalconv", source_split=split_of[d["dialog_id"]], dialogue_id=d["dialog_id"],
                   messages=[(m, True, True) for m in d["content"]], protected=set())


def load_cped(raw):
    """CPED (GitHub release, data/CPED): TV-series dialogue transcripts, one row per utterance."""
    import csv
    base = raw / "ime-predict-zh-v1/raw/CPED"
    for split in ("train", "valid", "test"):
        dialogues = collections.defaultdict(list)
        with open(base / f"{split}_split.csv", encoding="utf-8-sig", newline="") as stream:
            for row in csv.DictReader(stream):
                dialogues[row["Dialogue_ID"]].append((row["Utterance_ID"], row["Utterance"]))
        for did in sorted(dialogues):
            yield dict(source="cped", source_split=split, dialogue_id=did,
                       messages=[(t, True, True) for _, t in sorted(dialogues[did])], protected=set())


LOADERS = dict(crosswoz=load_crosswoz, risawoz=load_risawoz, kdconv=load_kdconv, naturalconv=load_naturalconv,
               cped=load_cped)


def key(d):
    return f'{d["source"]}|{d["source_split"]}|{d["dialogue_id"]}'


def candidates(sources):
    """All dialogues of the given sources, minus those drafted for ime-standard-zh-v1."""
    used = set()
    with open(STANDARD_DRAFT, encoding="utf-8") as stream:
        for line in stream:
            i = json.loads(line)
            used.add(f'{i["source"]}|{i["source_split"]}|{i["dialogue_id"]}')
    counts = collections.defaultdict(collections.Counter)
    for source in sources:
        for d in LOADERS[source](RAW):
            counts[source]["dialogues"] += 1
            if key(d) in used:
                counts[source]["used_by_standard"] += 1
                continue
            counts[source]["candidates"] += 1
            yield d, counts
    candidates.counts = counts


def hashes(codes):
    """64-bit polynomial hash of every N-window of a uint32 code array (wrapping arithmetic)."""
    import numpy as np
    n = len(codes) - N + 1
    h = np.zeros(n, dtype=np.uint64)
    wide = codes.astype(np.uint64)
    multiplier = np.uint64(MULTIPLIER)
    for k in range(N):
        h *= multiplier
        h += wide[k:k + n]
    return h


PATTERNS = None


def init(path):
    global PATTERNS
    import numpy as np
    PATTERNS = np.load(path)


def scan_file(path):
    """Return (file, documents, [(gram, document index, snippet)]): the first hit of each gram per document."""
    import numpy as np
    import pyarrow.parquet as pq
    documents, found = 0, []
    for batch in pq.ParquetFile(path).iter_batches(batch_size=20000, columns=["text"]):
        texts = [(t or "").replace("\x00", "") for t in batch.column(0).to_pylist()]
        codes = np.frombuffer("\x00".join(texts).encode("utf-32-le"), dtype=np.uint32)
        keep = ((codes >= 0x4E00) & (codes <= 0x9FFF)) | ((codes >= 0x3400) & (codes <= 0x4DBF)) \
            | ((codes >= 0xF900) & (codes <= 0xFAFF)) | (codes == 0x3007) | (codes == SEPARATOR)
        codes = codes[keep]
        if len(codes) >= N:
            h = hashes(codes)
            slot = np.searchsorted(PATTERNS, h)
            slot[slot == len(PATTERNS)] = 0
            hit = np.flatnonzero(PATTERNS[slot] == h)
            if len(hit):
                separators = np.flatnonzero(codes == SEPARATOR)
                seen = set()
                for i in hit.tolist():
                    window = codes[i:i + N]
                    if (window == SEPARATOR).any():
                        continue
                    document = int(np.searchsorted(separators, i))
                    gram = window.tobytes().decode("utf-32-le")
                    if (document, gram) in seen:
                        continue
                    seen.add((document, gram))
                    lo, hi = max(0, i - 40), i + N + 40
                    snippet = codes[lo:hi].tobytes().decode("utf-32-le").replace("\x00", "|")
                    found.append((gram, documents + document, snippet))
        documents += len(texts)
    return Path(path).name, documents, found


def scan(args):
    import numpy as np
    from multiprocessing import Pool
    out = Path(args.out)
    out.mkdir(parents=True)
    owners = collections.defaultdict(set)
    scannable = collections.Counter()
    total = collections.Counter()
    for d, counts in candidates(args.sources):
        any_message = False
        for text, _, _ in d["messages"]:
            h = han(norm(text))
            for k in range(len(h) - N + 1):
                owners[h[k:k + N]].add(key(d))
            any_message |= len(h) >= N
        scannable[d["source"]] += any_message
        total[d["source"]] += 1
    counts = candidates.counts
    grams = sorted(owners)
    codes = np.frombuffer("".join(grams).encode("utf-32-le"), dtype=np.uint32).reshape(len(grams), N)
    multiplier, h = np.uint64(MULTIPLIER), np.zeros(len(grams), dtype=np.uint64)
    for k in range(N):
        h *= multiplier
        h += codes[:, k].astype(np.uint64)
    if len(np.unique(h)) != len(h):
        raise SystemExit("hash collision among benchmark n-grams; change MULTIPLIER")
    np.save(out / "patterns.npy", np.sort(h))
    print(json.dumps({"candidates": dict(total), "grams": len(grams)}, ensure_ascii=False), flush=True)
    files = sorted(glob.glob(str(TRAIN / "*.parquet")))
    documents, hits = {}, []
    with Pool(args.workers, init, (str(out / "patterns.npy"),)) as pool:
        for name, n, found in pool.imap_unordered(scan_file, files):
            documents[name] = n
            for gram, document, snippet in found:
                if gram in owners:  # exact check; a bare hash match is not a hit
                    hits.append(dict(file=name, document=document, gram=gram, dialogues=sorted(owners[gram]), snippet=snippet))
            print(name, n, len(found), flush=True)
    hit_dialogues = collections.defaultdict(set)
    for hit in hits:
        for k in hit["dialogues"]:
            hit_dialogues[k.split("|")[0]].add(k)
    with open(out / "hits.jsonl", "w", encoding="utf-8", newline="\n") as stream:
        for hit in sorted(hits, key=lambda x: (x["file"], x["document"], x["gram"])):
            stream.write(json.dumps(hit, ensure_ascii=False) + "\n")
    summary = dict(
        n=N, train_files=len(files), train_documents=sum(documents.values()), grams=len(grams),
        hit_grams=len({h["gram"] for h in hits}), hit_documents=len({(h["file"], h["document"]) for h in hits}),
        sources={s: dict(counts[s], with_scannable_message=scannable[s], hit_dialogues=len(hit_dialogues[s]),
                         hit_rate=len(hit_dialogues[s]) / max(1, total[s])) for s in args.sources},
        dropped={s: sorted(hit_dialogues[s]) for s in args.sources})
    (out / "scan.json").write_text(json.dumps(summary, ensure_ascii=False, indent=1) + "\n", encoding="utf-8")
    print(json.dumps(summary["sources"], ensure_ascii=False, indent=1))


SEED = 20261111
PER_SOURCE = 1200
POSITIONS = 3
BAD = re.compile(r"https?://|www\.|@\w|<[a-z/]|&[a-z]+;|\[[A-Z_]+\]|�", re.I)
MIN_HAN_SHARE = 0.6
HAN_CHAR = re.compile("[〇㐀-䶿一-鿿豈-﫿]")
PARTICLES = set("的了呢吗吧啊呀哦哈嘛啦着过地得么呗哟喔噢嘞咯")


def clean(text):
    """Mechanical message filter: 4-120 characters, mostly Han, no markup, links or replacement characters."""
    return 4 <= len(text) <= 120 and not BAD.search(text) and len(HAN_CHAR.findall(text)) >= MIN_HAN_SHARE * len(text)


def blocked_positions(text, protected):
    blocked = set()
    for value in protected:
        for m in re.finditer(re.escape(value), text):
            blocked.update(range(m.start(), m.end()))
    return blocked


def next_positions(text, blocked):
    """Track A positions: between two Han characters, with at least two Han characters before the next
    non-Han character, and the first six characters of that continuation not on an annotated value."""
    out = []
    for run in HAN.finditer(text):
        for p in range(run.start() + 1, run.end() - 1):
            if not blocked.intersection(range(p, min(run.end(), p + 6))):
                out.append((p, run.end()))
    return out


def build(args):
    import hashlib
    import importlib.util
    import random
    spec = importlib.util.spec_from_file_location("standard_zh", ROOT / "scripts/benchmarks/prepare_standard_ime_zh.py")
    standard = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(standard)
    out = Path(args.out)
    out.mkdir(parents=True)
    scans = [json.loads((Path(s) / "scan.json").read_text(encoding="utf-8")) for s in args.scan]
    scan_summary = dict(n=scans[0]["n"], train_documents=scans[0]["train_documents"], sources={}, dropped={})
    for s in scans:
        if (s["n"], s["train_documents"]) != (scan_summary["n"], scan_summary["train_documents"]):
            raise SystemExit("scans used different n-gram lengths or training exports")
        scan_summary["sources"].update(s["sources"])
        scan_summary["dropped"].update(s["dropped"])
    missing = set(args.sources) - set(scan_summary["sources"])
    if missing:
        raise SystemExit(f"no leak scan for {sorted(missing)}")
    scan_summary["sources"] = {s: scan_summary["sources"][s] for s in args.sources}
    dropped = {k for keys in scan_summary["dropped"].values() for k in keys}
    reader = standard.Reader()
    pools = collections.defaultdict(list)
    for d, _ in candidates(args.sources):
        if d["source"] in args.sources and key(d) not in dropped:
            pools[d["source"]].append(d)
    items_a, items_c, messages_b = [], [], []
    stats = collections.defaultdict(collections.Counter)
    for source in args.sources:
        rng = random.Random(f"{SEED}/{source}")
        dialogues = sorted(pools[source], key=lambda d: (d["source_split"], d["dialogue_id"]))
        stats[source]["clean_candidates"] = len(dialogues)
        rng.shuffle(dialogues)
        kept = []
        for d in dialogues:
            if len(kept) >= PER_SOURCE:
                break
            texts = [norm(t) for t, _, _ in d["messages"]]
            eligible = [k for k, (_, human, ok) in enumerate(d["messages"]) if human and ok and clean(texts[k])]
            blocked = {k: blocked_positions(texts[k], d["protected"]) for k in eligible}
            positions = [(k, p, end) for k in eligible for p, end in next_positions(texts[k], blocked[k])]
            if len(positions) < POSITIONS:
                stats[source]["too_few_positions"] += 1
                continue
            # Track C: one clause-final 2-6 character span with an unambiguous automatic reading.
            spans = []
            for k in eligible:
                for s, e, r in standard.spans(texts[k], d["protected"], reader):
                    if 2 <= e - s <= 6:
                        spans.append((k, s, e, r))
            rng.shuffle(spans)
            completion = None
            for k, s, e, r in spans:
                syllables, readers = reader.readings(texts[k][r:e])
                if syllables is None or any(len(set(v)) > 1 for v in zip(*readers)):
                    continue  # no reading, or the three readers disagree somewhere in the clause
                completion = (k, s, e, syllables[s - r:])
                break
            if completion is None:
                stats[source]["no_completion_span"] += 1
                continue
            kept.append((d, texts, rng.sample(positions, POSITIONS), completion))
        if len(kept) < PER_SOURCE:
            raise SystemExit(f"{source}: only {len(kept)} usable dialogues")
        stats[source]["dialogues"] = len(kept)
        order = sorted(kept, key=lambda x: hashlib.sha256(f'{SEED}/{x[0]["dialogue_id"]}'.encode()).hexdigest())
        for n, (d, texts, positions, (ck, cs, ce, syllables)) in enumerate(order):
            split = ("development", "blind")[n % 2]
            base = f'zh-pred-{source}-{hashlib.sha256(key(d).encode()).hexdigest()[:10]}'
            common = dict(source=source, source_split=d["source_split"], dialogue_id=d["dialogue_id"], split=split,
                          dialogue=f"{source}:{d['dialogue_id']}")
            for j, (k, p, end) in enumerate(sorted(positions)):
                items_a.append(dict(id=f"{base}-a{j + 1}", **common, turn=k, previous_turns=texts[:k],
                                    prefix=texts[k][:p], continuation=texts[k][p:end]))
            for k in sorted({k for k, _, _ in positions}):
                messages_b.append(dict(id=f"{base}-b{k}", **common, turn=k, previous_turns=texts[:k], message=texts[k]))
            items_c.append(dict(id=f"{base}-c", **common, turn=ck, previous_turns=texts[:ck], prefix=texts[ck][:cs],
                                target=texts[ck][cs:ce], syllables=syllables,
                                keys_first_syllable=syllables[0], keys_initials="".join(x[0] for x in syllables)))
    files = {"items-a.jsonl": items_a, "messages-b.jsonl": messages_b, "items-c.jsonl": items_c}
    digests = {}
    for name, rows in files.items():
        rows.sort(key=lambda r: r["id"])
        if len({r["id"] for r in rows}) != len(rows):
            raise SystemExit(f"duplicate ids in {name}")
        with open(out / name, "w", encoding="utf-8", newline="\n") as stream:
            for row in rows:
                stream.write(json.dumps(row, ensure_ascii=False) + "\n")
        digests[name] = hashlib.sha256((out / name).read_bytes()).hexdigest()
    count = lambda rows: dict(sorted(collections.Counter(f'{r["source"]}/{r["split"]}' for r in rows).items()))
    manifest = dict(version="ime-predict-zh-v1-r1", seed=SEED, per_source=PER_SOURCE, positions=POSITIONS,
                    sources=list(args.sources), leak_scan=dict(n=scan_summary["n"], train_documents=scan_summary["train_documents"],
                                                               by_source=scan_summary["sources"]),
                    build={k: dict(v) for k, v in stats.items()},
                    counts={"track_a": count(items_a), "track_b_messages": count(messages_b), "track_c": count(items_c)},
                    labels="source text; track C readings are automatic (pypinyin, pypinyin over jieba, g2pM agree)",
                    sha256=digests, python=sys.version.split()[0])
    (out / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=1) + "\n", encoding="utf-8")
    print(json.dumps({k: manifest[k] for k in ("build", "counts", "sha256")}, ensure_ascii=False, indent=1))


def main():
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    commands = parser.add_subparsers(dest="command", required=True)
    p = commands.add_parser("scan")
    p.add_argument("out")
    p.add_argument("--workers", type=int, default=10)
    p.add_argument("--sources", nargs="+", choices=SOURCES, default=list(SOURCES))
    p = commands.add_parser("build")
    p.add_argument("out")
    p.add_argument("--scan", nargs="+", required=True, help="directories written by scan; together they cover --sources")
    p.add_argument("--sources", nargs="+", choices=SOURCES, default=list(SOURCES))
    args = parser.parse_args()
    {"scan": scan, "build": build}[args.command](args)


if __name__ == "__main__":
    main()

	"""Build the Chinese prediction benchmark ime-predict-zh-v1 (docs/zh-CN/benchmarks/predict-zh-plan.md).

	usage:
	python -I -X utf8 scripts/benchmarks/prepare_predict_ime_zh.py scan OUT [--sources ...] [--workers 10]
	Leak filter, before any sampling: every 16-Han-character n-gram of every message of every
	candidate dialogue is searched in export-r1/train (Han-normalized, as in the v1 leak check).
	Candidates are all dialogues of the chosen sources that never entered ime-standard-zh-v1
	(its draft, which is a superset of development and blind). Writes OUT/scan.json (aggregates
	and the hit dialogue keys) and OUT/hits.jsonl (local only).

	Dialogue texts stay in ignored directories. Nothing here reads a model score.
	"""
	import argparse
	import collections
	import glob
	import json
	import re
	import sys
	import unicodedata
	import zipfile
	from pathlib import Path

	ROOT = Path(__file__).resolve().parents[2]
	RAW = ROOT / "datasets"
	STANDARD_DRAFT = ROOT / "outputs/benchmarks/ime-standard-zh-v1/draft-r1/items.jsonl"
	TRAIN = ROOT / "outputs/zh-corpus/export-r1/train"
	SOURCES = ("crosswoz", "risawoz", "kdconv", "naturalconv", "cped")
	HAN = re.compile("[〇㐀-䶿一-鿿豈-﫿]+")
	N = 16
	MULTIPLIER = 0x9E3779B97F4A7C15
	SEPARATOR = 0


	def norm(s):
	s = unicodedata.normalize("NFC", s)
	s = "".join(c for c in s if unicodedata.category(c)[0] != "C" or c in "\n\t")
	return " ".join(s.split())


	def han(s):
	return "".join(HAN.findall(s))


	def protect_values(values):
	out = set()
	for v in values:
	if isinstance(v, (list, tuple)):
	out \|= protect_values(v)
	continue
	if not isinstance(v, str):
	continue
	v = re.sub(r"（[^）]）\|\([^)]\)", "", norm(v).replace(" ", ""))
	if len(v) >= 2 and HAN.search(v):
	out.add(v)
	return out


	# Every loader yields whole dialogues: messages are (text, human, eligible). `human` marks a person's
	# turn; `eligible` marks turns prediction targets may come from. `protected` holds annotated entity and
	# slot values, which targets must not cover.
	def load_crosswoz(raw):
	for split in ("train", "val", "test"):
	with zipfile.ZipFile(raw / "ime-standard-zh-v1/raw/CrossWOZ/data/crosswoz" / f"{split}.json.zip") as z:
	data = json.loads(z.read(z.namelist()[0]))
	for did, d in data.items():
	messages, protected = [], set()
	for m in d["messages"]:
	protected \|= protect_values([a[3] for a in m.get("dialog_act", [])])
	protected \|= protect_values([s[3] for s in m.get("user_state", [])])
	user = m["role"] == "usr"
	messages.append((m["content"], user, user))
	yield dict(source="crosswoz", source_split=split, dialogue_id=str(did), messages=messages, protected=protected)


	def load_risawoz(raw):
	names = {"all_train10000.json": "train", "all_dev600.json": "dev", "all_test600_new.json": "test"}
	with zipfile.ZipFile(raw / "ime-standard-zh-v1/raw/RiSAWOZ/RiSAWOZ-data/source-data.zip") as z:
	for n, split in names.items():
	for d in json.loads(z.read("source-data/" + n)):
	messages, protected = [], set()
	for t in d["dialogue"]:
	protected \|= protect_values([a[3] for a in t.get("user_actions", []) if len(a) > 3])
	protected \|= protect_values([a[3] for a in t.get("system_actions", []) if len(a) > 3])
	protected \|= protect_values(list(t.get("belief_state", {}).get("inform slot-values", {}).values()))
	messages.append((t["user_utterance"], True, True))
	if t.get("system_utterance"):
	messages.append((t["system_utterance"], False, False))
	yield dict(source="risawoz", source_split=split, dialogue_id=d["dialogue_id"], messages=messages,
	protected=protected)


	def load_kdconv(raw):
	for domain in ("film", "music", "travel"):
	for split in ("train", "dev", "test"):
	data = json.loads((raw / "ime-standard-zh-v1/raw/KdConv/data" / domain / f"{split}.json").read_text(encoding="utf-8"))
	for k, d in enumerate(data):
	protected, messages = protect_values([d.get("name", "")]), []
	for m in d["messages"]:
	for a in m.get("attrs", []):
	protected \|= protect_values([a.get("name", "")])
	if len(a.get("attrvalue", "")) <= 15:
	protected \|= protect_values([a["attrvalue"]])
	messages.append((m["message"], True, not m.get("attrs"))) # targets avoid turns quoting knowledge
	yield dict(source="kdconv", source_split=split, dialogue_id=f"{domain}/{split}/{k}", messages=messages,
	protected=protected)


	def load_naturalconv(raw):
	base = raw / "ime-predict-zh-v1/raw/NaturalConv"
	split_of = {}
	for split in ("train", "dev", "test"):
	for line in (base / f"{split}.txt").read_text(encoding="utf-8").split():
	split_of[line] = split
	for d in json.loads((base / "dialog_release.json").read_text(encoding="utf-8")):
	yield dict(source="naturalconv", source_split=split_of[d["dialog_id"]], dialogue_id=d["dialog_id"],
	messages=[(m, True, True) for m in d["content"]], protected=set())


	def load_cped(raw):
	"""CPED (GitHub release, data/CPED): TV-series dialogue transcripts, one row per utterance."""
	import csv
	base = raw / "ime-predict-zh-v1/raw/CPED"
	for split in ("train", "valid", "test"):
	dialogues = collections.defaultdict(list)
	with open(base / f"{split}_split.csv", encoding="utf-8-sig", newline="") as stream:
	for row in csv.DictReader(stream):
	dialogues[row["Dialogue_ID"]].append((row["Utterance_ID"], row["Utterance"]))
	for did in sorted(dialogues):
	yield dict(source="cped", source_split=split, dialogue_id=did,
	messages=[(t, True, True) for _, t in sorted(dialogues[did])], protected=set())


	LOADERS = dict(crosswoz=load_crosswoz, risawoz=load_risawoz, kdconv=load_kdconv, naturalconv=load_naturalconv,
	cped=load_cped)


	def key(d):
	return f'{d["source"]}\|{d["source_split"]}\|{d["dialogue_id"]}'


	def candidates(sources):
	"""All dialogues of the given sources, minus those drafted for ime-standard-zh-v1."""
	used = set()
	with open(STANDARD_DRAFT, encoding="utf-8") as stream:
	for line in stream:
	i = json.loads(line)
	used.add(f'{i["source"]}\|{i["source_split"]}\|{i["dialogue_id"]}')
	counts = collections.defaultdict(collections.Counter)
	for source in sources:
	for d in LOADERS[source](RAW):
	counts[source]["dialogues"] += 1
	if key(d) in used:
	counts[source]["used_by_standard"] += 1
	continue
	counts[source]["candidates"] += 1
	yield d, counts
	candidates.counts = counts


	def hashes(codes):
	"""64-bit polynomial hash of every N-window of a uint32 code array (wrapping arithmetic)."""
	import numpy as np
	n = len(codes) - N + 1
	h = np.zeros(n, dtype=np.uint64)
	wide = codes.astype(np.uint64)
	multiplier = np.uint64(MULTIPLIER)
	for k in range(N):
	h *= multiplier
	h += wide[k:k + n]
	return h


	PATTERNS = None


	def init(path):
	global PATTERNS
	import numpy as np
	PATTERNS = np.load(path)


	def scan_file(path):
	"""Return (file, documents, [(gram, document index, snippet)]): the first hit of each gram per document."""
	import numpy as np
	import pyarrow.parquet as pq
	documents, found = 0, []
	for batch in pq.ParquetFile(path).iter_batches(batch_size=20000, columns=["text"]):
	texts = [(t or "").replace("\x00", "") for t in batch.column(0).to_pylist()]
	codes = np.frombuffer("\x00".join(texts).encode("utf-32-le"), dtype=np.uint32)
	keep = ((codes >= 0x4E00) & (codes <= 0x9FFF)) \| ((codes >= 0x3400) & (codes <= 0x4DBF)) \
	\| ((codes >= 0xF900) & (codes <= 0xFAFF)) \| (codes == 0x3007) \| (codes == SEPARATOR)
	codes = codes[keep]
	if len(codes) >= N:
	h = hashes(codes)
	slot = np.searchsorted(PATTERNS, h)
	slot[slot == len(PATTERNS)] = 0
	hit = np.flatnonzero(PATTERNS[slot] == h)
	if len(hit):
	separators = np.flatnonzero(codes == SEPARATOR)
	seen = set()
	for i in hit.tolist():
	window = codes[i:i + N]
	if (window == SEPARATOR).any():
	continue
	document = int(np.searchsorted(separators, i))
	gram = window.tobytes().decode("utf-32-le")
	if (document, gram) in seen:
	continue
	seen.add((document, gram))
	lo, hi = max(0, i - 40), i + N + 40
	snippet = codes[lo:hi].tobytes().decode("utf-32-le").replace("\x00", "\|")
	found.append((gram, documents + document, snippet))
	documents += len(texts)
	return Path(path).name, documents, found


	def scan(args):
	import numpy as np
	from multiprocessing import Pool
	out = Path(args.out)
	out.mkdir(parents=True)
	owners = collections.defaultdict(set)
	scannable = collections.Counter()
	total = collections.Counter()
	for d, counts in candidates(args.sources):
	any_message = False
	for text, _, _ in d["messages"]:
	h = han(norm(text))
	for k in range(len(h) - N + 1):
	owners[h[k:k + N]].add(key(d))
	any_message \|= len(h) >= N
	scannable[d["source"]] += any_message
	total[d["source"]] += 1
	counts = candidates.counts
	grams = sorted(owners)
	codes = np.frombuffer("".join(grams).encode("utf-32-le"), dtype=np.uint32).reshape(len(grams), N)
	multiplier, h = np.uint64(MULTIPLIER), np.zeros(len(grams), dtype=np.uint64)
	for k in range(N):
	h *= multiplier
	h += codes[:, k].astype(np.uint64)
	if len(np.unique(h)) != len(h):
	raise SystemExit("hash collision among benchmark n-grams; change MULTIPLIER")
	np.save(out / "patterns.npy", np.sort(h))
	print(json.dumps({"candidates": dict(total), "grams": len(grams)}, ensure_ascii=False), flush=True)
	files = sorted(glob.glob(str(TRAIN / "*.parquet")))
	documents, hits = {}, []
	with Pool(args.workers, init, (str(out / "patterns.npy"),)) as pool:
	for name, n, found in pool.imap_unordered(scan_file, files):
	documents[name] = n
	for gram, document, snippet in found:
	if gram in owners: # exact check; a bare hash match is not a hit
	hits.append(dict(file=name, document=document, gram=gram, dialogues=sorted(owners[gram]), snippet=snippet))
	print(name, n, len(found), flush=True)
	hit_dialogues = collections.defaultdict(set)
	for hit in hits:
	for k in hit["dialogues"]:
	hit_dialogues[k.split("\|")[0]].add(k)
	with open(out / "hits.jsonl", "w", encoding="utf-8", newline="\n") as stream:
	for hit in sorted(hits, key=lambda x: (x["file"], x["document"], x["gram"])):
	stream.write(json.dumps(hit, ensure_ascii=False) + "\n")
	summary = dict(
	n=N, train_files=len(files), train_documents=sum(documents.values()), grams=len(grams),
	hit_grams=len({h["gram"] for h in hits}), hit_documents=len({(h["file"], h["document"]) for h in hits}),
	sources={s: dict(counts[s], with_scannable_message=scannable[s], hit_dialogues=len(hit_dialogues[s]),
	hit_rate=len(hit_dialogues[s]) / max(1, total[s])) for s in args.sources},
	dropped={s: sorted(hit_dialogues[s]) for s in args.sources})
	(out / "scan.json").write_text(json.dumps(summary, ensure_ascii=False, indent=1) + "\n", encoding="utf-8")
	print(json.dumps(summary["sources"], ensure_ascii=False, indent=1))


	SEED = 20261111
	PER_SOURCE = 1200
	POSITIONS = 3
	BAD = re.compile(r"https?://\|www\.\|@\w\|<[a-z/]\|&[a-z]+;\|\[[A-Z_]+\]\|�", re.I)
	MIN_HAN_SHARE = 0.6
	HAN_CHAR = re.compile("[〇㐀-䶿一-鿿豈-﫿]")
	PARTICLES = set("的了呢吗吧啊呀哦哈嘛啦着过地得么呗哟喔噢嘞咯")


	def clean(text):
	"""Mechanical message filter: 4-120 characters, mostly Han, no markup, links or replacement characters."""
	return 4 <= len(text) <= 120 and not BAD.search(text) and len(HAN_CHAR.findall(text)) >= MIN_HAN_SHARE * len(text)


	def blocked_positions(text, protected):
	blocked = set()
	for value in protected:
	for m in re.finditer(re.escape(value), text):
	blocked.update(range(m.start(), m.end()))
	return blocked


	def next_positions(text, blocked):
	"""Track A positions: between two Han characters, with at least two Han characters before the next
	non-Han character, and the first six characters of that continuation not on an annotated value."""
	out = []
	for run in HAN.finditer(text):
	for p in range(run.start() + 1, run.end() - 1):
	if not blocked.intersection(range(p, min(run.end(), p + 6))):
	out.append((p, run.end()))
	return out


	def build(args):
	import hashlib
	import importlib.util
	import random
	spec = importlib.util.spec_from_file_location("standard_zh", ROOT / "scripts/benchmarks/prepare_standard_ime_zh.py")
	standard = importlib.util.module_from_spec(spec)
	spec.loader.exec_module(standard)
	out = Path(args.out)
	out.mkdir(parents=True)
	scans = [json.loads((Path(s) / "scan.json").read_text(encoding="utf-8")) for s in args.scan]
	scan_summary = dict(n=scans[0]["n"], train_documents=scans[0]["train_documents"], sources={}, dropped={})
	for s in scans:
	if (s["n"], s["train_documents"]) != (scan_summary["n"], scan_summary["train_documents"]):
	raise SystemExit("scans used different n-gram lengths or training exports")
	scan_summary["sources"].update(s["sources"])
	scan_summary["dropped"].update(s["dropped"])
	missing = set(args.sources) - set(scan_summary["sources"])
	if missing:
	raise SystemExit(f"no leak scan for {sorted(missing)}")
	scan_summary["sources"] = {s: scan_summary["sources"][s] for s in args.sources}
	dropped = {k for keys in scan_summary["dropped"].values() for k in keys}
	reader = standard.Reader()
	pools = collections.defaultdict(list)
	for d, _ in candidates(args.sources):
	if d["source"] in args.sources and key(d) not in dropped:
	pools[d["source"]].append(d)
	items_a, items_c, messages_b = [], [], []
	stats = collections.defaultdict(collections.Counter)
	for source in args.sources:
	rng = random.Random(f"{SEED}/{source}")
	dialogues = sorted(pools[source], key=lambda d: (d["source_split"], d["dialogue_id"]))
	stats[source]["clean_candidates"] = len(dialogues)
	rng.shuffle(dialogues)
	kept = []
	for d in dialogues:
	if len(kept) >= PER_SOURCE:
	break
	texts = [norm(t) for t, _, _ in d["messages"]]
	eligible = [k for k, (_, human, ok) in enumerate(d["messages"]) if human and ok and clean(texts[k])]
	blocked = {k: blocked_positions(texts[k], d["protected"]) for k in eligible}
	positions = [(k, p, end) for k in eligible for p, end in next_positions(texts[k], blocked[k])]
	if len(positions) < POSITIONS:
	stats[source]["too_few_positions"] += 1
	continue
	# Track C: one clause-final 2-6 character span with an unambiguous automatic reading.
	spans = []
	for k in eligible:
	for s, e, r in standard.spans(texts[k], d["protected"], reader):
	if 2 <= e - s <= 6:
	spans.append((k, s, e, r))
	rng.shuffle(spans)
	completion = None
	for k, s, e, r in spans:
	syllables, readers = reader.readings(texts[k][r:e])
	if syllables is None or any(len(set(v)) > 1 for v in zip(*readers)):
	continue # no reading, or the three readers disagree somewhere in the clause
	completion = (k, s, e, syllables[s - r:])
	break
	if completion is None:
	stats[source]["no_completion_span"] += 1
	continue
	kept.append((d, texts, rng.sample(positions, POSITIONS), completion))
	if len(kept) < PER_SOURCE:
	raise SystemExit(f"{source}: only {len(kept)} usable dialogues")
	stats[source]["dialogues"] = len(kept)
	order = sorted(kept, key=lambda x: hashlib.sha256(f'{SEED}/{x[0]["dialogue_id"]}'.encode()).hexdigest())
	for n, (d, texts, positions, (ck, cs, ce, syllables)) in enumerate(order):
	split = ("development", "blind")[n % 2]
	base = f'zh-pred-{source}-{hashlib.sha256(key(d).encode()).hexdigest()[:10]}'
	common = dict(source=source, source_split=d["source_split"], dialogue_id=d["dialogue_id"], split=split,
	dialogue=f"{source}:{d['dialogue_id']}")
	for j, (k, p, end) in enumerate(sorted(positions)):
	items_a.append(dict(id=f"{base}-a{j + 1}", **common, turn=k, previous_turns=texts[:k],
	prefix=texts[k][:p], continuation=texts[k][p:end]))
	for k in sorted({k for k, _, _ in positions}):
	messages_b.append(dict(id=f"{base}-b{k}", **common, turn=k, previous_turns=texts[:k], message=texts[k]))
	items_c.append(dict(id=f"{base}-c", **common, turn=ck, previous_turns=texts[:ck], prefix=texts[ck][:cs],
	target=texts[ck][cs:ce], syllables=syllables,
	keys_first_syllable=syllables[0], keys_initials="".join(x[0] for x in syllables)))
	files = {"items-a.jsonl": items_a, "messages-b.jsonl": messages_b, "items-c.jsonl": items_c}
	digests = {}
	for name, rows in files.items():
	rows.sort(key=lambda r: r["id"])
	if len({r["id"] for r in rows}) != len(rows):
	raise SystemExit(f"duplicate ids in {name}")
	with open(out / name, "w", encoding="utf-8", newline="\n") as stream:
	for row in rows:
	stream.write(json.dumps(row, ensure_ascii=False) + "\n")
	digests[name] = hashlib.sha256((out / name).read_bytes()).hexdigest()
	count = lambda rows: dict(sorted(collections.Counter(f'{r["source"]}/{r["split"]}' for r in rows).items()))
	manifest = dict(version="ime-predict-zh-v1-r1", seed=SEED, per_source=PER_SOURCE, positions=POSITIONS,
	sources=list(args.sources), leak_scan=dict(n=scan_summary["n"], train_documents=scan_summary["train_documents"],
	by_source=scan_summary["sources"]),
	build={k: dict(v) for k, v in stats.items()},
	counts={"track_a": count(items_a), "track_b_messages": count(messages_b), "track_c": count(items_c)},
	labels="source text; track C readings are automatic (pypinyin, pypinyin over jieba, g2pM agree)",
	sha256=digests, python=sys.version.split()[0])
	(out / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=1) + "\n", encoding="utf-8")
	print(json.dumps({k: manifest[k] for k in ("build", "counts", "sha256")}, ensure_ascii=False, indent=1))


	def main():
	parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
	commands = parser.add_subparsers(dest="command", required=True)
	p = commands.add_parser("scan")
	p.add_argument("out")
	p.add_argument("--workers", type=int, default=10)
	p.add_argument("--sources", nargs="+", choices=SOURCES, default=list(SOURCES))
	p = commands.add_parser("build")
	p.add_argument("out")
	p.add_argument("--scan", nargs="+", required=True, help="directories written by scan; together they cover --sources")
	p.add_argument("--sources", nargs="+", choices=SOURCES, default=list(SOURCES))
	args = parser.parse_args()
	{"scan": scan, "build": build}[args.command](args)


	if __name__ == "__main__":
	main()