Spaces:
Sleeping
Sleeping
| """Utilities for discovering and loading precomputed patch metadata.""" | |
| from __future__ import annotations | |
| import json | |
| import os | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from typing import Dict, Iterable, Mapping, Optional | |
| APP_DIR = Path(__file__).resolve().parent | |
| DEFAULT_DATA_CANDIDATES = ( | |
| APP_DIR / "data" / "test-10", | |
| Path("data/test-10"), | |
| ) | |
| ENV_DATA_ROOT = "PATCH_VIEW_DATA_ROOT" | |
| def resolve_data_root(candidates: Iterable[Path] = DEFAULT_DATA_CANDIDATES) -> Path: | |
| """Locate the dataset root, honoring environment override when present.""" | |
| override = os.getenv(ENV_DATA_ROOT) | |
| if override: | |
| candidate = Path(override).expanduser().resolve() | |
| if not candidate.exists(): | |
| raise FileNotFoundError(f"PATCH_VIEW_DATA_ROOT={candidate} not found") | |
| return candidate | |
| for candidate in candidates: | |
| resolved = candidate.expanduser() | |
| if (resolved / "precomputed").exists(): | |
| return resolved | |
| # Fall back to the last candidate even if missing, to keep error messages consistent | |
| return Path(list(candidates)[-1]) | |
| class PrecomputedRecord: | |
| """Loaded metadata for a single audio file.""" | |
| file_id: str | |
| data: Dict[str, object] | |
| def audio_path(self) -> str: | |
| return str(self.data["audio_path"]) | |
| def reconstructed_audio_path(self) -> str: | |
| return str(self.data["reconstructed_audio_path"]) | |
| class PrecomputedRepository: | |
| """In-memory index of precomputed entropy/perplexity artifacts.""" | |
| def __init__(self, data_root: Optional[Path] = None): | |
| self.data_root = data_root or resolve_data_root() | |
| self.precomputed_dir = self.data_root / "precomputed" | |
| self._records: Dict[str, PrecomputedRecord] = {} | |
| def records(self) -> Mapping[str, PrecomputedRecord]: | |
| return self._records | |
| def load(self) -> None: | |
| """Load all JSON metadata into memory.""" | |
| if not self.precomputed_dir.exists(): | |
| raise FileNotFoundError( | |
| f"Expected precomputed dir at {self.precomputed_dir}, but it does not exist." | |
| ) | |
| for json_path in sorted(self.precomputed_dir.glob("*.json")): | |
| with open(json_path, "r") as handle: | |
| raw = json.load(handle) | |
| enriched = dict(raw) | |
| # Rewrite paths to prefer colocated dataset snapshot | |
| enriched["audio_path"] = self._resolve_audio_path(raw["audio_path"]) | |
| enriched["reconstructed_audio_path"] = str( | |
| self.precomputed_dir / Path(raw["reconstructed_audio_path"]).name | |
| ) | |
| record = PrecomputedRecord(file_id=enriched["file_id"], data=enriched) | |
| self._records[record.file_id] = record | |
| def ids(self) -> Iterable[str]: | |
| return sorted(self._records.keys()) | |
| def get(self, file_id: str) -> Dict[str, object]: | |
| if file_id not in self._records: | |
| raise KeyError(f"Unknown file id '{file_id}'") | |
| return self._records[file_id].data | |
| def _resolve_audio_path(self, original_path: str) -> str: | |
| original_name = Path(original_path).name | |
| local_candidate = self.data_root / original_name | |
| if local_candidate.exists(): | |
| return str(local_candidate) | |
| return str(Path(original_path)) | |