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.
- {hyper_models-0.2.0 → hyper_models-0.3.1}/PKG-INFO +18 -7
- {hyper_models-0.2.0 → hyper_models-0.3.1}/README.md +8 -4
- {hyper_models-0.2.0 → hyper_models-0.3.1}/pyproject.toml +14 -3
- hyper_models-0.3.1/src/hyper3_clip/__init__.py +3 -0
- hyper_models-0.3.1/src/hyper3_clip/models/__init__.py +3 -0
- hyper_models-0.3.1/src/hyper3_clip/models/encoders.py +181 -0
- hyper_models-0.3.1/src/hyper3_clip/models/experimental.py +691 -0
- hyper_models-0.3.1/src/hyper3_clip/models/himo.py +54 -0
- hyper_models-0.3.1/src/hyper3_clip/models/hyper3_clip.py +1032 -0
- hyper_models-0.3.1/src/hyper3_clip/models/lorentz.py +285 -0
- hyper_models-0.3.1/src/hyper3_clip/models/losses.py +1546 -0
- hyper_models-0.3.1/src/hyper3_clip/models/objectives.py +653 -0
- hyper_models-0.3.1/src/hyper3_clip/models/tren.py +282 -0
- hyper_models-0.3.1/src/hyper3_clip/training/__init__.py +1 -0
- hyper_models-0.3.1/src/hyper3_clip/training/distributed.py +153 -0
- {hyper_models-0.2.0 → hyper_models-0.3.1}/src/hyper_models/__init__.py +11 -3
- {hyper_models-0.2.0 → hyper_models-0.3.1}/src/hyper_models/loader.py +7 -1
- {hyper_models-0.2.0 → hyper_models-0.3.1}/src/hyper_models/loaders.py +12 -0
- {hyper_models-0.2.0 → hyper_models-0.3.1}/src/hyper_models/models.py +3 -1
- {hyper_models-0.2.0 → hyper_models-0.3.1}/src/hyper_models/registry.py +21 -0
- {hyper_models-0.2.0 → hyper_models-0.3.1}/src/hyper_models/torch_models.py +139 -2
- {hyper_models-0.2.0 → hyper_models-0.3.1}/.gitignore +0 -0
- {hyper_models-0.2.0 → hyper_models-0.3.1}/LICENSE +0 -0
- {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.
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
2
|
Name: hyper-models
|
|
3
|
-
Version: 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
|
|
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
|
|
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` | [](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` | [](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` | [](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` | [](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
|
|
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` | [](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` | [](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` | [](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` | [](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.
|
|
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
|
-
|
|
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,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)}")
|