Release selected document-purpose classifier and inference example

3055b70 verified about 2 months ago

15.9 kB

# /// script
# requires-python = ">=3.12"
# dependencies = [
#   "setfit==1.2.0", "sentence-transformers==6.0.1", "transformers==5.16.1",
#   "datasets==5.0.1", "torch==2.14.0", "huggingface-hub==1.30.0",
#   "scikit-learn==1.9.0", "duckdb==1.5.5", "pyarrow==25.0.1",
# ]
# ///
"""Classify a reproducible sample on Jobs; save one Parquet per source file.
The sampling manifest is prepared once from public Parquet footers. It contains
row-group ranges and per-file quotas, not document text. Outputs stay private.
"""

from collections import Counter, deque
from concurrent.futures import ThreadPoolExecutor
import hashlib
import json
import os
from pathlib import Path
import random
import threading
import time
from urllib.parse import urlsplit

import duckdb
from huggingface_hub import (
    HfApi,
    CommitOperationAdd,
    hf_hub_download,
    snapshot_download,
)
import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
from setfit import SetFitModel
import torch
from transformers import AutoTokenizer

REPO = "davanstrien/finepdfs-doctype-experiments-20260908"
PREFIX = "one-percent-20260909"
LABELS = [
    "administration_policy",
    "exercise_assessment",
    "instruction_reference",
    "news_promotion",
    "other",
    "research_analysis",
]
OUT = Path("outputs")
LOCAL = threading.local()
SCHEMA = pa.schema(
    [
        ("id", pa.string()),
        ("source_file", pa.string()),
        ("file_row_number", pa.int64()),
        ("source_row_group", pa.int64()),
        ("url", pa.string()),
        ("host", pa.string()),
        ("full_doc_lid", pa.string()),
        ("per_page_languages", pa.list_(pa.string())),
        ("is_truncated", pa.bool_()),
        ("token_count", pa.int64()),
        ("eligible", pa.bool_()),
        ("rejection_reasons", pa.list_(pa.string())),
        ("text_sha256", pa.string()),
        ("prepared_sha256", pa.string()),
        ("original_preparation_tokens", pa.int64()),
        ("excerpt_mode", pa.string()),
        ("label", pa.string()),
        ("scores", pa.list_(pa.float64())),
    ]
)


def sha(data):
    return hashlib.sha256(data).hexdigest()


def save(name, value):
    path = OUT / name
    path.write_text(json.dumps(value, indent=2) + "\n")
    return path


def prepare(text, tokenizer):
    tokens = tokenizer.encode(text, add_special_tokens=False, truncation=False, verbose=False)
    excerpt, mode = text, "full"
    if len(tokens) > 990:
        midpoint = (len(tokens) - 330) // 2
        windows = [tokens[:330], tokens[midpoint : midpoint + 330], tokens[-330:]]
        excerpt = "\n[...]\n".join(
            tokenizer.decode(part, skip_special_tokens=True, clean_up_tokenization_spaces=False)
            for part in windows
        )
        mode = "beginning_middle_end_330_each"
    prepared = "classification: " + excerpt
    return prepared, dict(
        prepared_sha256=sha(prepared.encode()),
        original_preparation_tokens=len(tokens),
        excerpt_mode=mode,
    )


def rejection_reasons(row):
    try:
        host = (urlsplit(row["url"]).hostname or "").lower().removeprefix("www.")
    except (TypeError, ValueError):
        host = ""
    reasons = []
    if (
        not isinstance(row["id"], str)
        or not row["id"]
        or not host
        or not isinstance(row["text"], str)
    ):
        reasons.append("invalid_identity_host_text")
    if not isinstance(row["text"], str) or len(row["text"].strip()) < 150:
        reasons.append("short_text")
    if row["full_doc_lid"] != "eng_Latn":
        reasons.append("document_language")
    if not isinstance(row["per_page_languages"], list) or any(
        lang not in ("eng_Latn", "unknown") for lang in row["per_page_languages"]
    ):
        reasons.append("page_language")
    return host, reasons


def read_file(part, tokenizer_path):
    started = time.monotonic()
    if not hasattr(LOCAL, "tokenizer"):
        LOCAL.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, local_files_only=True)
    rows = []
    with duckdb.connect() as con:
        con.execute("LOAD httpfs")
        con.execute("SET threads=1")
        con.execute("SET memory_limit='1GB'")
        for group in part["groups"]:
            sql = (
                "SELECT file_row_number,id,text,url,full_doc_lid,per_page_languages,is_truncated,token_count "
                f"FROM read_parquet('{part['url']}',file_row_number=true) "
                f"WHERE file_row_number>={group['start']} AND file_row_number<{group['end']} "
                "ORDER BY file_row_number"
            )
            for attempt in range(2):
                timer = threading.Timer(90, con.interrupt)
                timer.start()
                try:
                    chunk = con.execute(sql).fetch_arrow_table().to_pylist()
                    break
                except Exception:
                    if attempt:
                        raise
                finally:
                    timer.cancel()
                    timer.join()
            assert [r["file_row_number"] for r in chunk] == list(
                range(group["start"], group["end"])
            )
            for row in chunk:
                row["source_row_group"] = group["row_group_id"]
            rows.extend(chunk)
    read_seconds = time.monotonic() - started
    rows = sorted(
        random.Random(part["seed"]).sample(rows, part["quota"]),
        key=lambda r: r["file_row_number"],
    )
    texts = []
    for row in rows:
        row["host"], row["rejection_reasons"] = rejection_reasons(row)
        row["eligible"] = not row["rejection_reasons"]
        raw = row.pop("text")
        row.update(
            source_file=part["path"],
            text_sha256=sha(raw.encode()) if isinstance(raw, str) else None,
        )
        if row["eligible"]:
            prepared, provenance = prepare(raw, LOCAL.tokenizer)
            texts.append(prepared)
            row.update(provenance)
    return (
        rows,
        texts,
        dict(read_seconds=read_seconds, read_prepare_seconds=time.monotonic() - started),
    )


def predict(model, texts):
    if not texts:
        return np.empty((0, len(LABELS)))
    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
        vectors = model.model_body.encode(
            texts,
            batch_size=16,
            normalize_embeddings=False,
            show_progress_bar=False,
            convert_to_tensor=True,
        )
        scores = model.model_head.predict_proba(vectors.float().cpu().numpy())
    assert scores.shape == (len(texts), len(LABELS)) and np.isfinite(scores).all()
    assert np.allclose(scores.sum(axis=1), 1, atol=1e-6)
    return scores


def main():
    started = time.monotonic()
    OUT.mkdir(exist_ok=True)
    api = HfApi()
    assert api.dataset_info(REPO).private
    revision = os.environ["ONE_PERCENT_INPUT_REVISION"]
    inputs = {}
    for name in ("sample.json", "config.json"):
        path = hf_hub_download(
            REPO, f"{PREFIX}/inputs/{name}", repo_type="dataset", revision=revision
        )
        inputs[name] = json.loads(Path(path).read_text())
    sample, config = inputs["sample.json"], inputs["config.json"]
    assert sha(Path(__file__).read_bytes()) == config["script_sha256"]
    fingerprint = sha(json.dumps(inputs, sort_keys=True).encode())
    assert len(sample["files"]) == 100
    assert sum(p["source_rows"] for p in sample["files"]) == sample["source_rows"] == 23023372
    assert len({p["path"] for p in sample["files"]}) == 100
    assert sum(p["quota"] for p in sample["files"]) == sample["target_rows"] == 230234
    # One completion record per file. A partial commit is never treated as done.
    output_revision = api.dataset_info(REPO).sha
    existing = set(api.list_repo_files(REPO, repo_type="dataset", revision=output_revision))
    completed, pending = {}, []
    for part in sample["files"]:
        stem = Path(part["path"]).stem
        if f"{PREFIX}/outputs/{stem}.json" in existing:
            meta = json.loads(
                Path(
                    hf_hub_download(
                        REPO,
                        f"{PREFIX}/outputs/{stem}.json",
                        repo_type="dataset",
                        revision=output_revision,
                    )
                ).read_text()
            )
            data = Path(
                hf_hub_download(
                    REPO,
                    f"{PREFIX}/outputs/{stem}.parquet",
                    repo_type="dataset",
                    revision=output_revision,
                )
            )
            assert meta["fingerprint"] == fingerprint and meta["sampled_rows"] == part["quota"]
            assert (
                sha(data.read_bytes()) == meta["sha256"]
                and pq.read_metadata(data).num_rows == part["quota"]
            )
            completed[stem] = meta
        else:
            pending.append(part)
    model_root = snapshot_download(
        REPO,
        repo_type="dataset",
        revision=config["model_revision"],
        allow_patterns=[f"{config['model_path']}/*"],
    )
    model = SetFitModel.from_pretrained(
        str(Path(model_root) / config["model_path"]),
        device="cuda",
        local_files_only=True,
    )
    model.model_body.float().eval()
    model.model_body.max_seq_length = 1024
    model.model_body.default_prompt_name = None
    model.normalize_embeddings = False
    torch.set_float32_matmul_precision("high")
    assert torch.cuda.is_bf16_supported() and model.model_head.classes_.tolist() == LABELS
    digest = hashlib.sha256()
    for name, parameter in model.model_body.named_parameters():
        digest.update(name.encode())
        digest.update(str(parameter.dtype).encode())
        digest.update(parameter.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes())
    assert digest.hexdigest() == config["encoder_sha256"]
    tokenizer_path = snapshot_download(
        config["tokenizer"],
        revision=config["tokenizer_revision"],
        allow_patterns=["*.json", "*.txt", "*.model"],
    )
    # A small saved-reference check catches preprocessing or model drift at startup.
    fixture = []
    for name in ("pool.parquet", "predictions-fp32.parquet"):
        path = hf_hub_download(
            REPO,
            f"inference-pilot-20260908/{name}",
            repo_type="dataset",
            revision=config["reference_revision"],
        )
        assert sha(Path(path).read_bytes()) == config["reference_hashes"][name]
        fixture.append(pq.read_table(path).slice(0, 120).to_pylist())
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, local_files_only=True)
    prepared = [prepare(row["text"], tokenizer) for row in fixture[0]]
    assert [p[1]["prepared_sha256"] for p in prepared] == [r["prepared_sha256"] for r in fixture[1]]
    reference = np.array([r["scores"] for r in fixture[1]])
    actual = predict(model, [p[0] for p in prepared])
    parity = dict(
        label_agreement=float(np.mean(actual.argmax(1) == reference.argmax(1))),
        maximum_score_delta=float(np.abs(actual - reference).max()),
    )
    assert parity["label_agreement"] >= 0.99 and parity["maximum_score_delta"] < 0.01
    print(
        json.dumps(dict(event="startup_passed", parity=parity, remaining_files=len(pending))),
        flush=True,
    )
    with duckdb.connect() as con:
        con.execute("INSTALL httpfs FROM 'https://extensions.duckdb.org'")
    setup_seconds = time.monotonic() - started
    operations = []
    upload_seconds = 0.0

    def flush():
        nonlocal upload_seconds
        if operations:
            upload_start = time.monotonic()
            api.create_commit(
                repo_id=REPO,
                repo_type="dataset",
                operations=operations,
                commit_message="Save completed 1% inference files",
            )
            upload_seconds += time.monotonic() - upload_start
            operations.clear()

    # Only two files are prepared ahead; inference and commits run on this thread.
    with ThreadPoolExecutor(max_workers=2) as pool:
        parts = iter(pending)
        queue = deque(
            (part, pool.submit(read_file, part, tokenizer_path)) for part in list(pending)[:2]
        )
        for _ in range(min(2, len(pending))):
            next(parts)
        try:
            while queue and time.monotonic() - started < config["soft_limit_seconds"]:
                part, future = queue.popleft()
                rows, texts, timing = future.result()
                infer_start = time.monotonic()
                probabilities = iter(predict(model, texts))
                for row in rows:
                    if row["eligible"]:
                        scores = next(probabilities)
                        row.update(label=LABELS[int(scores.argmax())], scores=scores.tolist())
                timing["inference_seconds"] = time.monotonic() - infer_start
                stem = Path(part["path"]).stem
                path = OUT / f"{stem}.parquet"
                pq.write_table(pa.Table.from_pylist(rows, schema=SCHEMA), path, compression="zstd")
                meta = dict(
                    fingerprint=fingerprint,
                    sampled_rows=len(rows),
                    predicted_rows=len(texts),
                    source_file=part["path"],
                    source_revision=sample["source_revision"],
                    sha256=sha(path.read_bytes()),
                    **timing,
                    label_counts=dict(Counter(r["label"] for r in rows if r["eligible"])),
                    rejection_counts=dict(
                        Counter(reason for r in rows for reason in r["rejection_reasons"])
                    ),
                )
                for file in (path, save(f"{stem}.json", meta)):
                    operations.append(
                        CommitOperationAdd(
                            path_in_repo=f"{PREFIX}/outputs/{file.name}",
                            path_or_fileobj=file,
                        )
                    )
                if len(operations) >= 6:
                    flush()
                completed[stem] = meta
                print(
                    json.dumps(
                        dict(
                            event="file_complete",
                            file=stem,
                            files=len(completed),
                            predicted_rows=len(texts),
                            seconds=time.monotonic() - started,
                        )
                    ),
                    flush=True,
                )
                following = next(parts, None)
                if (
                    following is not None
                    and time.monotonic() - started < config["soft_limit_seconds"]
                ):
                    queue.append((following, pool.submit(read_file, following, tokenizer_path)))
        finally:
            flush()
    summary = dict(
        status="complete" if len(completed) == len(sample["files"]) else "partial",
        files=len(completed),
        sampled_rows=sum(x["sampled_rows"] for x in completed.values()),
        predicted_rows=sum(x["predicted_rows"] for x in completed.values()),
        source_revision=sample["source_revision"],
        fingerprint=fingerprint,
        label_order=LABELS,
        setup_seconds=setup_seconds,
        upload_seconds=upload_seconds,
        parity=parity,
        process_seconds=time.monotonic() - started,
    )
    if summary["status"] == "complete":
        assert summary["sampled_rows"] == sample["target_rows"]
    api.upload_file(
        repo_id=REPO,
        repo_type="dataset",
        path_or_fileobj=save("summary.json", summary),
        path_in_repo=f"{PREFIX}/summary.json",
    )
    print(json.dumps(summary), flush=True)


if __name__ == "__main__":
    main()

	# /// script
	# requires-python = ">=3.12"
	# dependencies = [
	# "setfit==1.2.0", "sentence-transformers==6.0.1", "transformers==5.16.1",
	# "datasets==5.0.1", "torch==2.14.0", "huggingface-hub==1.30.0",
	# "scikit-learn==1.9.0", "duckdb==1.5.5", "pyarrow==25.0.1",
	# ]
	# ///
	"""Classify a reproducible sample on Jobs; save one Parquet per source file.

	The sampling manifest is prepared once from public Parquet footers. It contains
	row-group ranges and per-file quotas, not document text. Outputs stay private.
	"""

	from collections import Counter, deque
	from concurrent.futures import ThreadPoolExecutor
	import hashlib
	import json
	import os
	from pathlib import Path
	import random
	import threading
	import time
	from urllib.parse import urlsplit

	import duckdb
	from huggingface_hub import (
	HfApi,
	CommitOperationAdd,
	hf_hub_download,
	snapshot_download,
	)
	import numpy as np
	import pyarrow as pa
	import pyarrow.parquet as pq
	from setfit import SetFitModel
	import torch
	from transformers import AutoTokenizer

	REPO = "davanstrien/finepdfs-doctype-experiments-20260908"
	PREFIX = "one-percent-20260909"
	LABELS = [
	"administration_policy",
	"exercise_assessment",
	"instruction_reference",
	"news_promotion",
	"other",
	"research_analysis",
	]
	OUT = Path("outputs")
	LOCAL = threading.local()
	SCHEMA = pa.schema(
	[
	("id", pa.string()),
	("source_file", pa.string()),
	("file_row_number", pa.int64()),
	("source_row_group", pa.int64()),
	("url", pa.string()),
	("host", pa.string()),
	("full_doc_lid", pa.string()),
	("per_page_languages", pa.list_(pa.string())),
	("is_truncated", pa.bool_()),
	("token_count", pa.int64()),
	("eligible", pa.bool_()),
	("rejection_reasons", pa.list_(pa.string())),
	("text_sha256", pa.string()),
	("prepared_sha256", pa.string()),
	("original_preparation_tokens", pa.int64()),
	("excerpt_mode", pa.string()),
	("label", pa.string()),
	("scores", pa.list_(pa.float64())),
	]
	)


	def sha(data):
	return hashlib.sha256(data).hexdigest()


	def save(name, value):
	path = OUT / name
	path.write_text(json.dumps(value, indent=2) + "\n")
	return path


	def prepare(text, tokenizer):
	tokens = tokenizer.encode(text, add_special_tokens=False, truncation=False, verbose=False)
	excerpt, mode = text, "full"
	if len(tokens) > 990:
	midpoint = (len(tokens) - 330) // 2
	windows = [tokens[:330], tokens[midpoint : midpoint + 330], tokens[-330:]]
	excerpt = "\n[...]\n".join(
	tokenizer.decode(part, skip_special_tokens=True, clean_up_tokenization_spaces=False)
	for part in windows
	)
	mode = "beginning_middle_end_330_each"
	prepared = "classification: " + excerpt
	return prepared, dict(
	prepared_sha256=sha(prepared.encode()),
	original_preparation_tokens=len(tokens),
	excerpt_mode=mode,
	)


	def rejection_reasons(row):
	try:
	host = (urlsplit(row["url"]).hostname or "").lower().removeprefix("www.")
	except (TypeError, ValueError):
	host = ""
	reasons = []
	if (
	not isinstance(row["id"], str)
	or not row["id"]
	or not host
	or not isinstance(row["text"], str)
	):
	reasons.append("invalid_identity_host_text")
	if not isinstance(row["text"], str) or len(row["text"].strip()) < 150:
	reasons.append("short_text")
	if row["full_doc_lid"] != "eng_Latn":
	reasons.append("document_language")
	if not isinstance(row["per_page_languages"], list) or any(
	lang not in ("eng_Latn", "unknown") for lang in row["per_page_languages"]
	):
	reasons.append("page_language")
	return host, reasons


	def read_file(part, tokenizer_path):
	started = time.monotonic()
	if not hasattr(LOCAL, "tokenizer"):
	LOCAL.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, local_files_only=True)
	rows = []
	with duckdb.connect() as con:
	con.execute("LOAD httpfs")
	con.execute("SET threads=1")
	con.execute("SET memory_limit='1GB'")
	for group in part["groups"]:
	sql = (
	"SELECT file_row_number,id,text,url,full_doc_lid,per_page_languages,is_truncated,token_count "
	f"FROM read_parquet('{part['url']}',file_row_number=true) "
	f"WHERE file_row_number>={group['start']} AND file_row_number<{group['end']} "
	"ORDER BY file_row_number"
	)
	for attempt in range(2):
	timer = threading.Timer(90, con.interrupt)
	timer.start()
	try:
	chunk = con.execute(sql).fetch_arrow_table().to_pylist()
	break
	except Exception:
	if attempt:
	raise
	finally:
	timer.cancel()
	timer.join()
	assert [r["file_row_number"] for r in chunk] == list(
	range(group["start"], group["end"])
	)
	for row in chunk:
	row["source_row_group"] = group["row_group_id"]
	rows.extend(chunk)
	read_seconds = time.monotonic() - started
	rows = sorted(
	random.Random(part["seed"]).sample(rows, part["quota"]),
	key=lambda r: r["file_row_number"],
	)
	texts = []
	for row in rows:
	row["host"], row["rejection_reasons"] = rejection_reasons(row)
	row["eligible"] = not row["rejection_reasons"]
	raw = row.pop("text")
	row.update(
	source_file=part["path"],
	text_sha256=sha(raw.encode()) if isinstance(raw, str) else None,
	)
	if row["eligible"]:
	prepared, provenance = prepare(raw, LOCAL.tokenizer)
	texts.append(prepared)
	row.update(provenance)
	return (
	rows,
	texts,
	dict(read_seconds=read_seconds, read_prepare_seconds=time.monotonic() - started),
	)


	def predict(model, texts):
	if not texts:
	return np.empty((0, len(LABELS)))
	with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
	vectors = model.model_body.encode(
	texts,
	batch_size=16,
	normalize_embeddings=False,
	show_progress_bar=False,
	convert_to_tensor=True,
	)
	scores = model.model_head.predict_proba(vectors.float().cpu().numpy())
	assert scores.shape == (len(texts), len(LABELS)) and np.isfinite(scores).all()
	assert np.allclose(scores.sum(axis=1), 1, atol=1e-6)
	return scores


	def main():
	started = time.monotonic()
	OUT.mkdir(exist_ok=True)
	api = HfApi()
	assert api.dataset_info(REPO).private
	revision = os.environ["ONE_PERCENT_INPUT_REVISION"]
	inputs = {}
	for name in ("sample.json", "config.json"):
	path = hf_hub_download(
	REPO, f"{PREFIX}/inputs/{name}", repo_type="dataset", revision=revision
	)
	inputs[name] = json.loads(Path(path).read_text())
	sample, config = inputs["sample.json"], inputs["config.json"]
	assert sha(Path(__file__).read_bytes()) == config["script_sha256"]
	fingerprint = sha(json.dumps(inputs, sort_keys=True).encode())
	assert len(sample["files"]) == 100
	assert sum(p["source_rows"] for p in sample["files"]) == sample["source_rows"] == 23023372
	assert len({p["path"] for p in sample["files"]}) == 100
	assert sum(p["quota"] for p in sample["files"]) == sample["target_rows"] == 230234
	# One completion record per file. A partial commit is never treated as done.
	output_revision = api.dataset_info(REPO).sha
	existing = set(api.list_repo_files(REPO, repo_type="dataset", revision=output_revision))
	completed, pending = {}, []
	for part in sample["files"]:
	stem = Path(part["path"]).stem
	if f"{PREFIX}/outputs/{stem}.json" in existing:
	meta = json.loads(
	Path(
	hf_hub_download(
	REPO,
	f"{PREFIX}/outputs/{stem}.json",
	repo_type="dataset",
	revision=output_revision,
	)
	).read_text()
	)
	data = Path(
	hf_hub_download(
	REPO,
	f"{PREFIX}/outputs/{stem}.parquet",
	repo_type="dataset",
	revision=output_revision,
	)
	)
	assert meta["fingerprint"] == fingerprint and meta["sampled_rows"] == part["quota"]
	assert (
	sha(data.read_bytes()) == meta["sha256"]
	and pq.read_metadata(data).num_rows == part["quota"]
	)
	completed[stem] = meta
	else:
	pending.append(part)
	model_root = snapshot_download(
	REPO,
	repo_type="dataset",
	revision=config["model_revision"],
	allow_patterns=[f"{config['model_path']}/*"],
	)
	model = SetFitModel.from_pretrained(
	str(Path(model_root) / config["model_path"]),
	device="cuda",
	local_files_only=True,
	)
	model.model_body.float().eval()
	model.model_body.max_seq_length = 1024
	model.model_body.default_prompt_name = None
	model.normalize_embeddings = False
	torch.set_float32_matmul_precision("high")
	assert torch.cuda.is_bf16_supported() and model.model_head.classes_.tolist() == LABELS
	digest = hashlib.sha256()
	for name, parameter in model.model_body.named_parameters():
	digest.update(name.encode())
	digest.update(str(parameter.dtype).encode())
	digest.update(parameter.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes())
	assert digest.hexdigest() == config["encoder_sha256"]
	tokenizer_path = snapshot_download(
	config["tokenizer"],
	revision=config["tokenizer_revision"],
	allow_patterns=[".json", ".txt", "*.model"],
	)
	# A small saved-reference check catches preprocessing or model drift at startup.
	fixture = []
	for name in ("pool.parquet", "predictions-fp32.parquet"):
	path = hf_hub_download(
	REPO,
	f"inference-pilot-20260908/{name}",
	repo_type="dataset",
	revision=config["reference_revision"],
	)
	assert sha(Path(path).read_bytes()) == config["reference_hashes"][name]
	fixture.append(pq.read_table(path).slice(0, 120).to_pylist())
	tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, local_files_only=True)
	prepared = [prepare(row["text"], tokenizer) for row in fixture[0]]
	assert [p[1]["prepared_sha256"] for p in prepared] == [r["prepared_sha256"] for r in fixture[1]]
	reference = np.array([r["scores"] for r in fixture[1]])
	actual = predict(model, [p[0] for p in prepared])
	parity = dict(
	label_agreement=float(np.mean(actual.argmax(1) == reference.argmax(1))),
	maximum_score_delta=float(np.abs(actual - reference).max()),
	)
	assert parity["label_agreement"] >= 0.99 and parity["maximum_score_delta"] < 0.01
	print(
	json.dumps(dict(event="startup_passed", parity=parity, remaining_files=len(pending))),
	flush=True,
	)
	with duckdb.connect() as con:
	con.execute("INSTALL httpfs FROM 'https://extensions.duckdb.org'")
	setup_seconds = time.monotonic() - started
	operations = []
	upload_seconds = 0.0

	def flush():
	nonlocal upload_seconds
	if operations:
	upload_start = time.monotonic()
	api.create_commit(
	repo_id=REPO,
	repo_type="dataset",
	operations=operations,
	commit_message="Save completed 1% inference files",
	)
	upload_seconds += time.monotonic() - upload_start
	operations.clear()

	# Only two files are prepared ahead; inference and commits run on this thread.
	with ThreadPoolExecutor(max_workers=2) as pool:
	parts = iter(pending)
	queue = deque(
	(part, pool.submit(read_file, part, tokenizer_path)) for part in list(pending)[:2]
	)
	for _ in range(min(2, len(pending))):
	next(parts)
	try:
	while queue and time.monotonic() - started < config["soft_limit_seconds"]:
	part, future = queue.popleft()
	rows, texts, timing = future.result()
	infer_start = time.monotonic()
	probabilities = iter(predict(model, texts))
	for row in rows:
	if row["eligible"]:
	scores = next(probabilities)
	row.update(label=LABELS[int(scores.argmax())], scores=scores.tolist())
	timing["inference_seconds"] = time.monotonic() - infer_start
	stem = Path(part["path"]).stem
	path = OUT / f"{stem}.parquet"
	pq.write_table(pa.Table.from_pylist(rows, schema=SCHEMA), path, compression="zstd")
	meta = dict(
	fingerprint=fingerprint,
	sampled_rows=len(rows),
	predicted_rows=len(texts),
	source_file=part["path"],
	source_revision=sample["source_revision"],
	sha256=sha(path.read_bytes()),
	**timing,
	label_counts=dict(Counter(r["label"] for r in rows if r["eligible"])),
	rejection_counts=dict(
	Counter(reason for r in rows for reason in r["rejection_reasons"])
	),
	)
	for file in (path, save(f"{stem}.json", meta)):
	operations.append(
	CommitOperationAdd(
	path_in_repo=f"{PREFIX}/outputs/{file.name}",
	path_or_fileobj=file,
	)
	)
	if len(operations) >= 6:
	flush()
	completed[stem] = meta
	print(
	json.dumps(
	dict(
	event="file_complete",
	file=stem,
	files=len(completed),
	predicted_rows=len(texts),
	seconds=time.monotonic() - started,
	)
	),
	flush=True,
	)
	following = next(parts, None)
	if (
	following is not None
	and time.monotonic() - started < config["soft_limit_seconds"]
	):
	queue.append((following, pool.submit(read_file, following, tokenizer_path)))
	finally:
	flush()
	summary = dict(
	status="complete" if len(completed) == len(sample["files"]) else "partial",
	files=len(completed),
	sampled_rows=sum(x["sampled_rows"] for x in completed.values()),
	predicted_rows=sum(x["predicted_rows"] for x in completed.values()),
	source_revision=sample["source_revision"],
	fingerprint=fingerprint,
	label_order=LABELS,
	setup_seconds=setup_seconds,
	upload_seconds=upload_seconds,
	parity=parity,
	process_seconds=time.monotonic() - started,
	)
	if summary["status"] == "complete":
	assert summary["sampled_rows"] == sample["target_rows"]
	api.upload_file(
	repo_id=REPO,
	repo_type="dataset",
	path_or_fileobj=save("summary.json", summary),
	path_in_repo=f"{PREFIX}/summary.json",
	)
	print(json.dumps(summary), flush=True)


	if __name__ == "__main__":
	main()