File size: 3,664 Bytes
69437d8 03b3d27 69437d8 cbd62a3 69437d8 03b3d27 69437d8 ee7e79e 69437d8 ee7e79e 69437d8 03b3d27 ee7e79e 69437d8 ee7e79e 69437d8 cbd62a3 9ca8ee5 ee7e79e 03b3d27 69437d8 03b3d27 69437d8 | 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 | """Weekly delta-sync entrypoint: embed new/changed papers and publish the snapshot to HF Hub."""
import json
import os
from datetime import datetime, timezone
from common.db import init_db, upsert_papers, mark_inactive_ids, update_metadata_fields
from common.vector_index import create_index, load_index, save_index, add_vector
from sync.anthology_source import iter_papers
from sync.delta import compute_delta, build_state
from sync.embeddings_client import EmbeddingsClient
from sync.hf_dataset_store import download_snapshot, upload_snapshot
def _embedding_text(paper: dict) -> str:
if paper["abstract"]:
return f"{paper['title']}\n\n{paper['abstract']}"
return paper["title"]
def run_sync(anthology_path: str, hf_repo_id: str, hf_token: str,
embedding_base_url: str, embedding_api_key: str, work_dir: str) -> None:
os.makedirs(work_dir, exist_ok=True)
downloaded_state = download_snapshot(hf_repo_id, work_dir, hf_token)
last_state = downloaded_state.get("papers", {})
db_path = os.path.join(work_dir, "papers.db")
index_path = os.path.join(work_dir, "index.faiss")
conn = init_db(db_path)
try:
index = load_index(index_path)
except FileNotFoundError:
index = None # created lazily once the embedding dimension is known
current_papers = list(iter_papers(anthology_path))
delta = compute_delta(current_papers, last_state)
if delta.new_or_changed:
client = EmbeddingsClient(base_url=embedding_base_url, api_key=embedding_api_key)
texts = [_embedding_text(p) for p in delta.new_or_changed]
vectors = client.embed_batch(texts) # raises RuntimeError on failure, propagates before any writes below
# qwen3-embedding-4b returns a fixed native dim (no matryoshka), so the
# index dimension is whatever the model emits, derived from the first vector.
embedding_dim = len(vectors[0])
if index is None:
index = create_index(dim=embedding_dim)
elif index.d != embedding_dim:
raise RuntimeError(
f"Embedding dimension {embedding_dim} does not match existing index "
f"dimension {index.d}; delete the stale index.faiss (locally and on "
f"the HF Hub) and run a fresh full sync."
)
upsert_rows = []
for paper, vector in zip(delta.new_or_changed, vectors):
faiss_id = add_vector(index, vector)
upsert_rows.append({**paper, "active": True, "faiss_id": faiss_id})
upsert_papers(conn, upsert_rows)
mark_inactive_ids(conn, list(delta.removed_ids))
# Keep bibtex/pdf_url current for all current papers without re-embedding.
# The embedding delta above only touched new/changed papers; this
# fills/refreshes these fields for the unchanged majority (incl. the
# one-time backfill after the columns were added).
update_metadata_fields(conn, current_papers)
if index is not None:
save_index(index, index_path)
state = {
"papers": build_state(current_papers),
"last_synced_at": datetime.now(timezone.utc).isoformat(),
}
with open(os.path.join(work_dir, "state.json"), "w") as f:
json.dump(state, f)
upload_snapshot(hf_repo_id, work_dir, hf_token)
if __name__ == "__main__":
run_sync(
anthology_path=os.environ["ANTHOLOGY_PATH"],
hf_repo_id=os.environ["HF_REPO_ID"],
hf_token=os.environ["HF_TOKEN"],
embedding_base_url=os.environ["EMBEDDING_BASE_URL"],
embedding_api_key=os.environ["EMBEDDING_API_KEY"],
work_dir=os.environ.get("WORK_DIR", "./work"),
)
|