File size: 6,737 Bytes
9d614f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | """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)} |