"""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 # allowed files beyond `limit` garbage_gate: int = 0 # rejected by the in-stream 4096-byte heuristics too_large: int = 0 read_errors: int = 0 filter_destroyed: int = 0 # single-line / 1-3 word / empty-after-JSON-extraction duplicate_files: int = 0 # 100 % duplicates dropped by the context deduplicator partial_dedup: int = 0 # files where only a sub-section was excised 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: # RepositoryNotFound / Gated / auth / network 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): # >2 % undecodable -> binary in disguise 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)}