scripts/upload_data_to_hf.py

"""
upload_data_to_hf.py · scripts/upload_data_to_hf.py

Upload local data assets to a Hugging Face **Dataset** repo so the Space can
hydrate from it at boot (see finagent/bootstrap.py). Run this once (and again
whenever you rebuild the corpus). The dataset repo has its own storage and does
NOT count against the Space's 1 GB repo limit.

Usage
-----
    # needs a write token: `huggingface-cli login` or HF_TOKEN env
    python scripts/upload_data_to_hf.py                 # uploads data/chroma → chroma/
    python scripts/upload_data_to_hf.py --pdfs --eval   # also archive PDFs + FinanceBench
    python scripts/upload_data_to_hf.py --repo me/mydata --private

Layout created in the dataset:
    chroma/        the vector store (what the Space downloads at boot)
    pdfs/          source filings              (optional, archival)
    financebench/  eval files                  (optional, archival)
"""

from __future__ import annotations

import argparse
import os
from pathlib import Path

DEFAULT_DATA_REPO = "Sarthak004/finagent-data"


from huggingface_hub import get_token

def _token():
    return (
        os.getenv("HF_TOKEN")
        or os.getenv("HUGGING_FACE_TOKEN")
        or get_token()
    )

def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--repo", default=os.getenv("FINAGENT_DATA_REPO", DEFAULT_DATA_REPO))
    p.add_argument("--private", action="store_true",
                   help="Create the dataset as private (needs HF_TOKEN to read at boot).")
    p.add_argument("--pdfs", action="store_true", help="Also upload data/india + data/us PDFs.")
    p.add_argument("--eval", action="store_true", help="Also upload FinanceBench eval files.")
    p.add_argument("--chroma-dir", default="data/chroma")
    args = p.parse_args()

    from huggingface_hub import HfApi

    token = _token()
    if not token:
        raise SystemExit("No HF token found. Run `huggingface-cli login` or set HF_TOKEN.")

    api = HfApi(token=token)
    api.create_repo(args.repo, repo_type="dataset", private=args.private, exist_ok=True)
    print(f"dataset repo: {args.repo} (private={args.private})")

    chroma = Path(args.chroma_dir)
    if not (chroma / "chroma.sqlite3").exists():
        raise SystemExit(f"No Chroma store at {chroma} (expected chroma.sqlite3).")
    print(f"uploading {chroma} → chroma/ …")
    api.upload_folder(folder_path=str(chroma), path_in_repo="chroma",
                      repo_id=args.repo, repo_type="dataset",
                      commit_message="Upload Chroma vector store")

    # Optional archival uploads — not needed by the Space at runtime.
    optional: list[tuple[bool, str, str]] = [
        (args.pdfs, "data/india", "pdfs/india"),
        (args.pdfs, "data/us", "pdfs/us"),
        (args.eval, "data/us/eval/financebench", "financebench"),
    ]
    for enabled, local, remote in optional:
        if enabled and Path(local).exists():
            print(f"uploading {local} → {remote} …")
            api.upload_folder(folder_path=local, path_in_repo=remote,
                              repo_id=args.repo, repo_type="dataset",
                              commit_message=f"Upload {remote}")

    print(f"\nDone. The Space will download chroma/ from {args.repo} on first boot.")


if __name__ == "__main__":
    main()