Spaces:
Sleeping
Sleeping
| """HF Dataset I/O for the retrieval-tuning loop. | |
| Three files are exchanged between the GitHub Actions tuner and the | |
| chat-logs HF Dataset: | |
| - turns.jsonl (chat writes; tuner annotates with judges) | |
| - retrieval_tuning.json (tuner writes; chat reads at boot) | |
| - tuning_log.jsonl (tuner-only audit trail) | |
| fetch_state() downloads all three to derived/student-chat/, materializing | |
| defaults for any file missing from the dataset (cold start). | |
| push_state() uploads the same three local paths back. | |
| Lazy imports of huggingface_hub via re-exported wrappers so tests can patch | |
| this module's symbols and so this module imports without huggingface_hub | |
| installed (local dev). | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import logging | |
| import shutil | |
| from pathlib import Path | |
| from scripts.lib.chat_eval.tuning import empty_state | |
| from scripts.lib.turns_merge import is_not_found_error | |
| logger = logging.getLogger(__name__) | |
| _DATASET_FILES = ("turns.jsonl", "retrieval_tuning.json", "tuning_log.jsonl") | |
| def hf_hub_download(*args, **kwargs): | |
| """Re-export so tests can patch this module's reference.""" | |
| from huggingface_hub import hf_hub_download as _real | |
| return _real(*args, **kwargs) | |
| def HfApi(*args, **kwargs): | |
| """Re-export so tests can patch this module's reference.""" | |
| from huggingface_hub import HfApi as _real | |
| return _real(*args, **kwargs) | |
| def fetch_state(*, repo_id: str, token: str, derived_dir: Path) -> None: | |
| """Download the three dataset files into derived_dir. | |
| Missing files get sensible defaults so the rest of the loop has real | |
| files to read: | |
| - turns.jsonl missing -> empty file (no turns to judge) | |
| - retrieval_tuning.json missing -> empty_state() | |
| - tuning_log.jsonl missing -> empty file | |
| """ | |
| derived_dir.mkdir(parents=True, exist_ok=True) | |
| for fname in _DATASET_FILES: | |
| local = derived_dir / fname | |
| try: | |
| tmp_path = hf_hub_download( | |
| repo_id=repo_id, | |
| filename=fname, | |
| repo_type="dataset", | |
| token=token, | |
| ) | |
| shutil.copyfile(tmp_path, local) | |
| logger.info("state_io: fetched %s", fname) | |
| except Exception as e: | |
| if not is_not_found_error(e): | |
| raise | |
| if fname == "retrieval_tuning.json": | |
| local.write_text(json.dumps(empty_state()), encoding="utf-8") | |
| else: | |
| local.write_text("", encoding="utf-8") | |
| logger.info("state_io: %s missing on remote; wrote default.", fname) | |
| def push_state(*, repo_id: str, token: str, derived_dir: Path, | |
| commit_message: str) -> None: | |
| """Upload the three files from derived_dir to the dataset. | |
| For turns.jsonl: refetch remote and merge by turn_id (local wins on | |
| conflict; remote-only turns preserved). retrieval_tuning.json and | |
| tuning_log.jsonl are tuner-only writes and overwrite remote. | |
| Raises FileNotFoundError if any file is missing locally — we never | |
| push an empty over a populated remote by accident. | |
| """ | |
| api = HfApi() | |
| for fname in _DATASET_FILES: | |
| local = derived_dir / fname | |
| if not local.is_file(): | |
| raise FileNotFoundError( | |
| f"state_io.push_state: missing local file {local}; " | |
| f"refusing to push." | |
| ) | |
| staging: Path | None = None | |
| if fname == "turns.jsonl": | |
| merged_bytes = _merge_turns_with_remote( | |
| local, repo_id=repo_id, token=token, | |
| ) | |
| staging = local.with_suffix(local.suffix + ".merged") | |
| staging.write_bytes(merged_bytes) | |
| upload_path = str(staging) | |
| else: | |
| upload_path = str(local) | |
| try: | |
| api.upload_file( | |
| path_or_fileobj=upload_path, | |
| path_in_repo=fname, | |
| repo_id=repo_id, | |
| repo_type="dataset", | |
| token=token, | |
| commit_message=commit_message, | |
| ) | |
| logger.info("state_io: pushed %s", fname) | |
| finally: | |
| if staging is not None: | |
| try: | |
| staging.unlink() | |
| except OSError: | |
| pass | |
| def _merge_turns_with_remote(local_path: Path, *, repo_id: str, | |
| token: str) -> bytes: | |
| """Tuner-side adapter: read local file, delegate to turns_merge.""" | |
| from scripts.lib.turns_merge import merge_turns_with_remote | |
| return merge_turns_with_remote(local_path.read_bytes(), | |
| repo_id=repo_id, token=token) | |