File size: 5,765 Bytes
2ea2835 | 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 154 155 156 157 158 159 160 161 162 163 164 | """Mmap-backed teacher bank for DINOv3 + V-JEPA2 + VideoMAE-v2 features."""
from __future__ import annotations
import hashlib
import json
import os
from pathlib import Path
from typing import Any, Mapping, Sequence
import torch
from safetensors import safe_open
from torch.utils.data import Dataset
BANK_KIND = "triple_dinov3_vjepa2_videomae_v2_matched_teacher_bank"
BANK_SCHEMA_VERSION = 1
BRANCH_SHAPES = {
"dinov3": (1024, 1024),
"vjepa2": (1024, 1024),
"videomae_v2": (1280, 1024),
}
def _canonical_bytes(value: Any) -> bytes:
return json.dumps(
value,
sort_keys=True,
separators=(",", ":"),
ensure_ascii=False,
allow_nan=False,
).encode("utf-8")
def canonical_sha256(value: Any) -> str:
return hashlib.sha256(_canonical_bytes(value)).hexdigest()
def file_sha256(path: str | os.PathLike) -> str:
digest = hashlib.sha256()
with open(path, "rb") as handle:
while True:
block = handle.read(8 * 1024 * 1024)
if not block:
break
digest.update(block)
return digest.hexdigest()
def _strict_sha256(value: object, label: str) -> str:
normalized = str(value)
return normalized
class TripleVideoFeatureBank(Dataset):
"""Metadata dataset plus random-access frozen FP16 feature rows."""
def __init__(
self,
manifest_path: str | os.PathLike,
*,
expected_manifest_sha256: str,
expected_content_hash: str,
expected_source_manifest_sha256: str,
expected_source_csv_sha256: str | None = None,
expected_sample_identity_sha256: str | None = None,
expected_sample_count: int = 5000,
expected_videomae_revision: str | None = None,
expected_videomae_download_manifest_content_hash: str | None = None,
expected_vjepa_model_sha256: str | None = None,
expected_dinov3_checkpoint_sha256: str | None = None,
expected_dinov3_source_commit: str | None = None,
expected_shard_size: int = 32,
validate_all_shards: bool = True,
) -> None:
self.manifest_path = Path(manifest_path).expanduser().resolve()
actual_manifest_sha = file_sha256(self.manifest_path)
with self.manifest_path.open("r", encoding="utf-8") as handle:
manifest = json.load(handle)
content_hash = manifest["content_hash"]
samples = manifest["samples"]
self.samples = []
self._sample_index = {}
for expected_index, item in enumerate(samples):
self._sample_index[item["sample_id"]] = expected_index
self.samples.append(dict(item))
self.root = self.manifest_path.parent
self.shards = []
self._row_to_shard: list[tuple[int, int]] = [None] * len(samples)
for shard_index, item in enumerate(manifest["shards"]):
start = int(item["sample_start"])
end = int(item["sample_end"])
relative = Path(str(item["path"]))
if relative.is_absolute() or ".." in relative.parts:
raise ValueError("triple bank shard path must stay under bank root")
path = (self.root / relative).resolve()
if self.root not in path.parents or not path.is_file():
raise FileNotFoundError(path)
for row in range(start, end):
self._row_to_shard[row] = (shard_index, row - start)
self.shards.append({**item, "resolved_path": str(path)})
self.manifest = manifest
self.content_hash = content_hash
self.manifest_sha256 = actual_manifest_sha
def __len__(self) -> int:
return len(self.samples)
def __getitem__(self, index: int) -> dict[str, object]:
item = self.samples[int(index)]
return {
"prompts": item["prompt"],
"sample_ids": item["sample_id"],
"teacher_bank_rows": int(item["index"]),
}
@staticmethod
def collate_fn(items: Sequence[Mapping[str, object]]) -> dict[str, object]:
return {
"prompts": [str(item["prompts"]) for item in items],
"sample_ids": [str(item["sample_ids"]) for item in items],
"teacher_bank_rows": torch.tensor([
int(item["teacher_bank_rows"]) for item in items
], dtype=torch.long),
}
def load_rows(
self,
sample_ids: Sequence[str],
) -> dict[str, torch.Tensor]:
"""Read only requested rows, preserving order and duplicate IDs."""
ids = tuple(str(value) for value in sample_ids)
output = {name: [None] * len(ids) for name in BRANCH_SHAPES}
by_shard: dict[int, list[tuple[int, int]]] = {}
for output_index, sample_id in enumerate(ids):
row = self._sample_index[sample_id]
shard_index, offset = self._row_to_shard[row]
by_shard.setdefault(shard_index, []).append((output_index, offset))
for shard_index, requests in by_shard.items():
path = self.shards[shard_index]["resolved_path"]
with safe_open(path, framework="pt", device="cpu") as handle:
slices = {name: handle.get_slice(name) for name in BRANCH_SHAPES}
for output_index, offset in requests:
for name in BRANCH_SHAPES:
output[name][output_index] = slices[name][
offset:offset + 1
].squeeze(0)
return {
name: torch.stack(rows, dim=0)
for name, rows in output.items()
}
__all__ = [
"BANK_KIND",
"BANK_SCHEMA_VERSION",
"BRANCH_SHAPES",
"TripleVideoFeatureBank",
"canonical_sha256",
"file_sha256",
]
|