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
|
from __future__ import annotations
import torch
from .common import normalized
def load_feature_pair(
vision_path: str, text_path: str
) -> tuple[dict, dict, dict[int, int], dict[int, int]]:
vision = torch.load(vision_path, map_location="cpu", weights_only=False)
text = torch.load(text_path, map_location="cpu", weights_only=False)
vision["features"] = normalized(vision["features"])
text["features"] = normalized(text["features"])
vision_lookup = {int(row): i for i, row in enumerate(vision["rows"])}
text_lookup = {int(row): i for i, row in enumerate(text["rows"])}
return vision, text, vision_lookup, text_lookup
def select_rows(
feature: torch.Tensor, lookup: dict[int, int], rows: list[int]
) -> torch.Tensor:
missing = [int(row) for row in rows if int(row) not in lookup]
if missing:
raise KeyError(
f"{len(missing)} requested rows are absent from feature cache; "
f"first missing rows: {missing[:5]}"
)
return feature[[lookup[int(row)] for row in rows]]
|
