# /// 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()