hyper-models 0.1.0__tar.gz → 0.3.0__tar.gz

{hyper_models-0.1.0 → hyper_models-0.3.0}/PKG-INFO

@@ -1,12 +1,12 @@
 Metadata-Version: 2.4
 Name: hyper-models
-Version: 0.1.0
+Version: 0.3.0
 Summary: A model zoo for non-Euclidean embedding models (hyperbolic, spherical, product manifolds)
 Project-URL: Homepage, https://github.com/Hyper3Labs/hyper-models
 Project-URL: Repository, https://github.com/Hyper3Labs/hyper-models
 Project-URL: Documentation, https://github.com/Hyper3Labs/hyper-models#readme
 Project-URL: Issues, https://github.com/Hyper3Labs/hyper-models/issues
-Author: Hyper3Labs
+Author: hyper3labs
 License: MIT
 License-File: LICENSE
 Keywords: clip,embeddings,hyperbolic,model-zoo,non-euclidean,onnx
@@ -27,6 +27,14 @@ Requires-Dist: pillow>=10.0
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == 'dev'
 Requires-Dist: ruff>=0.1; extra == 'dev'
+Provides-Extra: ml
+Requires-Dist: protobuf>=4.25.0; extra == 'ml'
+Requires-Dist: pyyaml>=6.0.0; extra == 'ml'
+Requires-Dist: safetensors>=0.4.0; extra == 'ml'
+Requires-Dist: sentencepiece>=0.2.0; extra == 'ml'
+Requires-Dist: timm>=1.0.0; extra == 'ml'
+Requires-Dist: torch>=2.9.1; extra == 'ml'
+Requires-Dist: transformers<6.0,>=5.0.0; extra == 'ml'
 Description-Content-Type: text/markdown
 
 # hyper-models
@@ -51,13 +59,22 @@ Description-Content-Type: text/markdown
 ## Why?
 
 - **Standardized access** to non-Euclidean embedding models
-- **Torch-free runtime** via ONNX (models published to Hugging Face Hub)
+- **One catalog surface**: model names map to internal loaders such as ONNX or optional torch-backed runtimes
 - **Simple API** — `load()` and `encode_images()`
 
 ## Installation
 
 ```bash
-pip install hyper-models
+uv pip install hyper-models
+```
+
+This base install is the simple path: it stays **torch-free** and is enough for
+ONNX-backed catalog entries such as HyCoCLIP and MERU.
+
+For torch-backed checkpoints (for example UNCHA and Hyper3-CLIP):
+
+```bash
+uv pip install "hyper-models[ml]"
 ```
 
 ## Usage
@@ -68,7 +85,11 @@ from PIL import Image
 
 # List available models
 hyper_models.list_models()
-# ['hycoclip-vit-s', 'hycoclip-vit-b', 'meru-vit-s', 'meru-vit-b']
+# ['hycoclip-vit-s', 'hycoclip-vit-b', 'meru-vit-s', 'meru-vit-b', 'uncha-vit-s', 'uncha-vit-b', 'hyper3-clip-v0.5']
+
+# Inspect supported internal loader kinds
+hyper_models.list_loaders()
+# ['hyper3-clip-torch', 'onnx', 'uncha-image-torch']
 
 # Load model (auto-downloads from Hugging Face Hub)
 model = hyper_models.load("hycoclip-vit-s")
@@ -82,6 +103,7 @@ embeddings = model.encode_images(images)  # (1, 513) ndarray
 # Get model info
 info = hyper_models.get_model_info("hycoclip-vit-s")
 info.hub_id     # 'mnm-matin/hyperbolic-clip'
+info.loader     # 'onnx'
 info.license    # 'CC-BY-NC'
 
 # Low-level: preprocess images yourself
@@ -89,6 +111,54 @@ batch = hyper_models.preprocess_images(images)  # (B, 3, 224, 224)
 embeddings = model.encode(batch)
 ```
 
+### Architecture
+
+`hyper-models` is intended to be a timm-like catalog for non-Euclidean models.
+
+- The public abstraction is the catalog entry name, for example `hycoclip-vit-s`.
+- Each entry declares metadata such as geometry, dimensionality, artifact path,
+  and an internal loader kind.
+- Internal loaders may differ by model family:
+  - `onnx` for exported, torch-free runtimes
+  - `uncha-image-torch` for raw checkpoints that need a PyTorch image runtime
+  - `hyper3-clip-torch` for Hyper3-CLIP safetensors checkpoints
+
+This keeps callers on one stable API:
+
+```python
+model = hyper_models.load("hycoclip-vit-s")
+model = hyper_models.load("uncha-vit-b")
+model = hyper_models.load("hyper3-clip-v0.5")
+```
+
+Callers do not need to know which internal loader is used, except for optional
+dependency installation when choosing entries that need `hyper-models[ml]`.
+
+### HyperView integration
+
+HyperView auto-detects `hyper-models` names and routes them to the `hyper-models` provider.
+
+```python
+import hyperview as hv
+
+dataset = hv.Dataset.from_huggingface(
+  name="demo",
+  hf_dataset="uoft-cs/cifar10",
+  split="train",
+  image_key="img",
+)
+
+# Uses provider='hyper-models' automatically.
+space_key = dataset.compute_embeddings(model="uncha-vit-b")
+layout_key = dataset.compute_visualization(space_key=space_key, layout="poincare")
+```
+
+HyperView's simple path remains torch-free. If you use the default ONNX-backed
+`hyper-models` entries or the default `embed-anything` provider, HyperView does
+not need PyTorch. PyTorch is only needed when you explicitly select a
+torch-backed catalog entry such as `uncha-vit-s`, `uncha-vit-b`, or
+`hyper3-clip-v0.5`.
+
 ## Models
 
 ### Hyperbolic
@@ -99,16 +169,21 @@ embeddings = model.encode(batch)
 | `hycoclip-vit-b` | [![HF](https://img.shields.io/badge/🤗-HuggingFace-yellow)](https://huggingface.co/mnm-matin/hyperbolic-clip/tree/main/hycoclip-vit-b) | [ICLR 2025](https://arxiv.org/abs/2410.06912) | [PalAvik/hycoclip](https://github.com/PalAvik/hycoclip) |
 | `meru-vit-s` | [![HF](https://img.shields.io/badge/🤗-HuggingFace-yellow)](https://huggingface.co/mnm-matin/hyperbolic-clip/tree/main/meru-vit-s) | [ICML 2023](https://arxiv.org/abs/2304.09172) | [facebookresearch/meru](https://github.com/facebookresearch/meru) |
 | `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) |
+| `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) |
+| `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) |
+| `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) |
 | `hyp-vit` | — | [CVPR 2022](https://arxiv.org/abs/2203.10833) | [htdt/hyp_metric](https://github.com/htdt/hyp_metric) |
 | `hie` | — | [CVPR 2020](https://arxiv.org/abs/1904.02239) | [leymir/hyperbolic-image-embeddings](https://github.com/leymir/hyperbolic-image-embeddings) |
 | `hcnn` | — | [ICLR 2024](https://openreview.net/forum?id=ekz1hN5QNh) | [kschwethelm/HyperbolicCV](https://github.com/kschwethelm/HyperbolicCV) |
 
-### Spherical
+### Hyperspherical
+
+| Model            | Available | Paper | Code |
+|------------------|:---------:|-------|------|
+| `megadescriptor` (via timm) | [![HF](https://img.shields.io/badge/🤗-HuggingFace-yellow)](https://huggingface.co/BVRA/MegaDescriptor-L-384) | [WACV 2024](https://openaccess.thecvf.com/content/WACV2024/papers/Cermak_WildlifeDatasets_An_Open-Source_Toolkit_for_Animal_Re-Identification_WACV_2024_paper.pdf) | [WildlifeDatasets/wildlife-datasets](https://github.com/WildlifeDatasets/wildlife-datasets) |
+| `sphereface`     | — | [CVPR 2017](https://arxiv.org/abs/1704.08063) | [wy1iu/sphereface](https://github.com/wy1iu/sphereface) |
+| `arcface`       | — | [CVPR 2019](https://arxiv.org/abs/1801.07698) | [deepinsight/insightface](https://github.com/deepinsight/insightface) |
 
-| Model | Available | Paper | Code |
-|-------|:---------:|-------|------|
-| `sphereface` | — | [CVPR 2017](https://arxiv.org/abs/1704.08063) | [wy1iu/sphereface](https://github.com/wy1iu/sphereface) |
-| `arcface` | — | [CVPR 2019](https://arxiv.org/abs/1801.07698) | [deepinsight/insightface](https://github.com/deepinsight/insightface) |
 
 ### Product Manifolds
 

hyper_models-0.3.0/README.md

@@ -0,0 +1,170 @@
+# hyper-models
+
+<p align="center">
+  <strong>A model zoo for non-Euclidean embedding models</strong>
+  <br>
+  <em>Hyperbolic · Spherical · Product Manifolds</em>
+</p>
+
+<p align="center">
+  <a href="https://huggingface.co/mnm-matin/hyperbolic-clip">
+    <img src="https://img.shields.io/badge/🤗_Models-hyperbolic--clip-orange" alt="Hugging Face">
+  </a>
+  <a href="LICENSE">
+    <img src="https://img.shields.io/badge/License-MIT-blue" alt="License: MIT">
+  </a>
+</p>
+
+---
+
+## Why?
+
+- **Standardized access** to non-Euclidean embedding models
+- **One catalog surface**: model names map to internal loaders such as ONNX or optional torch-backed runtimes
+- **Simple API** — `load()` and `encode_images()`
+
+## Installation
+
+```bash
+uv pip install hyper-models
+```
+
+This base install is the simple path: it stays **torch-free** and is enough for
+ONNX-backed catalog entries such as HyCoCLIP and MERU.
+
+For torch-backed checkpoints (for example UNCHA and Hyper3-CLIP):
+
+```bash
+uv pip install "hyper-models[ml]"
+```
+
+## Usage
+
+```python
+import hyper_models
+from PIL import Image
+
+# List available models
+hyper_models.list_models()
+# ['hycoclip-vit-s', 'hycoclip-vit-b', 'meru-vit-s', 'meru-vit-b', 'uncha-vit-s', 'uncha-vit-b', 'hyper3-clip-v0.5']
+
+# Inspect supported internal loader kinds
+hyper_models.list_loaders()
+# ['hyper3-clip-torch', 'onnx', 'uncha-image-torch']
+
+# Load model (auto-downloads from Hugging Face Hub)
+model = hyper_models.load("hycoclip-vit-s")
+model.geometry  # 'hyperboloid'
+model.dim       # 513
+
+# Encode PIL images
+images = [Image.open("image.jpg")]
+embeddings = model.encode_images(images)  # (1, 513) ndarray
+
+# Get model info
+info = hyper_models.get_model_info("hycoclip-vit-s")
+info.hub_id     # 'mnm-matin/hyperbolic-clip'
+info.loader     # 'onnx'
+info.license    # 'CC-BY-NC'
+
+# Low-level: preprocess images yourself
+batch = hyper_models.preprocess_images(images)  # (B, 3, 224, 224)
+embeddings = model.encode(batch)
+```
+
+### Architecture
+
+`hyper-models` is intended to be a timm-like catalog for non-Euclidean models.
+
+- The public abstraction is the catalog entry name, for example `hycoclip-vit-s`.
+- Each entry declares metadata such as geometry, dimensionality, artifact path,
+  and an internal loader kind.
+- Internal loaders may differ by model family:
+  - `onnx` for exported, torch-free runtimes
+  - `uncha-image-torch` for raw checkpoints that need a PyTorch image runtime
+  - `hyper3-clip-torch` for Hyper3-CLIP safetensors checkpoints
+
+This keeps callers on one stable API:
+
+```python
+model = hyper_models.load("hycoclip-vit-s")
+model = hyper_models.load("uncha-vit-b")
+model = hyper_models.load("hyper3-clip-v0.5")
+```
+
+Callers do not need to know which internal loader is used, except for optional
+dependency installation when choosing entries that need `hyper-models[ml]`.
+
+### HyperView integration
+
+HyperView auto-detects `hyper-models` names and routes them to the `hyper-models` provider.
+
+```python
+import hyperview as hv
+
+dataset = hv.Dataset.from_huggingface(
+  name="demo",
+  hf_dataset="uoft-cs/cifar10",
+  split="train",
+  image_key="img",
+)
+
+# Uses provider='hyper-models' automatically.
+space_key = dataset.compute_embeddings(model="uncha-vit-b")
+layout_key = dataset.compute_visualization(space_key=space_key, layout="poincare")
+```
+
+HyperView's simple path remains torch-free. If you use the default ONNX-backed
+`hyper-models` entries or the default `embed-anything` provider, HyperView does
+not need PyTorch. PyTorch is only needed when you explicitly select a
+torch-backed catalog entry such as `uncha-vit-s`, `uncha-vit-b`, or
+`hyper3-clip-v0.5`.
+
+## Models
+
+### Hyperbolic
+
+| Model | Available | Paper | Code |
+|-------|:---------:|-------|------|
+| `hycoclip-vit-s` | [![HF](https://img.shields.io/badge/🤗-HuggingFace-yellow)](https://huggingface.co/mnm-matin/hyperbolic-clip/tree/main/hycoclip-vit-s) | [ICLR 2025](https://arxiv.org/abs/2410.06912) | [PalAvik/hycoclip](https://github.com/PalAvik/hycoclip) |
+| `hycoclip-vit-b` | [![HF](https://img.shields.io/badge/🤗-HuggingFace-yellow)](https://huggingface.co/mnm-matin/hyperbolic-clip/tree/main/hycoclip-vit-b) | [ICLR 2025](https://arxiv.org/abs/2410.06912) | [PalAvik/hycoclip](https://github.com/PalAvik/hycoclip) |
+| `meru-vit-s` | [![HF](https://img.shields.io/badge/🤗-HuggingFace-yellow)](https://huggingface.co/mnm-matin/hyperbolic-clip/tree/main/meru-vit-s) | [ICML 2023](https://arxiv.org/abs/2304.09172) | [facebookresearch/meru](https://github.com/facebookresearch/meru) |
+| `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) |
+| `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) |
+| `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) |
+| `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) |
+| `hyp-vit` | — | [CVPR 2022](https://arxiv.org/abs/2203.10833) | [htdt/hyp_metric](https://github.com/htdt/hyp_metric) |
+| `hie` | — | [CVPR 2020](https://arxiv.org/abs/1904.02239) | [leymir/hyperbolic-image-embeddings](https://github.com/leymir/hyperbolic-image-embeddings) |
+| `hcnn` | — | [ICLR 2024](https://openreview.net/forum?id=ekz1hN5QNh) | [kschwethelm/HyperbolicCV](https://github.com/kschwethelm/HyperbolicCV) |
+
+### Hyperspherical
+
+| Model            | Available | Paper | Code |
+|------------------|:---------:|-------|------|
+| `megadescriptor` (via timm) | [![HF](https://img.shields.io/badge/🤗-HuggingFace-yellow)](https://huggingface.co/BVRA/MegaDescriptor-L-384) | [WACV 2024](https://openaccess.thecvf.com/content/WACV2024/papers/Cermak_WildlifeDatasets_An_Open-Source_Toolkit_for_Animal_Re-Identification_WACV_2024_paper.pdf) | [WildlifeDatasets/wildlife-datasets](https://github.com/WildlifeDatasets/wildlife-datasets) |
+| `sphereface`     | — | [CVPR 2017](https://arxiv.org/abs/1704.08063) | [wy1iu/sphereface](https://github.com/wy1iu/sphereface) |
+| `arcface`       | — | [CVPR 2019](https://arxiv.org/abs/1801.07698) | [deepinsight/insightface](https://github.com/deepinsight/insightface) |
+
+
+### Product Manifolds
+
+| Model | Available | Paper | Code |
+|-------|:---------:|-------|------|
+| `hyperbolics` | — | [ICLR 2019](https://openreview.net/forum?id=HJxeWnCcF7) | [HazyResearch/hyperbolics](https://github.com/HazyResearch/hyperbolics) |
+
+## Export Tooling
+
+This repo also contains tooling to export PyTorch models to ONNX:
+
+```bash
+cd export/hycoclip
+uv run python export_onnx.py --checkpoint model.pth --onnx model.onnx
+```
+
+See [export/hycoclip/README.md](export/hycoclip/README.md) for details.
+
+## References
+
+- [HyCoCLIP](https://github.com/PalAvik/hycoclip)
+- [MERU](https://github.com/facebookresearch/meru)
+- [geoopt](https://github.com/geoopt/geoopt)

{hyper_models-0.1.0 → hyper_models-0.3.0}/pyproject.toml

@@ -1,11 +1,11 @@
 [project]
 name = "hyper-models"
-version = "0.1.0"
+version = "0.3.0"
 description = "A model zoo for non-Euclidean embedding models (hyperbolic, spherical, product manifolds)"
 readme = "README.md"
 license = { text = "MIT" }
 requires-python = ">=3.10"
-authors = [{ name = "Hyper3Labs" }]
+authors = [{ name = "hyper3labs" }]
 keywords = ["embeddings", "hyperbolic", "onnx", "clip", "model-zoo", "non-euclidean"]
 classifiers = [
     "Development Status :: 3 - Alpha",
@@ -30,6 +30,15 @@ dev = [
     "pytest>=7.0",
     "ruff>=0.1",
 ]
+ml = [
+    "torch>=2.9.1",
+    "timm>=1.0.0",
+    "transformers>=5.0.0,<6.0",
+    "safetensors>=0.4.0",
+    "pyyaml>=6.0.0",
+    "sentencepiece>=0.2.0",
+    "protobuf>=4.25.0",
+]
 
 [project.urls]
 Homepage = "https://github.com/Hyper3Labs/hyper-models"
@@ -42,7 +51,7 @@ requires = ["hatchling"]
 build-backend = "hatchling.build"
 
 [tool.hatch.build.targets.wheel]
-packages = ["src/hyper_models"]
+packages = ["src/hyper_models", "src/hyper3_clip"]
 
 [tool.hatch.build.targets.sdist]
 include = ["/src", "/README.md", "/LICENSE"]

hyper_models-0.3.0/src/hyper3_clip/__init__.py

@@ -0,0 +1,3 @@
+from hyper3_clip.models.hyper3_clip import Hyper3CLIP
+
+__all__ = ["Hyper3CLIP"]

hyper_models-0.3.0/src/hyper3_clip/models/__init__.py

@@ -0,0 +1,3 @@
+from hyper3_clip.models.hyper3_clip import Hyper3CLIP
+
+__all__ = ["Hyper3CLIP"]

hyper_models-0.3.0/src/hyper3_clip/models/encoders.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+import timm
+import torch
+from torch import nn
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoTokenizer,
+    CLIPTextConfig,
+    CLIPTextModel,
+    CLIPTextModelWithProjection,
+    CLIPVisionConfig,
+    CLIPVisionModel,
+    CLIPVisionModelWithProjection,
+    SiglipTextConfig,
+    SiglipTextModel,
+    SiglipVisionConfig,
+    SiglipVisionModel,
+)
+
+
+class VisionEncoder(nn.Module):
+    def __init__(self, backbone_name: str, pretrained: bool = True) -> None:
+        super().__init__()
+        self.kind = "timm"
+        if backbone_name.startswith("hf_clip_projected:"):
+            self.kind = "hf_clip_projected"
+            model_name = backbone_name.removeprefix("hf_clip_projected:")
+            self.backbone = (
+                CLIPVisionModelWithProjection.from_pretrained(model_name)
+                if pretrained
+                else CLIPVisionModelWithProjection(CLIPVisionConfig.from_pretrained(model_name))
+            )
+            self.output_dim = self.backbone.config.projection_dim
+        elif backbone_name.startswith("hf_clip:"):
+            self.kind = "hf_vision"
+            model_name = backbone_name.removeprefix("hf_clip:")
+            self.backbone = (
+                CLIPVisionModel.from_pretrained(model_name)
+                if pretrained
+                else CLIPVisionModel(CLIPVisionConfig.from_pretrained(model_name))
+            )
+            self.output_dim = self.backbone.config.hidden_size
+        elif backbone_name.startswith("hf_siglip:"):
+            self.kind = "hf_vision"
+            model_name = backbone_name.removeprefix("hf_siglip:")
+            self.backbone = (
+                SiglipVisionModel.from_pretrained(model_name)
+                if pretrained
+                else SiglipVisionModel(SiglipVisionConfig.from_pretrained(model_name))
+            )
+            self.output_dim = self.backbone.config.hidden_size
+        else:
+            self.backbone = timm.create_model(
+                backbone_name,
+                pretrained=pretrained,
+                num_classes=0,
+                global_pool="avg",
+            )
+            self.output_dim = self.backbone.num_features
+
+    def forward(self, image: torch.Tensor) -> torch.Tensor:
+        if self.kind == "hf_clip_projected":
+            return self.backbone(pixel_values=image).image_embeds
+        if self.kind == "hf_vision":
+            out = self.backbone(pixel_values=image)
+            if hasattr(out, "pooler_output") and out.pooler_output is not None:
+                return out.pooler_output
+            return out.last_hidden_state[:, 0]
+        return self.backbone(image)
+
+    def forward_with_tokens(self, image: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.kind == "hf_clip_projected":
+            out = self.backbone(pixel_values=image)
+            tokens = getattr(out, "last_hidden_state", None)
+            if tokens is None and hasattr(out, "vision_model_output"):
+                tokens = out.vision_model_output.last_hidden_state
+            if tokens is None:
+                raise RuntimeError("Projected CLIP vision output did not include patch tokens")
+            return out.image_embeds, tokens
+        if self.kind == "hf_vision":
+            out = self.backbone(pixel_values=image)
+            if hasattr(out, "pooler_output") and out.pooler_output is not None:
+                pooled = out.pooler_output
+            else:
+                pooled = out.last_hidden_state[:, 0]
+            return pooled, out.last_hidden_state
+
+        if not hasattr(self.backbone, "forward_features"):
+            pooled = self.backbone(image)
+            return pooled, pooled[:, None, :]
+        features = self.backbone.forward_features(image)
+        if hasattr(self.backbone, "forward_head"):
+            pooled = self.backbone.forward_head(features, pre_logits=False)
+        else:
+            pooled = self.backbone(image)
+        return pooled, _tokens_from_features(features)
+
+
+class TextEncoder(nn.Module):
+    def __init__(self, model_name: str, pretrained: bool = True, pooling: str = "auto") -> None:
+        super().__init__()
+        if pooling not in {"auto", "pooler", "cls", "mean"}:
+            raise ValueError(
+                f"Unsupported text pooling {pooling!r}; expected auto, pooler, cls, or mean"
+            )
+        self.kind = "hf_text"
+        self.pooling = pooling
+        tokenizer_name = model_name.removeprefix("hf_clip_projected:")
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+        model_name_lower = model_name.lower()
+        if model_name.startswith("hf_clip_projected:"):
+            self.kind = "hf_clip_projected"
+            projected_model_name = model_name.removeprefix("hf_clip_projected:")
+            if pretrained:
+                self.backbone = CLIPTextModelWithProjection.from_pretrained(projected_model_name)
+            else:
+                self.backbone = CLIPTextModelWithProjection(
+                    CLIPTextConfig.from_pretrained(projected_model_name)
+                )
+            self.output_dim = self.backbone.config.projection_dim
+        elif "siglip" in model_name_lower:
+            if pretrained:
+                self.backbone = SiglipTextModel.from_pretrained(model_name)
+            else:
+                self.backbone = SiglipTextModel(SiglipTextConfig.from_pretrained(model_name))
+            self.output_dim = self.backbone.config.hidden_size
+        elif "clip" in model_name_lower:
+            if pretrained:
+                self.backbone = CLIPTextModel.from_pretrained(model_name)
+            else:
+                self.backbone = CLIPTextModel(CLIPTextConfig.from_pretrained(model_name))
+            self.output_dim = self.backbone.config.hidden_size
+        else:
+            if pretrained:
+                self.backbone = AutoModel.from_pretrained(model_name)
+            else:
+                self.backbone = AutoModel.from_config(AutoConfig.from_pretrained(model_name))
+            hidden_size = getattr(self.backbone.config, "hidden_size", None)
+            if hidden_size is None:
+                raise ValueError(f"Unsupported text model config for {model_name}")
+            self.output_dim = hidden_size
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        out = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
+        if self.kind == "hf_clip_projected":
+            return out.text_embeds
+        if self.pooling == "mean":
+            mask = attention_mask.to(dtype=out.last_hidden_state.dtype).unsqueeze(-1)
+            summed = (out.last_hidden_state * mask).sum(dim=1)
+            denom = mask.sum(dim=1).clamp_min(1.0)
+            return summed / denom
+        if (
+            self.pooling in {"auto", "pooler"}
+            and hasattr(out, "pooler_output")
+            and out.pooler_output is not None
+        ):
+            return out.pooler_output
+        return out.last_hidden_state[:, 0]
+
+
+def _tokens_from_features(features: torch.Tensor | dict | tuple | list) -> torch.Tensor:
+    if isinstance(features, dict):
+        for key in ("x", "last_hidden_state", "features"):
+            if key in features:
+                features = features[key]
+                break
+        else:
+            features = next(iter(features.values()))
+    if isinstance(features, tuple | list):
+        features = features[0]
+    if not torch.is_tensor(features):
+        raise TypeError(f"Expected tensor features, got {type(features)!r}")
+    if features.ndim == 4:
+        return features.flatten(2).transpose(1, 2)
+    if features.ndim == 3:
+        return features
+    if features.ndim == 2:
+        return features[:, None, :]
+    raise ValueError(f"Unsupported feature tensor shape {tuple(features.shape)}")