Spaces:
Sleeping
Sleeping
File size: 5,704 Bytes
b92b744 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | """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
|