Spaces:
Sleeping
Sleeping
| """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() | |