"""Run with ``python -m app.ingest --help`` from backend/.""" import argparse import asyncio import hashlib import json import os from pathlib import Path from app.ingest.fetch import atomic_json, fetch_pages from app.ingest.scope import PROFILES def read_jsonl(path: Path): with path.open() as handle: for line_number, line in enumerate(handle, 1): if not line.strip(): continue try: row = json.loads(line) if not isinstance(row, dict): raise TypeError("record is not an object") yield row except (json.JSONDecodeError, ValueError, TypeError) as exc: raise ValueError(f"Invalid JSONL at {path.name}:{line_number}") from exc def transform_file(source: Path, output: Path, transform, *, resume: bool, stage: str, batch_size: int = 100) -> dict: """Transform batches with output truncation and an input-record checkpoint.""" if source.resolve() == output.resolve(): raise ValueError("Input and output paths must differ") signature = hashlib.sha256(json.dumps([str(source.resolve()), source.stat().st_size, source.stat().st_mtime_ns, stage]).encode()).hexdigest() checkpoint = output.with_suffix(output.suffix + ".checkpoint.json") state = {"signature": signature, "input_records": 0, "output_records": 0, "bytes": 0, "done": False} if resume and checkpoint.exists(): state = json.loads(checkpoint.read_text()) if state.get("signature") != signature: raise ValueError("Input or transformation changed since checkpoint") elif output.exists(): raise FileExistsError("Output already exists; use --resume or choose a new path") output.parent.mkdir(parents=True, exist_ok=True) if not output.exists(): if state["bytes"]: raise ValueError("Checkpoint output is missing") output.touch() if output.stat().st_size < state["bytes"]: raise ValueError("Checkpoint output is truncated") if output.stat().st_size != state["bytes"]: with output.open("r+b") as handle: handle.truncate(state["bytes"]) atomic_json(checkpoint, state) if state["done"]: return state skipped = state["input_records"] batch = [] def commit(rows): result = transform(rows) with output.open("ab") as handle: for row in result: handle.write((json.dumps(row, ensure_ascii=False) + "\n").encode()) handle.flush() os.fsync(handle.fileno()) state["bytes"] = handle.tell() state["input_records"] += len(rows) state["output_records"] += len(result) atomic_json(checkpoint, state) for index, row in enumerate(read_jsonl(source)): if index < skipped: continue batch.append(row) if len(batch) >= batch_size: commit(batch) batch = [] if batch: commit(batch) state["done"] = True atomic_json(checkpoint, state) return state def parser() -> argparse.ArgumentParser: root = argparse.ArgumentParser(description="NullMap resumable ingestion; all intermediate stages use JSONL") commands = root.add_subparsers(dest="command", required=True) fetch = commands.add_parser("fetch", help="Download ClinicalTrials.gov records") fetch.add_argument("source", choices=["ctgov"]) fetch.add_argument("--output", type=Path, required=True) fetch.add_argument("--query", default="") fetch.add_argument("--filter", default=None) fetch.add_argument("--limit", type=int, default=1000) fetch.add_argument("--all", action="store_true", help="Continue until the registry cursor is exhausted") fetch.add_argument("--page-size", type=int, default=1000) fetch.add_argument("--resume", action="store_true") for name in ("normalize", "classify", "embed"): command = commands.add_parser(name) command.add_argument("--input", type=Path, required=True) command.add_argument("--output", type=Path, required=True) command.add_argument("--resume", action="store_true") command.add_argument("--batch-size", type=int, default=100) if name == "classify": command.add_argument("--model", type=Path) link = commands.add_parser("link", help="Merge normalized JSONL files by PMID and NCT") link.add_argument("--input", type=Path, nargs="+", required=True) link.add_argument("--output", type=Path, required=True) link.add_argument("--resume", action="store_true") train = commands.add_parser("train-classifier") train.add_argument("--input", type=Path, required=True) train.add_argument("--output", type=Path, required=True) train.add_argument("--label-field", default="human_label") label = commands.add_parser("label", help="Create instrumented structured LLM labels for training") label.add_argument("--input", type=Path, required=True) label.add_argument("--output", type=Path, required=True) label.add_argument("--limit", type=int, default=100) label.add_argument("--batch-size", type=int, default=10) label.add_argument("--resume", action="store_true") index = commands.add_parser("index") index.add_argument("--input", type=Path, required=True) index.add_argument("--batch-size", type=int, default=200) index.add_argument("--resume", action="store_true") index.add_argument("--sesoi", type=float, default=0.2) index.add_argument("--effect-type", default="SMD") pagerank = commands.add_parser( "pagerank", help="Recompute within-corpus citation PageRank and update authority rank features", ) pagerank.add_argument("--damping", type=float, default=0.85) pagerank.add_argument("--iterations", type=int, default=30) pagerank.add_argument("--batch-size", type=int, default=500) references = commands.add_parser("reference-ids", help="Collect missing review-reference IDs for S3 backfill") references.add_argument("--input", type=Path, nargs="+", required=True) references.add_argument("--output", type=Path, required=True) tei = commands.add_parser("tei", help="Extract null-hypothesis evidence from GROBID TEI XML full text") tei.add_argument("--input", type=Path, nargs="+", required=True, help="TEI .xml files or directories") tei.add_argument("--output", type=Path, required=True) tei.add_argument("--source", default="openalex") snapshot = commands.add_parser("snapshot", help="Read the public OpenAlex S3 Parquet snapshot") snapshot.add_argument("--input", nargs="+", default=[]) snapshot.add_argument("--manifest", help="Public S3 manifest URL or a pinned local manifest; defaults to public works") snapshot.add_argument("--max-files", type=int, help="Select at most this many manifest files before enforcing the byte budget") snapshot.add_argument("--largest-first", action="store_true", help="Order parts by descending size; a partial run then covers more of the corpus") snapshot.add_argument("--skip-files", type=int, default=0, help="Skip this many leading manifest files, to continue after an earlier --max-files run in a new output directory") snapshot.add_argument("--output", type=Path, default=Path("data/snapshot/works.jsonl")) snapshot.add_argument("--max-bytes", type=int, default=5_000_000_000) snapshot.add_argument("--limit", type=int, help="Optional explicit row cap; absent means exhaust selected parts") snapshot.add_argument("--topic", default="", help="Optional additional topic-name filter") snapshot.add_argument("--profile", choices=list(PROFILES), default="hypertension-kidney", help="The -pubmed profile keeps only PubMed-indexed works, which registry linking and full text need") snapshot.add_argument("--work-ids", type=Path, help="Backfill these newline-separated W IDs from S3; overrides topic/profile filtering") snapshot.add_argument("--min-free-bytes", type=int, default=1_000_000_000) snapshot.add_argument("--plan", action="store_true", help="Inspect the manifest without reading Parquet") snapshot.add_argument("--resume", action="store_true") return root async def run(args) -> dict: from app.config import settings if getattr(args, "batch_size", 1) <= 0: raise ValueError("batch-size must be positive") if args.command == "fetch": state = await fetch_pages(args.source, args.output, query=args.query, filters=args.filter, limit=None if args.all else args.limit, resume=args.resume, page_size=args.page_size) return {"records": state["count"], "source_exhausted": state["done"], "output": str(args.output)} if args.command == "label": from app.ingest.label import label_file return await label_file(args.input, args.output, limit=args.limit, batch_size=args.batch_size, resume=args.resume) if args.command == "normalize": from app.ingest import flatten_trial, normalize_work def transform(rows): return [result for row in rows if (result := ( flatten_trial(row) if "protocolSection" in row else normalize_work(row))) is not None] elif args.command == "classify": from app.classifier import ResultClassifier classifier = ResultClassifier(args.model or settings.classifier_model_path or None) def transform(rows): return [{**row, **classifier.classify(row)} if row.get("source") == "openalex" else row for row in rows] elif args.command == "embed": from app.embeddings import get_embedder embedder = get_embedder(settings.embedding_model, settings.embedding_device) def transform(rows): texts = [" ".join(filter(None, [row.get("title", ""), row.get("abstract", ""), row.get("intervention", ""), row.get("outcome", "")])) for row in rows] vectors = embedder.encode(texts) return [{**row, "embedding": vector, "embedding_model": settings.embedding_model} for row, vector in zip(rows, vectors, strict=True)] elif args.command == "link": from app.ingest.linker import link_studies signature = hashlib.sha256(json.dumps([ [str(path.resolve()), path.stat().st_size, path.stat().st_mtime_ns] for path in args.input ]).encode()).hexdigest() checkpoint = args.output.with_suffix(args.output.suffix + ".checkpoint.json") if args.output.exists(): if args.resume and checkpoint.exists(): state = json.loads(checkpoint.read_text()) if state.get("signature") == signature and state.get("done"): return state raise FileExistsError("Link output exists with no matching completed checkpoint") rows = link_studies([row for path in args.input for row in read_jsonl(path)]) args.output.parent.mkdir(parents=True, exist_ok=True) temporary = args.output.with_suffix(args.output.suffix + ".tmp") with temporary.open("w") as handle: for row in rows: handle.write(json.dumps(row) + "\n") temporary.replace(args.output) state = {"records": len(rows), "signature": signature, "done": True} atomic_json(checkpoint, state) return state elif args.command == "train-classifier": from app.classifier import train_classifier return train_classifier(list(read_jsonl(args.input)), args.output, label_field=args.label_field) elif args.command == "reference-ids": from app.ingest.openalex import normalize_work known, referenced = set(), set() for path in args.input: for row in read_jsonl(path): study = normalize_work(row) if "source" not in row else row known.add(str(study["id"]).rsplit("/", 1)[-1]) if study.get("is_review"): referenced.update(str(x).rsplit("/", 1)[-1] for x in study.get("referenced_works", [])) missing = sorted(referenced - known) args.output.parent.mkdir(parents=True, exist_ok=True) temporary = args.output.with_suffix(args.output.suffix + ".tmp") temporary.write_text("".join(identifier + "\n" for identifier in missing)) temporary.replace(args.output) return {"missing_reference_ids": len(missing), "output": str(args.output)} elif args.command == "tei": from app.ingest.tei import extract_tei_files return extract_tei_files(args.input, args.output, source=args.source) elif args.command == "snapshot": from app.ingest.snapshot import fingerprint, load_manifest, plan_manifest, scan_snapshot if args.input and args.manifest: raise ValueError("Choose --input local files or --manifest, not both") location = args.manifest or settings.openalex_snapshot_manifest if args.input: manifest = {"files": [{"url": str(Path(x).resolve()), "size_bytes": Path(x).stat().st_size} for x in args.input]} else: manifest = await load_manifest(location) plan = plan_manifest(manifest, max_files=args.max_files, largest_first=getattr(args, "largest_first", False), skip_files=getattr(args, "skip_files", 0)) if args.plan: return {k: v for k, v in plan.items() if k != "files"} pinned = args.output.with_suffix(args.output.suffix + ".manifest.json") if pinned.exists() and fingerprint(json.loads(pinned.read_text())) != plan["manifest_sha256"]: raise ValueError("The snapshot manifest changed; use a new output directory for this release") atomic_json(pinned, manifest) files = [part["url"] for part in plan["files"]] work_ids = {line.strip() for line in args.work_ids.read_text().splitlines() if line.strip()} if args.work_ids else None result = await asyncio.to_thread( scan_snapshot, files, args.output, max_bytes=args.max_bytes, limit=args.limit, topic=args.topic, profile=args.profile, remote_sizes={part["url"]: part["size_bytes"] for part in plan["files"]}, resume=args.resume, snapshot_id=plan["manifest_sha256"], snapshot_date=plan["snapshot_date"], all_parts_selected=plan["all_parts_selected"], work_ids=work_ids, min_free_bytes=args.min_free_bytes, progress=lambda state: print(json.dumps({"stage": "snapshot", "part": state["part"], "parts": state["selected_parts"], "records": state["records"]}), flush=True), ) if not args.input and "://" in location and result["selected_parts_complete"]: if fingerprint(await load_manifest(location)) != plan["manifest_sha256"]: result["complete_snapshot_scope"] = False result["manifest_changed_during_scan"] = True atomic_json(args.output.with_suffix(args.output.suffix + ".checkpoint.json"), result) raise RuntimeError("Public snapshot changed during the scan; result is not a complete release") return result elif args.command == "index": from app.repository import ElasticRepository from app.statistics import assign_bucket repository = ElasticRepository() signature = hashlib.sha256(json.dumps([str(args.input.resolve()), args.input.stat().st_size, args.input.stat().st_mtime_ns, settings.elastic_index, args.sesoi, args.effect_type]).encode()).hexdigest() checkpoint = args.input.with_suffix(args.input.suffix + ".index-checkpoint.json") state = {"signature": signature, "records": 0} if args.resume and checkpoint.exists(): state = json.loads(checkpoint.read_text()) if state.get("signature") != signature: raise ValueError("Index checkpoint does not match input or settings") skipped = state["records"] batch = [] merged = [] try: await repository.ensure_index() for position, row in enumerate(read_jsonl(args.input)): if row.get("source") == "merged" and row.get("linked_papers"): merged.append(row) if position < skipped: continue row = {**row, **assign_bucket(row, sesoi=args.sesoi, effect_type=args.effect_type)} batch.append(row) if len(batch) >= args.batch_size: await repository.bulk_upsert(batch) state["records"] += len(batch) atomic_json(checkpoint, state) batch = [] if batch: await repository.bulk_upsert(batch) state["records"] += len(batch) atomic_json(checkpoint, state) # A linked JSONL removes duplicate paper rows, but the index can still # contain papers from earlier runs. Persist canonical redirects after # every batch is durable, then checkpoint this migration separately. if state.get("links_version") != 1: await repository.persist_links(merged) state["links_version"] = 1 atomic_json(checkpoint, state) return state finally: await repository.close() elif args.command == "pagerank": from app.ingest.pagerank import update_authority from app.repository import ElasticRepository repository = ElasticRepository() try: return await update_authority( repository, damping=args.damping, iterations=args.iterations, batch_size=args.batch_size, ) finally: await repository.close() else: raise ValueError("Unsupported command") stage = args.command if args.command == "normalize": from app.ingest.openalex import NORMALIZER_VERSION stage += ":" + NORMALIZER_VERSION if args.command == "classify": from app.classifier import WEAK_CLASSIFIER_VERSION model_path = args.model or settings.classifier_model_path stage += ":" + (hashlib.sha256(Path(model_path).read_bytes()).hexdigest() if model_path else WEAK_CLASSIFIER_VERSION) # Trained models still use the weak lexicon for evidence and rare-class fallbacks. if model_path: stage += ":" + WEAK_CLASSIFIER_VERSION if args.command == "embed": stage += ":" + settings.embedding_model return transform_file(args.input, args.output, transform, resume=args.resume, stage=stage, batch_size=args.batch_size) def main(): arguments = parser().parse_args() try: print(json.dumps(asyncio.run(run(arguments)), indent=2)) except (ValueError, RuntimeError, FileExistsError, TypeError) as exc: raise SystemExit(str(exc)) from None if __name__ == "__main__": main()