hyper-models 0.2.0__tar.gz → 0.3.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (24) hide show
  1. {hyper_models-0.2.0 → hyper_models-0.3.1}/PKG-INFO +18 -7
  2. {hyper_models-0.2.0 → hyper_models-0.3.1}/README.md +8 -4
  3. {hyper_models-0.2.0 → hyper_models-0.3.1}/pyproject.toml +14 -3
  4. hyper_models-0.3.1/src/hyper3_clip/__init__.py +3 -0
  5. hyper_models-0.3.1/src/hyper3_clip/models/__init__.py +3 -0
  6. hyper_models-0.3.1/src/hyper3_clip/models/encoders.py +181 -0
  7. hyper_models-0.3.1/src/hyper3_clip/models/experimental.py +691 -0
  8. hyper_models-0.3.1/src/hyper3_clip/models/himo.py +54 -0
  9. hyper_models-0.3.1/src/hyper3_clip/models/hyper3_clip.py +1032 -0
  10. hyper_models-0.3.1/src/hyper3_clip/models/lorentz.py +285 -0
  11. hyper_models-0.3.1/src/hyper3_clip/models/losses.py +1546 -0
  12. hyper_models-0.3.1/src/hyper3_clip/models/objectives.py +653 -0
  13. hyper_models-0.3.1/src/hyper3_clip/models/tren.py +282 -0
  14. hyper_models-0.3.1/src/hyper3_clip/training/__init__.py +1 -0
  15. hyper_models-0.3.1/src/hyper3_clip/training/distributed.py +153 -0
  16. {hyper_models-0.2.0 → hyper_models-0.3.1}/src/hyper_models/__init__.py +11 -3
  17. {hyper_models-0.2.0 → hyper_models-0.3.1}/src/hyper_models/loader.py +7 -1
  18. {hyper_models-0.2.0 → hyper_models-0.3.1}/src/hyper_models/loaders.py +12 -0
  19. {hyper_models-0.2.0 → hyper_models-0.3.1}/src/hyper_models/models.py +3 -1
  20. {hyper_models-0.2.0 → hyper_models-0.3.1}/src/hyper_models/registry.py +21 -0
  21. {hyper_models-0.2.0 → hyper_models-0.3.1}/src/hyper_models/torch_models.py +139 -2
  22. {hyper_models-0.2.0 → hyper_models-0.3.1}/.gitignore +0 -0
  23. {hyper_models-0.2.0 → hyper_models-0.3.1}/LICENSE +0 -0
  24. {hyper_models-0.2.0 → hyper_models-0.3.1}/src/hyper_models/preprocessing.py +0 -0
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.4
1
+ Metadata-Version: 2.5
2
2
  Name: hyper-models
3
- Version: 0.2.0
3
+ Version: 0.3.1
4
4
  Summary: A model zoo for non-Euclidean embedding models (hyperbolic, spherical, product manifolds)
5
5
  Project-URL: Homepage, https://github.com/Hyper3Labs/hyper-models
6
6
  Project-URL: Repository, https://github.com/Hyper3Labs/hyper-models
@@ -22,14 +22,21 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
22
22
  Requires-Python: >=3.10
23
23
  Requires-Dist: huggingface-hub>=1.0
24
24
  Requires-Dist: numpy>=1.26
25
- Requires-Dist: onnxruntime>=1.20
25
+ Requires-Dist: onnxruntime<1.24,>=1.20; python_version < '3.11'
26
+ Requires-Dist: onnxruntime>=1.20; python_version >= '3.11'
26
27
  Requires-Dist: pillow>=10.0
27
28
  Provides-Extra: dev
28
29
  Requires-Dist: pytest>=7.0; extra == 'dev'
29
30
  Requires-Dist: ruff>=0.1; extra == 'dev'
31
+ Requires-Dist: tomli>=2.0; (python_version < '3.11') and extra == 'dev'
30
32
  Provides-Extra: ml
33
+ Requires-Dist: protobuf>=4.25.0; extra == 'ml'
34
+ Requires-Dist: pyyaml>=6.0.0; extra == 'ml'
35
+ Requires-Dist: safetensors>=0.4.0; extra == 'ml'
36
+ Requires-Dist: sentencepiece>=0.2.0; extra == 'ml'
31
37
  Requires-Dist: timm>=1.0.0; extra == 'ml'
32
38
  Requires-Dist: torch>=2.9.1; extra == 'ml'
39
+ Requires-Dist: transformers<6.0,>=5.0.0; extra == 'ml'
33
40
  Description-Content-Type: text/markdown
34
41
 
35
42
  # hyper-models
@@ -66,7 +73,7 @@ uv pip install hyper-models
66
73
  This base install is the simple path: it stays **torch-free** and is enough for
67
74
  ONNX-backed catalog entries such as HyCoCLIP and MERU.
68
75
 
69
- For torch-backed checkpoints (for example UNCHA):
76
+ For torch-backed checkpoints (for example UNCHA and Hyper3-CLIP):
70
77
 
71
78
  ```bash
72
79
  uv pip install "hyper-models[ml]"
@@ -80,11 +87,11 @@ from PIL import Image
80
87
 
81
88
  # List available models
82
89
  hyper_models.list_models()
83
- # ['hycoclip-vit-s', 'hycoclip-vit-b', 'meru-vit-s', 'meru-vit-b', 'uncha-vit-s', 'uncha-vit-b']
90
+ # ['hycoclip-vit-s', 'hycoclip-vit-b', 'meru-vit-s', 'meru-vit-b', 'uncha-vit-s', 'uncha-vit-b', 'hyper3-clip-v0.5']
84
91
 
85
92
  # Inspect supported internal loader kinds
86
93
  hyper_models.list_loaders()
87
- # ['onnx', 'uncha-image-torch']
94
+ # ['hyper3-clip-torch', 'onnx', 'uncha-image-torch']
88
95
 
89
96
  # Load model (auto-downloads from Hugging Face Hub)
90
97
  model = hyper_models.load("hycoclip-vit-s")
@@ -116,12 +123,14 @@ embeddings = model.encode(batch)
116
123
  - Internal loaders may differ by model family:
117
124
  - `onnx` for exported, torch-free runtimes
118
125
  - `uncha-image-torch` for raw checkpoints that need a PyTorch image runtime
126
+ - `hyper3-clip-torch` for Hyper3-CLIP safetensors checkpoints
119
127
 
120
128
  This keeps callers on one stable API:
121
129
 
122
130
  ```python
123
131
  model = hyper_models.load("hycoclip-vit-s")
124
132
  model = hyper_models.load("uncha-vit-b")
133
+ model = hyper_models.load("hyper3-clip-v0.5")
125
134
  ```
126
135
 
127
136
  Callers do not need to know which internal loader is used, except for optional
@@ -149,7 +158,8 @@ layout_key = dataset.compute_visualization(space_key=space_key, layout="poincare
149
158
  HyperView's simple path remains torch-free. If you use the default ONNX-backed
150
159
  `hyper-models` entries or the default `embed-anything` provider, HyperView does
151
160
  not need PyTorch. PyTorch is only needed when you explicitly select a
152
- torch-backed catalog entry such as `uncha-vit-s` or `uncha-vit-b`.
161
+ torch-backed catalog entry such as `uncha-vit-s`, `uncha-vit-b`, or
162
+ `hyper3-clip-v0.5`.
153
163
 
154
164
  ## Models
155
165
 
@@ -163,6 +173,7 @@ torch-backed catalog entry such as `uncha-vit-s` or `uncha-vit-b`.
163
173
  | `meru-vit-b` | [![HF](https://img.shields.io/badge/🤗-HuggingFace-yellow)](https://huggingface.co/mnm-matin/hyperbolic-clip/tree/main/meru-vit-b) | [ICML 2023](https://arxiv.org/abs/2304.09172) | [facebookresearch/meru](https://github.com/facebookresearch/meru) |
164
174
  | `uncha-vit-s` | [![HF](https://img.shields.io/badge/🤗-HuggingFace-yellow)](https://huggingface.co/hayeonkim/uncha/blob/main/uncha_vit_s.pth) | [CVPR 2026](https://arxiv.org/abs/2603.22042) | [jeeit17/UNCHA](https://github.com/jeeit17/UNCHA) |
165
175
  | `uncha-vit-b` | [![HF](https://img.shields.io/badge/🤗-HuggingFace-yellow)](https://huggingface.co/hayeonkim/uncha/blob/main/uncha_vit_b.pth) | [CVPR 2026](https://arxiv.org/abs/2603.22042) | [jeeit17/UNCHA](https://github.com/jeeit17/UNCHA) |
176
+ | `hyper3-clip-v0.5` | [![HF](https://img.shields.io/badge/🤗-HuggingFace-yellow)](https://huggingface.co/hyper3labs/hyper3-clip-v0.5) | — | [Hyper3Labs/hyper3-clip](https://github.com/Hyper3Labs/hyper3-clip) |
166
177
  | `hyp-vit` | — | [CVPR 2022](https://arxiv.org/abs/2203.10833) | [htdt/hyp_metric](https://github.com/htdt/hyp_metric) |
167
178
  | `hie` | — | [CVPR 2020](https://arxiv.org/abs/1904.02239) | [leymir/hyperbolic-image-embeddings](https://github.com/leymir/hyperbolic-image-embeddings) |
168
179
  | `hcnn` | — | [ICLR 2024](https://openreview.net/forum?id=ekz1hN5QNh) | [kschwethelm/HyperbolicCV](https://github.com/kschwethelm/HyperbolicCV) |
@@ -32,7 +32,7 @@ uv pip install hyper-models
32
32
  This base install is the simple path: it stays **torch-free** and is enough for
33
33
  ONNX-backed catalog entries such as HyCoCLIP and MERU.
34
34
 
35
- For torch-backed checkpoints (for example UNCHA):
35
+ For torch-backed checkpoints (for example UNCHA and Hyper3-CLIP):
36
36
 
37
37
  ```bash
38
38
  uv pip install "hyper-models[ml]"
@@ -46,11 +46,11 @@ from PIL import Image
46
46
 
47
47
  # List available models
48
48
  hyper_models.list_models()
49
- # ['hycoclip-vit-s', 'hycoclip-vit-b', 'meru-vit-s', 'meru-vit-b', 'uncha-vit-s', 'uncha-vit-b']
49
+ # ['hycoclip-vit-s', 'hycoclip-vit-b', 'meru-vit-s', 'meru-vit-b', 'uncha-vit-s', 'uncha-vit-b', 'hyper3-clip-v0.5']
50
50
 
51
51
  # Inspect supported internal loader kinds
52
52
  hyper_models.list_loaders()
53
- # ['onnx', 'uncha-image-torch']
53
+ # ['hyper3-clip-torch', 'onnx', 'uncha-image-torch']
54
54
 
55
55
  # Load model (auto-downloads from Hugging Face Hub)
56
56
  model = hyper_models.load("hycoclip-vit-s")
@@ -82,12 +82,14 @@ embeddings = model.encode(batch)
82
82
  - Internal loaders may differ by model family:
83
83
  - `onnx` for exported, torch-free runtimes
84
84
  - `uncha-image-torch` for raw checkpoints that need a PyTorch image runtime
85
+ - `hyper3-clip-torch` for Hyper3-CLIP safetensors checkpoints
85
86
 
86
87
  This keeps callers on one stable API:
87
88
 
88
89
  ```python
89
90
  model = hyper_models.load("hycoclip-vit-s")
90
91
  model = hyper_models.load("uncha-vit-b")
92
+ model = hyper_models.load("hyper3-clip-v0.5")
91
93
  ```
92
94
 
93
95
  Callers do not need to know which internal loader is used, except for optional
@@ -115,7 +117,8 @@ layout_key = dataset.compute_visualization(space_key=space_key, layout="poincare
115
117
  HyperView's simple path remains torch-free. If you use the default ONNX-backed
116
118
  `hyper-models` entries or the default `embed-anything` provider, HyperView does
117
119
  not need PyTorch. PyTorch is only needed when you explicitly select a
118
- torch-backed catalog entry such as `uncha-vit-s` or `uncha-vit-b`.
120
+ torch-backed catalog entry such as `uncha-vit-s`, `uncha-vit-b`, or
121
+ `hyper3-clip-v0.5`.
119
122
 
120
123
  ## Models
121
124
 
@@ -129,6 +132,7 @@ torch-backed catalog entry such as `uncha-vit-s` or `uncha-vit-b`.
129
132
  | `meru-vit-b` | [![HF](https://img.shields.io/badge/🤗-HuggingFace-yellow)](https://huggingface.co/mnm-matin/hyperbolic-clip/tree/main/meru-vit-b) | [ICML 2023](https://arxiv.org/abs/2304.09172) | [facebookresearch/meru](https://github.com/facebookresearch/meru) |
130
133
  | `uncha-vit-s` | [![HF](https://img.shields.io/badge/🤗-HuggingFace-yellow)](https://huggingface.co/hayeonkim/uncha/blob/main/uncha_vit_s.pth) | [CVPR 2026](https://arxiv.org/abs/2603.22042) | [jeeit17/UNCHA](https://github.com/jeeit17/UNCHA) |
131
134
  | `uncha-vit-b` | [![HF](https://img.shields.io/badge/🤗-HuggingFace-yellow)](https://huggingface.co/hayeonkim/uncha/blob/main/uncha_vit_b.pth) | [CVPR 2026](https://arxiv.org/abs/2603.22042) | [jeeit17/UNCHA](https://github.com/jeeit17/UNCHA) |
135
+ | `hyper3-clip-v0.5` | [![HF](https://img.shields.io/badge/🤗-HuggingFace-yellow)](https://huggingface.co/hyper3labs/hyper3-clip-v0.5) | — | [Hyper3Labs/hyper3-clip](https://github.com/Hyper3Labs/hyper3-clip) |
132
136
  | `hyp-vit` | — | [CVPR 2022](https://arxiv.org/abs/2203.10833) | [htdt/hyp_metric](https://github.com/htdt/hyp_metric) |
133
137
  | `hie` | — | [CVPR 2020](https://arxiv.org/abs/1904.02239) | [leymir/hyperbolic-image-embeddings](https://github.com/leymir/hyperbolic-image-embeddings) |
134
138
  | `hcnn` | — | [ICLR 2024](https://openreview.net/forum?id=ekz1hN5QNh) | [kschwethelm/HyperbolicCV](https://github.com/kschwethelm/HyperbolicCV) |
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "hyper-models"
3
- version = "0.2.0"
3
+ version = "0.3.1"
4
4
  description = "A model zoo for non-Euclidean embedding models (hyperbolic, spherical, product manifolds)"
5
5
  readme = "README.md"
6
6
  license = { text = "MIT" }
@@ -20,7 +20,11 @@ classifiers = [
20
20
  ]
21
21
  dependencies = [
22
22
  "numpy>=1.26",
23
- "onnxruntime>=1.20",
23
+ # onnxruntime stopped publishing cp310 wheels after 1.23.x; without this
24
+ # split a universal resolution picks a release that cannot install on the
25
+ # 3.10 this package still supports.
26
+ "onnxruntime>=1.20,<1.24; python_version<'3.11'",
27
+ "onnxruntime>=1.20; python_version>='3.11'",
24
28
  "huggingface-hub>=1.0",
25
29
  "pillow>=10.0",
26
30
  ]
@@ -29,10 +33,17 @@ dependencies = [
29
33
  dev = [
30
34
  "pytest>=7.0",
31
35
  "ruff>=0.1",
36
+ # tomllib is stdlib from 3.11; the version test reads pyproject.toml on 3.10 too.
37
+ "tomli>=2.0; python_version<'3.11'",
32
38
  ]
33
39
  ml = [
34
40
  "torch>=2.9.1",
35
41
  "timm>=1.0.0",
42
+ "transformers>=5.0.0,<6.0",
43
+ "safetensors>=0.4.0",
44
+ "pyyaml>=6.0.0",
45
+ "sentencepiece>=0.2.0",
46
+ "protobuf>=4.25.0",
36
47
  ]
37
48
 
38
49
  [project.urls]
@@ -46,7 +57,7 @@ requires = ["hatchling"]
46
57
  build-backend = "hatchling.build"
47
58
 
48
59
  [tool.hatch.build.targets.wheel]
49
- packages = ["src/hyper_models"]
60
+ packages = ["src/hyper_models", "src/hyper3_clip"]
50
61
 
51
62
  [tool.hatch.build.targets.sdist]
52
63
  include = ["/src", "/README.md", "/LICENSE"]
@@ -0,0 +1,3 @@
1
+ from hyper3_clip.models.hyper3_clip import Hyper3CLIP
2
+
3
+ __all__ = ["Hyper3CLIP"]
@@ -0,0 +1,3 @@
1
+ from hyper3_clip.models.hyper3_clip import Hyper3CLIP
2
+
3
+ __all__ = ["Hyper3CLIP"]
@@ -0,0 +1,181 @@
1
+ from __future__ import annotations
2
+
3
+ import timm
4
+ import torch
5
+ from torch import nn
6
+ from transformers import (
7
+ AutoConfig,
8
+ AutoModel,
9
+ AutoTokenizer,
10
+ CLIPTextConfig,
11
+ CLIPTextModel,
12
+ CLIPTextModelWithProjection,
13
+ CLIPVisionConfig,
14
+ CLIPVisionModel,
15
+ CLIPVisionModelWithProjection,
16
+ SiglipTextConfig,
17
+ SiglipTextModel,
18
+ SiglipVisionConfig,
19
+ SiglipVisionModel,
20
+ )
21
+
22
+
23
+ class VisionEncoder(nn.Module):
24
+ def __init__(self, backbone_name: str, pretrained: bool = True) -> None:
25
+ super().__init__()
26
+ self.kind = "timm"
27
+ if backbone_name.startswith("hf_clip_projected:"):
28
+ self.kind = "hf_clip_projected"
29
+ model_name = backbone_name.removeprefix("hf_clip_projected:")
30
+ self.backbone = (
31
+ CLIPVisionModelWithProjection.from_pretrained(model_name)
32
+ if pretrained
33
+ else CLIPVisionModelWithProjection(CLIPVisionConfig.from_pretrained(model_name))
34
+ )
35
+ self.output_dim = self.backbone.config.projection_dim
36
+ elif backbone_name.startswith("hf_clip:"):
37
+ self.kind = "hf_vision"
38
+ model_name = backbone_name.removeprefix("hf_clip:")
39
+ self.backbone = (
40
+ CLIPVisionModel.from_pretrained(model_name)
41
+ if pretrained
42
+ else CLIPVisionModel(CLIPVisionConfig.from_pretrained(model_name))
43
+ )
44
+ self.output_dim = self.backbone.config.hidden_size
45
+ elif backbone_name.startswith("hf_siglip:"):
46
+ self.kind = "hf_vision"
47
+ model_name = backbone_name.removeprefix("hf_siglip:")
48
+ self.backbone = (
49
+ SiglipVisionModel.from_pretrained(model_name)
50
+ if pretrained
51
+ else SiglipVisionModel(SiglipVisionConfig.from_pretrained(model_name))
52
+ )
53
+ self.output_dim = self.backbone.config.hidden_size
54
+ else:
55
+ self.backbone = timm.create_model(
56
+ backbone_name,
57
+ pretrained=pretrained,
58
+ num_classes=0,
59
+ global_pool="avg",
60
+ )
61
+ self.output_dim = self.backbone.num_features
62
+
63
+ def forward(self, image: torch.Tensor) -> torch.Tensor:
64
+ if self.kind == "hf_clip_projected":
65
+ return self.backbone(pixel_values=image).image_embeds
66
+ if self.kind == "hf_vision":
67
+ out = self.backbone(pixel_values=image)
68
+ if hasattr(out, "pooler_output") and out.pooler_output is not None:
69
+ return out.pooler_output
70
+ return out.last_hidden_state[:, 0]
71
+ return self.backbone(image)
72
+
73
+ def forward_with_tokens(self, image: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
74
+ if self.kind == "hf_clip_projected":
75
+ out = self.backbone(pixel_values=image)
76
+ tokens = getattr(out, "last_hidden_state", None)
77
+ if tokens is None and hasattr(out, "vision_model_output"):
78
+ tokens = out.vision_model_output.last_hidden_state
79
+ if tokens is None:
80
+ raise RuntimeError("Projected CLIP vision output did not include patch tokens")
81
+ return out.image_embeds, tokens
82
+ if self.kind == "hf_vision":
83
+ out = self.backbone(pixel_values=image)
84
+ if hasattr(out, "pooler_output") and out.pooler_output is not None:
85
+ pooled = out.pooler_output
86
+ else:
87
+ pooled = out.last_hidden_state[:, 0]
88
+ return pooled, out.last_hidden_state
89
+
90
+ if not hasattr(self.backbone, "forward_features"):
91
+ pooled = self.backbone(image)
92
+ return pooled, pooled[:, None, :]
93
+ features = self.backbone.forward_features(image)
94
+ if hasattr(self.backbone, "forward_head"):
95
+ pooled = self.backbone.forward_head(features, pre_logits=False)
96
+ else:
97
+ pooled = self.backbone(image)
98
+ return pooled, _tokens_from_features(features)
99
+
100
+
101
+ class TextEncoder(nn.Module):
102
+ def __init__(self, model_name: str, pretrained: bool = True, pooling: str = "auto") -> None:
103
+ super().__init__()
104
+ if pooling not in {"auto", "pooler", "cls", "mean"}:
105
+ raise ValueError(
106
+ f"Unsupported text pooling {pooling!r}; expected auto, pooler, cls, or mean"
107
+ )
108
+ self.kind = "hf_text"
109
+ self.pooling = pooling
110
+ tokenizer_name = model_name.removeprefix("hf_clip_projected:")
111
+ self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
112
+ model_name_lower = model_name.lower()
113
+ if model_name.startswith("hf_clip_projected:"):
114
+ self.kind = "hf_clip_projected"
115
+ projected_model_name = model_name.removeprefix("hf_clip_projected:")
116
+ if pretrained:
117
+ self.backbone = CLIPTextModelWithProjection.from_pretrained(projected_model_name)
118
+ else:
119
+ self.backbone = CLIPTextModelWithProjection(
120
+ CLIPTextConfig.from_pretrained(projected_model_name)
121
+ )
122
+ self.output_dim = self.backbone.config.projection_dim
123
+ elif "siglip" in model_name_lower:
124
+ if pretrained:
125
+ self.backbone = SiglipTextModel.from_pretrained(model_name)
126
+ else:
127
+ self.backbone = SiglipTextModel(SiglipTextConfig.from_pretrained(model_name))
128
+ self.output_dim = self.backbone.config.hidden_size
129
+ elif "clip" in model_name_lower:
130
+ if pretrained:
131
+ self.backbone = CLIPTextModel.from_pretrained(model_name)
132
+ else:
133
+ self.backbone = CLIPTextModel(CLIPTextConfig.from_pretrained(model_name))
134
+ self.output_dim = self.backbone.config.hidden_size
135
+ else:
136
+ if pretrained:
137
+ self.backbone = AutoModel.from_pretrained(model_name)
138
+ else:
139
+ self.backbone = AutoModel.from_config(AutoConfig.from_pretrained(model_name))
140
+ hidden_size = getattr(self.backbone.config, "hidden_size", None)
141
+ if hidden_size is None:
142
+ raise ValueError(f"Unsupported text model config for {model_name}")
143
+ self.output_dim = hidden_size
144
+
145
+ def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
146
+ out = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
147
+ if self.kind == "hf_clip_projected":
148
+ return out.text_embeds
149
+ if self.pooling == "mean":
150
+ mask = attention_mask.to(dtype=out.last_hidden_state.dtype).unsqueeze(-1)
151
+ summed = (out.last_hidden_state * mask).sum(dim=1)
152
+ denom = mask.sum(dim=1).clamp_min(1.0)
153
+ return summed / denom
154
+ if (
155
+ self.pooling in {"auto", "pooler"}
156
+ and hasattr(out, "pooler_output")
157
+ and out.pooler_output is not None
158
+ ):
159
+ return out.pooler_output
160
+ return out.last_hidden_state[:, 0]
161
+
162
+
163
+ def _tokens_from_features(features: torch.Tensor | dict | tuple | list) -> torch.Tensor:
164
+ if isinstance(features, dict):
165
+ for key in ("x", "last_hidden_state", "features"):
166
+ if key in features:
167
+ features = features[key]
168
+ break
169
+ else:
170
+ features = next(iter(features.values()))
171
+ if isinstance(features, tuple | list):
172
+ features = features[0]
173
+ if not torch.is_tensor(features):
174
+ raise TypeError(f"Expected tensor features, got {type(features)!r}")
175
+ if features.ndim == 4:
176
+ return features.flatten(2).transpose(1, 2)
177
+ if features.ndim == 3:
178
+ return features
179
+ if features.ndim == 2:
180
+ return features[:, None, :]
181
+ raise ValueError(f"Unsupported feature tensor shape {tuple(features.shape)}")