Spaces:
Sleeping
Sleeping
| """Bounded normalization/embedding/linking directly into the shared ES index.""" | |
| import asyncio | |
| import json | |
| from pathlib import Path | |
| from elasticsearch.helpers import async_scan | |
| from app.classifier import WEAK_CLASSIFIER_VERSION, weak_classify | |
| from app.config import Settings, settings | |
| from app.ingest.ctgov import flatten_trial | |
| from app.ingest.fetch import atomic_json | |
| from app.ingest.linker import link_studies | |
| from app.ingest.openalex import NORMALIZER_VERSION, normalize_work | |
| from app.ingest.snapshot import fingerprint | |
| async def link_indexed_batch(repository, documents: list[dict]) -> None: | |
| """Link across batches and sources; distinct registered trials stay distinct.""" | |
| pmids = sorted({str(pmid) for doc in documents for pmid in doc.get("pmids", [])}) | |
| nct_ids = sorted({nct for doc in documents for nct in doc.get("nct_ids", [])}) | |
| result_pmids = sorted({str(pmid) for doc in documents for pmid in doc.get("result_pmids", [])}) | |
| clauses = [] | |
| for name, values in (("pmids", result_pmids), ("result_pmids", pmids), ("nct_ids", nct_ids)): | |
| if values: | |
| clauses.append({"terms": {name: values}}) | |
| if not clauses: | |
| return | |
| current = await repository.get_many([doc["id"] for doc in documents], include_vectors=False) | |
| query = { | |
| "query": { | |
| "bool": { | |
| "filter": [{"term": {"record_kind": "study"}}], | |
| "should": clauses, | |
| "minimum_should_match": 1, | |
| "must_not": [{"term": {"is_review": True}}], | |
| } | |
| }, | |
| "_source": {"excludes": ["embedding", "attachments"]}, | |
| } | |
| related = [ | |
| hit["_source"] | |
| async for hit in async_scan( | |
| repository.client, | |
| index=repository.index, | |
| query=query, | |
| size=500, | |
| ) | |
| ] | |
| linked = link_studies([*related, *current]) | |
| merged = [doc for doc in linked if doc.get("source") == "merged"] | |
| await repository.persist_links(merged) | |
| async def materialize( | |
| path: Path, | |
| repository, | |
| *, | |
| embedder=None, | |
| batch_size: int = 100, | |
| config: Settings = settings, | |
| progress=None, | |
| ) -> dict: | |
| """A committed byte offset avoids rereading a million rows on resume. | |
| No vector JSONL intermediate or all-corpus in-memory linker is created. | |
| A crash before checkpoint repeats idempotent writes, never loses a batch. | |
| """ | |
| if batch_size <= 0: | |
| raise ValueError("batch_size must be positive") | |
| source_checkpoint = path.with_suffix(path.suffix + ".checkpoint.json") | |
| source_identity = ( | |
| json.loads(source_checkpoint.read_text())["signature"] | |
| if source_checkpoint.exists() | |
| else [path.stat().st_size, path.stat().st_mtime_ns] | |
| ) | |
| signature = fingerprint( | |
| [ | |
| str(path.resolve()), | |
| source_identity, | |
| repository.index, | |
| config.elastic_url, | |
| NORMALIZER_VERSION, | |
| WEAK_CLASSIFIER_VERSION, | |
| config.embedding_model if embedder is not None else "no_vectors", | |
| ] | |
| ) | |
| checkpoint = path.with_suffix(path.suffix + ".materialize.json") | |
| state = { | |
| "signature": signature, | |
| "offset": 0, | |
| "records": 0, | |
| "metadata_only": 0, | |
| "embedded": 0, | |
| "complete": False, | |
| } | |
| if checkpoint.exists(): | |
| state = json.loads(checkpoint.read_text()) | |
| if state["signature"] != signature: | |
| raise ValueError( | |
| "Source, target index or model changed; use a new materialization checkpoint" | |
| ) | |
| if path.stat().st_size < state["offset"]: | |
| raise ValueError("Materialized source was truncated") | |
| if state["complete"] and path.stat().st_size == state["offset"]: | |
| return state | |
| state["complete"] = False | |
| with path.open("rb") as handle: | |
| handle.seek(state["offset"]) | |
| while True: | |
| batch = [] | |
| while len(batch) < batch_size: | |
| line = handle.readline() | |
| if not line: | |
| break | |
| if not line.strip(): | |
| continue | |
| raw = json.loads(line) | |
| doc = flatten_trial(raw) if "protocolSection" in raw else normalize_work(raw) | |
| if doc["source"] == "openalex": | |
| doc.update(weak_classify(doc["abstract"], is_review=doc["is_review"])) | |
| batch.append(doc) | |
| if not batch: | |
| break | |
| if embedder is not None: | |
| texts = [ | |
| " ".join( | |
| str(doc.get(key) or "") | |
| for key in ("title", "abstract", "intervention", "outcome") | |
| ) | |
| for doc in batch | |
| ] | |
| vectors = await asyncio.to_thread(embedder.encode, texts) | |
| for doc, vector in zip(batch, vectors, strict=True): | |
| doc.update(embedding=vector, embedding_model=config.embedding_model) | |
| await repository.bulk_upsert(batch) | |
| await link_indexed_batch(repository, batch) | |
| state.update( | |
| offset=handle.tell(), | |
| records=state["records"] + len(batch), | |
| metadata_only=state["metadata_only"] | |
| + sum(doc.get("abstract_available") is False for doc in batch), | |
| embedded=state["embedded"] + (len(batch) if embedder is not None else 0), | |
| ) | |
| atomic_json(checkpoint, state) | |
| if progress: | |
| progress(dict(state)) | |
| state["offset"] = handle.tell() | |
| state["complete"] = True | |
| atomic_json(checkpoint, state) | |
| return state | |