| """Final export of the vimeml-tiny-zh-v1 corpus for upload (zh-v13-export-r1). |
| |
| web stage-a-v2 decision=keep documents, minus exact duplicates and every document on a |
| listed spam host or spam farm (same host -> site rule and spam_allow as `zh_stage_a.py |
| report`). Review documents are never part of this command. |
| review stage-a-v2 decision=review documents whose id is in --accept (one id per line). The |
| accept list comes from the private review tooling (all votes keep); a review document |
| that is not on the list is dropped, reviewed or not. |
| chat LCCC v2 train/validation/test utterances. lccc-v2/review.jsonl (weak-term hits that were |
| never reviewed) is not exported. |
| finish SHA256SUMS and manifest.json over everything written. |
| |
| A fixed 0.05% of web and review documents, chosen by SHA-256 of the id, goes to validation/. |
| Output files are zstd-compressed Parquet of about --file-bytes of text each, so a failed download |
| costs one file. Each input shard is written to a temporary name and renamed when complete; |
| completed shards are skipped, nothing is overwritten. |
| |
| usage: |
| python -I -X utf8 scripts/corpus/zh_export.py web OUT [--workers 8] |
| python -I -X utf8 scripts/corpus/zh_export.py review OUT --accept accept-ids.txt |
| python -I -X utf8 scripts/corpus/zh_export.py chat OUT |
| python -I -X utf8 scripts/corpus/zh_export.py finish OUT |
| """ |
| import argparse |
| import hashlib |
| import json |
| import os |
| import sys |
| import time |
| import tomllib |
| from concurrent.futures import ProcessPoolExecutor |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[2] |
| sys.path.insert(0, str(ROOT / 'src')) |
| from vimeml.corpus.zh_filter import in_domains, registrable |
|
|
| VERSION = 'zh-v13-export-r1' |
| STAGE_A = ROOT / 'outputs/zh-corpus/stage-a-v2' |
| LCCC = ROOT / 'outputs/zh-corpus/lccc-v2' |
| TERMS = ROOT / 'configs/zh-filter-terms.toml' |
| VALIDATION_RATE = 0.0005 |
| COLUMNS = ['id', 'host', 'script', 'han', 'text'] |
|
|
|
|
| def is_validation(doc_id): |
| digest = hashlib.sha256(('zh-v13-validation\0' + doc_id).encode()).digest() |
| return int.from_bytes(digest[:8], 'big') < int(VALIDATION_RATE * (1 << 64)) |
|
|
|
|
| def sha256(path): |
| h = hashlib.sha256() |
| with open(path, 'rb') as stream: |
| while block := stream.read(1 << 22): |
| h.update(block) |
| return h.hexdigest() |
|
|
|
|
| def spam_rules(): |
| spam = json.loads((STAGE_A / 'spam-hosts.json').read_text(encoding='utf-8')) |
| allow = tuple(tomllib.loads(TERMS.read_text(encoding='utf-8')).get('domains', {}).get('spam_allow', [])) |
| return {h['host'] for h in spam['hosts']}, {f['domain'] for f in spam['farms']}, allow |
|
|
|
|
| def duplicate_ids(): |
| import pyarrow.parquet as pq |
| return set(pq.read_table(STAGE_A / 'duplicates.parquet', columns=['id']).column('id').to_pylist()) |
|
|
|
|
| class Writer: |
| """Rolls Parquet files of about `limit` text bytes under out/<split>/<prefix>-NNN.parquet.""" |
|
|
| def __init__(self, out, split, prefix, limit, level, schema): |
| self.dir, self.prefix, self.limit, self.level, self.schema = out / split, prefix, limit, level, schema |
| self.dir.mkdir(parents=True, exist_ok=True) |
| self.writer, self.index, self.bytes, self.files = None, 0, 0, [] |
|
|
| def write(self, table, text_bytes): |
| import pyarrow.parquet as pq |
| if self.writer is None: |
| path = self.dir / f'{self.prefix}-{self.index:03d}.parquet.part' |
| self.writer = pq.ParquetWriter(path, self.schema, compression='zstd', compression_level=self.level) |
| self.files.append(path) |
| self.writer.write_table(table, row_group_size=len(table)) |
| self.bytes += text_bytes |
| if self.bytes >= self.limit: |
| self.close_file() |
|
|
| def close_file(self): |
| if self.writer is not None: |
| self.writer.close() |
| self.writer, self.index, self.bytes = None, self.index + 1, 0 |
|
|
| def finish(self): |
| self.close_file() |
| done = [] |
| for part in self.files: |
| final = part.with_suffix('') |
| os.replace(part, final) |
| done.append(final) |
| return done |
|
|
|
|
| def export_shard(job): |
| import pyarrow as pa |
| import pyarrow.compute as pc |
| import pyarrow.parquet as pq |
| shard, out, kind, limit, level, accept_path = job |
| out = Path(out) |
| marker = out / 'parts' / f'{kind}-{shard}.json' |
| if marker.exists(): |
| return json.loads(marker.read_text(encoding='utf-8')) |
| for stale in out.glob(f'*/{kind}-{shard}-*.parquet.part'): |
| stale.unlink() |
| hosts, farms, allow = spam_rules() |
| duplicates = duplicate_ids() |
| accept = set(Path(accept_path).read_text(encoding='utf-8').split()) if accept_path else None |
| schema = pa.schema([('id', pa.string()), ('host', pa.string()), ('script', pa.string()), |
| ('han', pa.int64()), ('text', pa.string())]) |
| writers = {split: Writer(out, split, f'{kind}-{shard}', limit, level, schema) for split in ('train', 'validation')} |
| counts = {'input': 0, 'duplicate': 0, 'spam_host': 0, 'spam_farm': 0, 'not_accepted': 0, |
| 'train': 0, 'validation': 0, 'train_han': 0, 'validation_han': 0} |
| site_cache = {} |
| source = pq.ParquetFile(STAGE_A / f'{shard}.parquet') |
| decision = 'keep' if kind == 'web' else 'review' |
| started = time.time() |
| for group in range(source.metadata.num_row_groups): |
| table = source.read_row_group(group, columns=COLUMNS + ['decision']) |
| table = table.filter(pc.equal(table['decision'], decision)).select(COLUMNS) |
| ids, host_values = table['id'].to_pylist(), table['host'].to_pylist() |
| keep, split = [], [] |
| for doc_id, host in zip(ids, host_values): |
| counts['input'] += 1 |
| if doc_id in duplicates: |
| counts['duplicate'] += 1; keep.append(False); split.append(None); continue |
| if host not in site_cache: |
| allowed = in_domains(host, allow) |
| site_cache[host] = ('spam_host' if host in hosts and not allowed else |
| 'spam_farm' if not allowed and registrable(host) in farms else None) |
| if site_cache[host]: |
| counts[site_cache[host]] += 1; keep.append(False); split.append(None); continue |
| if accept is not None and doc_id not in accept: |
| counts['not_accepted'] += 1; keep.append(False); split.append(None); continue |
| keep.append(True); split.append('validation' if is_validation(doc_id) else 'train') |
| for name in ('train', 'validation'): |
| mask = pa.array([k and s == name for k, s in zip(keep, split)]) |
| part = table.filter(mask) |
| if len(part): |
| text_bytes = pc.sum(pc.binary_length(part['text'])).as_py() |
| writers[name].write(part, text_bytes) |
| counts[name] += len(part) |
| counts[f'{name}_han'] += pc.sum(part['han']).as_py() |
| files = [] |
| for writer in writers.values(): |
| files += [p.relative_to(out).as_posix() for p in writer.finish()] |
| result = {'shard': shard, 'kind': kind, **counts, 'files': files, 'seconds': round(time.time() - started, 1)} |
| marker.parent.mkdir(parents=True, exist_ok=True) |
| marker.write_text(json.dumps(result, indent=1) + '\n', encoding='utf-8') |
| return result |
|
|
|
|
| def run_shards(args, kind): |
| out = Path(args.out) |
| shards = sorted(p.name[:-len('.parquet')] for p in STAGE_A.glob('[0-9]*_*[0-9].parquet')) |
| if len(shards) != 21: |
| raise SystemExit(f'expected 21 stage-a-v2 shards, found {len(shards)}') |
| accept = str(Path(args.accept).resolve()) if kind == 'review' else None |
| jobs = [(s, str(out), kind, args.file_bytes, args.level, accept) for s in shards] |
| total = {} |
| with ProcessPoolExecutor(args.workers) as pool: |
| for result in pool.map(export_shard, jobs): |
| print(f'{result["shard"]}: train {result["train"]:,} validation {result["validation"]:,} ' |
| f'dropped dup {result["duplicate"]:,} host {result["spam_host"]:,} farm {result["spam_farm"]:,} ' |
| f'not accepted {result["not_accepted"]:,} ({result["seconds"]}s)', flush=True) |
| for k, v in result.items(): |
| if isinstance(v, int): |
| total[k] = total.get(k, 0) + v |
| print(kind, json.dumps(total)) |
|
|
|
|
| def cmd_chat(args): |
| import pyarrow as pa |
| import pyarrow.parquet as pq |
| out = Path(args.out) |
| marker = out / 'parts' / 'chat.json' |
| if marker.exists(): |
| print('chat already exported'); return |
| result = {} |
| for split in ('train', 'validation', 'test'): |
| target = out / split / f'chat-lccc-{split}.parquet' |
| target.parent.mkdir(parents=True, exist_ok=True) |
| rows = {'text': [], 'dialogue': [], 'turn': [], 'script': []} |
| with open(LCCC / f'{split}.jsonl', encoding='utf-8') as stream: |
| for line in stream: |
| r = json.loads(line) |
| for k in rows: |
| rows[k].append(r[k]) |
| part = target.with_name(target.name + '.part') |
| pq.write_table(pa.table(rows), part, compression='zstd', compression_level=args.level, row_group_size=200_000) |
| os.replace(part, target) |
| result[split] = {'utterances': len(rows['text']), 'chars': sum(map(len, rows['text'])), |
| 'file': target.relative_to(out).as_posix()} |
| print(split, result[split], flush=True) |
| marker.parent.mkdir(parents=True, exist_ok=True) |
| marker.write_text(json.dumps(result, indent=1) + '\n', encoding='utf-8') |
|
|
|
|
| def cmd_finish(args): |
| out = Path(args.out) |
| parts = {p.stem: json.loads(p.read_text(encoding='utf-8')) for p in sorted((out / 'parts').glob('*.json'))} |
| files = sorted(p for p in out.glob('*/*.parquet')) |
| if list(out.glob('*/*.part')): |
| raise SystemExit('unfinished .part files remain') |
| lines = [f'{sha256(p)} {p.relative_to(out).as_posix()}' for p in files] |
| (out / 'SHA256SUMS').write_text('\n'.join(lines) + '\n', encoding='utf-8', newline='\n') |
| totals = {} |
| for name, r in parts.items(): |
| kind = r.get('kind', 'chat') |
| if kind == 'chat': |
| totals['chat'] = {s: v['utterances'] for s, v in r.items()} |
| continue |
| t = totals.setdefault(kind, {}) |
| for k, v in r.items(): |
| if isinstance(v, int): |
| t[k] = t.get(k, 0) + v |
| manifest = { |
| 'version': VERSION, 'created': time.strftime('%Y-%m-%dT%H:%M:%S%z'), |
| 'inputs': {'stage_a_run_sha256': sha256(STAGE_A / 'run.json'), |
| 'spam_hosts_sha256': sha256(STAGE_A / 'spam-hosts.json'), |
| 'duplicates_sha256': sha256(STAGE_A / 'duplicates.parquet'), |
| 'terms_sha256': sha256(TERMS), 'lccc_manifest_sha256': sha256(LCCC / 'manifest.json'), |
| 'script_sha256': sha256(Path(__file__))}, |
| 'validation_rate': VALIDATION_RATE, 'columns': {'web/review': COLUMNS, 'chat': ['text', 'dialogue', 'turn', 'script']}, |
| 'totals': totals, 'files': len(files), 'bytes': sum(p.stat().st_size for p in files), |
| } |
| (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 ('totals', 'files', 'bytes')}, ensure_ascii=False, indent=1)) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__.split('\n')[0]) |
| commands = parser.add_subparsers(dest='command', required=True) |
| for name in ('web', 'review', 'chat', 'finish'): |
| p = commands.add_parser(name) |
| p.add_argument('out') |
| p.add_argument('--level', type=int, default=9) |
| p.add_argument('--file-bytes', type=int, default=4 << 30, help='uncompressed text bytes per output file') |
| p.add_argument('--workers', type=int, default=8) |
| if name == 'review': |
| p.add_argument('--accept', required=True) |
| args = parser.parse_args() |
| if args.command in ('web', 'review'): |
| run_shards(args, args.command) |
| else: |
| {'chat': cmd_chat, 'finish': cmd_finish}[args.command](args) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
