sync/hf_dataset_store.py

"""Persists the FAISS index, SQLite DB, and sync state to an HF Hub dataset repo."""
import json
import os
from huggingface_hub import hf_hub_download, upload_folder
from huggingface_hub.errors import EntryNotFoundError

FILES = ["index.faiss", "papers.db", "state.json"]

EMPTY_STATE = {"papers": {}, "last_synced_at": None}


def download_snapshot(repo_id: str, local_dir: str, token: str) -> dict:
    os.makedirs(local_dir, exist_ok=True)
    state = dict(EMPTY_STATE)
    for filename in FILES:
        try:
            hf_hub_download(repo_id=repo_id, repo_type="dataset", filename=filename, local_dir=local_dir, token=token)
        except (FileNotFoundError, EntryNotFoundError):
            continue
        if filename == "state.json":
            with open(os.path.join(local_dir, filename)) as f:
                content = f.read().strip()
            state = json.loads(content) if content else dict(EMPTY_STATE)
    return state


def upload_snapshot(repo_id: str, local_dir: str, token: str) -> None:
    upload_folder(
        repo_id=repo_id,
        repo_type="dataset",
        folder_path=local_dir,
        token=token,
    )