| """Universal Hugging Face repo prober + deep folder streamer (PRIVATE to Data_Module). |
| |
| * Probes repo_type="dataset" first and falls back to repo_type="model" automatically. |
| * Lists the WHOLE repo tree recursively (metadata only - nothing is downloaded yet) and prints the |
| total number of files discovered. |
| * Applies the path blacklist BEFORE any byte is fetched, then streams the allowed files one at a time |
| (first 4096 bytes -> in-stream garbage gate -> full read), so the repo never sits in RAM. |
| * `limit` = maximum number of ALLOWED files to stream (None = "full"). When more files are allowed than |
| the limit, they are sampled at an even stride so the sample spans all sub-folders instead of only the first. |
| """ |
| from __future__ import annotations |
|
|
| import time |
| from dataclasses import dataclass |
| from typing import Dict, Iterator, List, Optional, Tuple |
|
|
| from .blacklist_gatekeeper import SAMPLE_BYTES, BlacklistGatekeeper |
|
|
| PLACEHOLDER_MARK = "PUT_YOUR_TOKEN" |
| MAX_FILE_BYTES = 4 * 1024 * 1024 |
|
|
|
|
| @dataclass |
| class RepoTelemetry: |
| repo: str |
| repo_type: str = "?" |
| total_discovered: int = 0 |
| skipped_blacklist: int = 0 |
| not_streamed: int = 0 |
| garbage_gate: int = 0 |
| too_large: int = 0 |
| read_errors: int = 0 |
| filter_destroyed: int = 0 |
| duplicate_files: int = 0 |
| partial_dedup: int = 0 |
| kept: int = 0 |
| error: str = "" |
|
|
| @property |
| def destroyed(self) -> int: |
| return (self.garbage_gate + self.too_large + self.read_errors |
| + self.filter_destroyed + self.duplicate_files) |
|
|
|
|
| def _clean_token(token: Optional[str]) -> Optional[str]: |
| if not token or PLACEHOLDER_MARK in token: |
| return None |
| return token |
|
|
|
|
| class _HubBackend: |
| """Thin wrapper over huggingface_hub so the streamer can be unit-tested with a fake.""" |
|
|
| def __init__(self, token: Optional[str]): |
| from huggingface_hub import HfApi, HfFileSystem |
| self.token = token |
| self.api = HfApi(token=token) |
| self.fs = HfFileSystem(token=token) |
|
|
| def probe(self, repo: str) -> Tuple[str, List[str]]: |
| problems: List[str] = [] |
| for repo_type in ("dataset", "model"): |
| try: |
| self.api.repo_info(repo, repo_type=repo_type) |
| return repo_type, problems |
| except Exception as exc: |
| problems.append(f"{repo_type}: {type(exc).__name__}: {str(exc).splitlines()[0][:120]}") |
| raise RuntimeError("repo not reachable as dataset or model -> " + " | ".join(problems)) |
|
|
| def list_files(self, repo: str, repo_type: str) -> List[Tuple[str, Optional[int]]]: |
| files: List[Tuple[str, Optional[int]]] = [] |
| for item in self.api.list_repo_tree(repo, repo_type=repo_type, recursive=True): |
| if type(item).__name__ == "RepoFolder": |
| continue |
| files.append((item.path, getattr(item, "size", None))) |
| return files |
|
|
| def open_file(self, repo: str, repo_type: str, path: str): |
| prefix = "datasets/" if repo_type == "dataset" else "" |
| return self.fs.open(f"{prefix}{repo}/{path}", "rb") |
|
|
|
|
| class HFUniversalStreamer: |
| def __init__(self, token: Optional[str], gate: BlacklistGatekeeper, backend=None, |
| max_file_bytes: int = MAX_FILE_BYTES): |
| self.gate = gate |
| self.max_file_bytes = max_file_bytes |
| self.backend = backend or _HubBackend(_clean_token(token)) |
|
|
| @staticmethod |
| def _pick(paths: List[Tuple[str, Optional[int]]], limit: Optional[int]) -> List[Tuple[str, Optional[int]]]: |
| if limit is None or len(paths) <= limit: |
| return paths |
| stride = len(paths) / float(limit) |
| return [paths[int(i * stride)] for i in range(limit)] |
|
|
| def _read(self, repo: str, repo_type: str, path: str, size: Optional[int], tel: RepoTelemetry |
| ) -> Optional[bytes]: |
| for attempt in (1, 2): |
| try: |
| with self.backend.open_file(repo, repo_type, path) as handle: |
| head = handle.read(SAMPLE_BYTES) |
| reason = self.gate.content_garbage_reason(head) |
| if reason: |
| tel.garbage_gate += 1 |
| return None |
| rest = handle.read(self.max_file_bytes + 1 - len(head)) if len(head) >= SAMPLE_BYTES else b"" |
| data = head + rest |
| if len(data) > self.max_file_bytes: |
| tel.too_large += 1 |
| return None |
| return data |
| except Exception: |
| if attempt == 2: |
| tel.read_errors += 1 |
| return None |
| time.sleep(1.0) |
| return None |
|
|
| def stream(self, repo: str, limit: Optional[int], tel: RepoTelemetry) -> Iterator[Dict]: |
| repo_type, problems = self.backend.probe(repo) |
| tel.repo_type = repo_type |
| listing = self.backend.list_files(repo, repo_type) |
| tel.total_discovered = len(listing) |
| print(f" [stream] {repo} resolved as '{repo_type}'" |
| + (f" (dataset probe failed: {problems[0]})" if problems else "") |
| + f" -> {tel.total_discovered} files discovered") |
|
|
| allowed: List[Tuple[str, Optional[int]]] = [] |
| for path, size in sorted(listing): |
| if self.gate.path_block_reason(path): |
| tel.skipped_blacklist += 1 |
| else: |
| allowed.append((path, size)) |
| chosen = self._pick(allowed, limit) |
| tel.not_streamed = len(allowed) - len(chosen) |
| print(f" [stream] {tel.skipped_blacklist} blocked by blacklist, {len(allowed)} allowed, " |
| f"streaming {len(chosen)} (limit={'full' if limit is None else limit})") |
|
|
| for path, size in chosen: |
| if size is not None and size > self.max_file_bytes: |
| tel.too_large += 1 |
| continue |
| data = self._read(repo, repo_type, path, size, tel) |
| if data is None: |
| continue |
| text = data.decode("utf-8", errors="replace") |
| if text.count("\ufffd") > max(8, len(text) // 50): |
| tel.garbage_gate += 1 |
| continue |
| dot = path.rfind(".") |
| ext = path[dot:].lower() if dot > path.rfind("/") and dot > 0 else "" |
| yield {"path": path, "ext": ext, "text": text, "size": len(data)} |
