Release Vela Omni Nano with multimodal embeddings and evaluation scores

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

NOTICE

@@ -0,0 +1,25 @@
+Vela Omni
+
+This artifact includes components from Google SigLIP (Apache License 2.0), Sentence Transformers (Apache License 2.0), and OpenAI Whisper. OpenAI Whisper is distributed under the following MIT notice.
+
+MIT License
+
+Copyright (c) 2022 OpenAI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,81 @@
+---
+library_name: pytorch
+license: apache-2.0
+pipeline_tag: feature-extraction
+tags:
+- vela
+- semantic-router
+- multimodal
+- image-text-retrieval
+- audio-text-retrieval
+---
+
+<div align="center">
+  <img src="https://vllm-sr.ai/img/vllm-sr-logo.social.png" alt="vLLM Semantic Router" width="560" />
+  <p>
+    <a href="https://vllm-sr.ai/"><strong>Docs</strong></a> |
+    <a href="https://vllm-sr.ai/blog/"><strong>Blog</strong></a> |
+    <a href="https://vllm-dev.slack.com/archives/C09CTGF8KCN"><strong>Slack</strong></a> |
+    <a href="https://github.com/vllm-project/semantic-router"><strong>GitHub</strong></a>
+  </p>
+</div>
+
+# Vela Omni Nano
+
+Vela Omni Nano maps text, images, and speech into a shared embedding space for multimodal search and matching.
+
+**134M parameters · 384 dimensions · L2-normalized embeddings.**
+
+[Try it in Vela Studio](https://huggingface.co/spaces/llm-semantic-router/vela-studio).
+
+## Evaluation
+
+Scores use a 0–100 scale; higher is better. All applicable models use the same examples and retrieval pools. N/A denotes a modality the text-only model does not support.
+
+| Metric | [Vela-1.0-Encoder-307M-Embedding](https://huggingface.co/llm-semantic-router/Vela-1.0-Encoder-307M-Embedding) | [multi-modal-embed-small](https://huggingface.co/llm-semantic-router/multi-modal-embed-small) | [multi-modal-embed-large](https://huggingface.co/llm-semantic-router/multi-modal-embed-large) | [Vela-1.0-Omni-134M-Nano](https://huggingface.co/llm-semantic-router/Vela-1.0-Omni-134M-Nano) |
+|---|---:|---:|---:|---:|
+| Banking77 · Accuracy | 80.00 | 70.42 | 75.78 | 70.42 |
+| MASSIVE English · Accuracy | 75.64 | 65.95 | 72.31 | 65.95 |
+| COCO · Image → text · R@1 | N/A | 40.83 | 42.53 | 44.23 |
+| COCO · Image → text · R@5 | N/A | 67.19 | 75.21 | 73.03 |
+| COCO · Image → text · R@10 | N/A | 78.49 | 87.61 | 86.03 |
+| COCO · Text → image · R@1 | N/A | 30.18 | 35.04 | 36.11 |
+| COCO · Text → image · R@5 | N/A | 59.42 | 70.09 | 68.72 |
+| COCO · Text → image · R@10 | N/A | 74.29 | 83.91 | 82.02 |
+| LibriSpeech · Audio → text · R@1 | N/A | 4.21 | 56.99 | 5.29 |
+| LibriSpeech · Audio → text · R@5 | N/A | 11.99 | 81.85 | 15.82 |
+| LibriSpeech · Audio → text · R@10 | N/A | 19.03 | 87.94 | 23.90 |
+| LibriSpeech · Text → audio · R@1 | N/A | 9.58 | 78.58 | 11.61 |
+| LibriSpeech · Text → audio · R@5 | N/A | 22.53 | 94.02 | 26.55 |
+| LibriSpeech · Text → audio · R@10 | N/A | 30.69 | 97.01 | 35.44 |
+
+Text evaluation uses fixed class prototypes: 3,080 Banking77 and 2,972 MASSIVE English queries. This is nearest-prototype classification, not the MTEB classification protocol.
+
+Image retrieval uses the COCO Karpathy CC-BY2 image subset: 823 images and 4,115 captions. Speech retrieval uses 2,611 LibriSpeech test-clean clips and 2,610 unique transcripts, with clips limited to 30 seconds. Recall is measured against each complete candidate pool, accepting all matching positives.
+
+Evaluation uses FP32 embeddings, cosine similarity, and a common 128-token text limit. [Full scores and protocol](./scores.json) include model revisions, sample counts, and paired uncertainty.
+
+## Quick start
+
+Use PyTorch, Transformers 4.57.6, Hugging Face Hub, safetensors, NumPy, and Pillow:
+
+```python
+import sys
+from huggingface_hub import snapshot_download
+
+path = snapshot_download("llm-semantic-router/Vela-1.0-Omni-134M-Nano")
+sys.path.insert(0, path)
+from vela_omni import VelaOmni
+
+model = VelaOmni.from_pretrained(path, device="cpu")
+vectors = model.encode_text(["A bicycle beside a building.", "Someone is reading aloud."])
+print(vectors.shape)  # (2, 384)
+```
+
+Pass a list of Pillow images to `model.encode_image(images)`. Pass a list of mono NumPy waveforms to `model.encode_audio(waveforms, sampling_rate=16000)`; each waveform must be at most 30 seconds. Compare normalized vectors with their dot product. Similarity scores are rankings, not calibrated probabilities.
+
+Cross-modal alignment uses COCO image–caption pairs from the CC-BY 2.0 image subset and [LibriSpeech](https://www.openslr.org/12/) speech–transcript pairs (CC-BY 4.0). COCO annotations are provided under CC-BY 4.0.
+
+The repository includes the native model code, component configurations, and tokenizer and processor files. See [NOTICE](./NOTICE) and [LICENSE](./LICENSE) for component attribution and license terms.
+
+[Explore the Vela model collection](https://huggingface.co/collections/llm-semantic-router/vela-10-router-models-6aa555ba70cc6997d6d67798)

components/audio/config.json

@@ -0,0 +1,143 @@
+{
+  "activation_dropout": 0.0,
+  "activation_function": "gelu",
+  "architectures": [
+    "WhisperForConditionalGeneration"
+  ],
+  "attention_dropout": 0.0,
+  "begin_suppress_tokens": [
+    220,
+    50257
+  ],
+  "bos_token_id": 50257,
+  "d_model": 384,
+  "decoder_attention_heads": 6,
+  "decoder_ffn_dim": 1536,
+  "decoder_layerdrop": 0.0,
+  "decoder_layers": 4,
+  "decoder_start_token_id": 50258,
+  "dropout": 0.0,
+  "encoder_attention_heads": 6,
+  "encoder_ffn_dim": 1536,
+  "encoder_layerdrop": 0.0,
+  "encoder_layers": 4,
+  "eos_token_id": 50257,
+  "forced_decoder_ids": [
+    [
+      1,
+      50259
+    ],
+    [
+      2,
+      50359
+    ],
+    [
+      3,
+      50363
+    ]
+  ],
+  "init_std": 0.02,
+  "is_encoder_decoder": true,
+  "max_length": 448,
+  "max_source_positions": 1500,
+  "max_target_positions": 448,
+  "model_type": "whisper",
+  "num_hidden_layers": 4,
+  "num_mel_bins": 80,
+  "pad_token_id": 50257,
+  "scale_embedding": false,
+  "suppress_tokens": [
+    1,
+    2,
+    7,
+    8,
+    9,
+    10,
+    14,
+    25,
+    26,
+    27,
+    28,
+    29,
+    31,
+    58,
+    59,
+    60,
+    61,
+    62,
+    63,
+    90,
+    91,
+    92,
+    93,
+    359,
+    503,
+    522,
+    542,
+    873,
+    893,
+    902,
+    918,
+    922,
+    931,
+    1350,
+    1853,
+    1982,
+    2460,
+    2627,
+    3246,
+    3253,
+    3268,
+    3536,
+    3846,
+    3961,
+    4183,
+    4667,
+    6585,
+    6647,
+    7273,
+    9061,
+    9383,
+    10428,
+    10929,
+    11938,
+    12033,
+    12331,
+    12562,
+    13793,
+    14157,
+    14635,
+    15265,
+    15618,
+    16553,
+    16604,
+    18362,
+    18956,
+    20075,
+    21675,
+    22520,
+    26130,
+    26161,
+    26435,
+    28279,
+    29464,
+    31650,
+    32302,
+    32470,
+    36865,
+    42863,
+    47425,
+    49870,
+    50254,
+    50258,
+    50358,
+    50359,
+    50360,
+    50361,
+    50362
+  ],
+  "torch_dtype": "float32",
+  "transformers_version": "4.27.0.dev0",
+  "use_cache": true,
+  "vocab_size": 51865
+}

components/image/config.json

@@ -0,0 +1,16 @@
+{
+  "architectures": [
+    "SiglipModel"
+  ],
+  "initializer_factor": 1.0,
+  "model_type": "siglip",
+  "text_config": {
+    "model_type": "siglip_text_model"
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.37.0.dev0",
+  "vision_config": {
+    "image_size": 512,
+    "model_type": "siglip_vision_model"
+  }
+}

components/image/preprocessor_config.json

@@ -0,0 +1,23 @@
+{
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_processor_type": "SiglipImageProcessor",
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "processor_class": "SiglipProcessor",
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "height": 512,
+    "width": 512
+  }
+}

components/text/1_Pooling/config.json

@@ -0,0 +1,7 @@
+{
+  "word_embedding_dimension": 384,
+  "pooling_mode_cls_token": false,
+  "pooling_mode_mean_tokens": true,
+  "pooling_mode_max_tokens": false,
+  "pooling_mode_mean_sqrt_len_tokens": false
+}

components/text/config.json

@@ -0,0 +1,23 @@
+{
+  "architectures": [
+    "BertModel"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "gradient_checkpointing": false,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 384,
+  "initializer_range": 0.02,
+  "intermediate_size": 1536,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 6,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "transformers_version": "4.8.2",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 30522
+}

components/text/modules.json

@@ -0,0 +1,20 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.models.Transformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_Pooling",
+    "type": "sentence_transformers.models.Pooling"
+  },
+  {
+    "idx": 2,
+    "name": "2",
+    "path": "2_Normalize",
+    "type": "sentence_transformers.models.Normalize"
+  }
+]

components/text/sentence_bert_config.json

@@ -0,0 +1,4 @@
+{
+  "max_seq_length": 256,
+  "do_lower_case": false
+}

components/text/special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "unk_token": "[UNK]",
+  "sep_token": "[SEP]",
+  "pad_token": "[PAD]",
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]"
+}

components/text/tokenizer_config.json

@@ -0,0 +1,14 @@
+{
+  "do_lower_case": true,
+  "unk_token": "[UNK]",
+  "sep_token": "[SEP]",
+  "pad_token": "[PAD]",
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "tokenize_chinese_chars": true,
+  "strip_accents": null,
+  "do_basic_tokenize": true,
+  "never_split": null,
+  "tokenizer_class": "BertTokenizer",
+  "model_max_length": 512
+}

config.json

@@ -0,0 +1,15 @@
+{
+  "format_version": 2,
+  "variant": "nano",
+  "parameter_count": 133907328,
+  "embedding_dim": 384,
+  "max_text_length": 128,
+  "architectures": [
+    "VelaOmni"
+  ],
+  "torch_dtype": "float32",
+  "library_name": "pytorch",
+  "audio_sampling_rate": 16000,
+  "audio_max_seconds": 30,
+  "audio_projection": "linear"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa28600e25bdfb3796fc173fd9504c5a4687c07211d7e3dbb751dceef1a8c341
+size 535689632

omni_components/__init__.py

@@ -0,0 +1 @@
+from .embedder import MultimodalEmbedder

omni_components/audio_encoder.py

@@ -0,0 +1,223 @@
+"""Audio encoder based on Whisper architecture."""
+
+import logging
+from typing import Any
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional
+from transformers import WhisperConfig, WhisperFeatureExtractor, WhisperModel
+
+logger = logging.getLogger(__name__)
+
+
+class AudioEncoder(nn.Module):
+    """
+    Audio encoder using Whisper encoder architecture.
+
+    Whisper's encoder is a sequential transformer stack, making it
+    compatible with adaptive layer exit (2DMSE).
+    """
+
+    def __init__(self, model_name_or_path, revision=None, output_dim=384,
+                 pooling_mode="mean", normalize=True, enable_layer_outputs=True):
+        super().__init__()
+        from transformers.models.whisper.modeling_whisper import WhisperEncoder
+        self.model_name = str(model_name_or_path)
+        self.revision = revision
+        self.output_dim = output_dim
+        self.pooling_mode = pooling_mode
+        self.normalize = normalize
+        self.enable_layer_outputs = enable_layer_outputs
+        self.config = WhisperConfig.from_pretrained(model_name_or_path, local_files_only=True)
+        self.config._attn_implementation = "sdpa"
+        self.encoder = WhisperEncoder(self.config)
+        self.feature_extractor = WhisperFeatureExtractor.from_pretrained(
+            model_name_or_path, local_files_only=True)
+        self.hidden_size = self.config.d_model
+        self.num_layers = self.config.encoder_layers
+        self.projection = nn.Linear(self.hidden_size, output_dim) if self.hidden_size != output_dim else nn.Identity()
+        self.layer_projections = nn.ModuleList([
+            nn.Linear(self.hidden_size, output_dim) for _ in range(self.num_layers)
+        ]) if enable_layer_outputs else None
+        self.layer_norms = nn.ModuleList([
+            nn.LayerNorm(self.hidden_size) for _ in range(self.num_layers)
+        ]) if enable_layer_outputs else None
+
+
+    def _pool(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Pool hidden states to get audio embedding."""
+        if self.pooling_mode == "mean":
+            return hidden_states.mean(dim=1)
+        elif self.pooling_mode == "first":
+            return hidden_states[:, 0]
+        elif self.pooling_mode == "last":
+            return hidden_states[:, -1]
+        else:
+            raise ValueError(f"Unknown pooling mode: {self.pooling_mode}")
+
+    def preprocess(
+        self,
+        audio: np.ndarray | list[np.ndarray] | torch.Tensor,
+        sampling_rate: int = 16000,
+    ) -> torch.Tensor:
+        """
+        Preprocess audio for the encoder.
+
+        Args:
+            audio: Audio waveform(s) as numpy arrays or tensors
+            sampling_rate: Audio sampling rate (should be 16kHz for Whisper)
+
+        Returns:
+            input_features: Mel spectrogram features
+        """
+        if isinstance(audio, torch.Tensor):
+            audio = audio.cpu().numpy()
+
+        if isinstance(audio, np.ndarray) and audio.ndim == 1:
+            audio = [audio]
+
+        # Use feature extractor to get mel spectrograms
+        features = self.feature_extractor(
+            audio,
+            sampling_rate=sampling_rate,
+            return_tensors="pt",
+        )
+        return features["input_features"]
+
+    def _prepare_input_features(
+        self,
+        input_features: torch.Tensor | None,
+        audio: np.ndarray | list[np.ndarray] | None,
+        sampling_rate: int,
+    ) -> torch.Tensor | None:
+        """Preprocess raw audio and move it to the encoder device."""
+        if input_features is not None or audio is None:
+            return input_features
+        prepared = self.preprocess(audio, sampling_rate)
+        parameter = next(self.parameters(), None)
+        device = (
+            parameter.device
+            if parameter is not None
+            else torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        )
+        return prepared.to(device)
+
+    @staticmethod
+    def _hidden_states(outputs: Any) -> tuple[torch.Tensor, ...] | None:
+        """Return encoder-layer states without the input embedding state."""
+        hidden_states = getattr(outputs, "hidden_states", None)
+        return hidden_states[1:] if hidden_states is not None else None
+
+    def _project_layer(
+        self,
+        hidden_states: torch.Tensor,
+        layer_idx: int,
+    ) -> torch.Tensor:
+        """Pool and project one intermediate encoder layer."""
+        if self.layer_norms is not None and self.layer_projections is not None:
+            normalized = self.layer_norms[layer_idx](hidden_states)
+            return self.layer_projections[layer_idx](self._pool(normalized))
+        return self.projection(self._pool(hidden_states))
+
+    def _select_embedding(
+        self,
+        outputs: Any,
+        all_hidden_states: tuple[torch.Tensor, ...] | None,
+        target_layer: int | None,
+    ) -> torch.Tensor:
+        """Select either an intermediate 2DMSE layer or the final layer."""
+        if target_layer is None or all_hidden_states is None:
+            return self.projection(self._pool(outputs.last_hidden_state))
+        layer_idx = min(target_layer, len(all_hidden_states) - 1)
+        return self._project_layer(all_hidden_states[layer_idx], layer_idx)
+
+    def _all_layer_embeddings(
+        self,
+        all_hidden_states: tuple[torch.Tensor, ...],
+    ) -> list[torch.Tensor]:
+        """Project every encoder layer for confidence-based exit training."""
+        embeddings = []
+        for layer_idx, layer_hidden_states in enumerate(all_hidden_states):
+            embedding = self._project_layer(layer_hidden_states, layer_idx)
+            if self.normalize:
+                embedding = functional.normalize(embedding, p=2, dim=-1)
+            embeddings.append(embedding)
+        return embeddings
+
+    def forward(
+        self,
+        input_features: torch.Tensor | None = None,
+        audio: np.ndarray | list[np.ndarray] | None = None,
+        sampling_rate: int = 16000,
+        target_layer: int | None = None,  # For 2DMSE
+        target_dim: int | None = None,  # For MRL
+        return_all_layers: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
+        """
+        Forward pass for audio encoding.
+
+        Args:
+            input_features: Preprocessed mel spectrogram features
+            audio: Raw audio waveforms (will be preprocessed if input_features not provided)
+            sampling_rate: Audio sampling rate
+            target_layer: Exit at this layer (None = full model)
+            target_dim: Truncate to this dimension (None = full dim)
+            return_all_layers: Return embeddings from all layers
+
+        Returns:
+            embeddings: [batch_size, output_dim] or truncated
+            all_layer_embeddings: (optional) List of [batch_size, output_dim]
+        """
+        input_features = self._prepare_input_features(
+            input_features,
+            audio,
+            sampling_rate,
+        )
+
+        # Forward through encoder
+        outputs = self.encoder(
+            input_features=input_features,
+            output_hidden_states=self.enable_layer_outputs or return_all_layers,
+        )
+
+        all_hidden_states = self._hidden_states(outputs)
+        embedding = self._select_embedding(outputs, all_hidden_states, target_layer)
+
+        # Dimension truncation (MRL)
+        if target_dim is not None:
+            embedding = embedding[:, :target_dim]
+
+        # Normalize
+        if self.normalize:
+            embedding = functional.normalize(embedding, p=2, dim=-1)
+
+        # Return all layer embeddings if requested
+        if return_all_layers and all_hidden_states is not None:
+            return embedding, self._all_layer_embeddings(all_hidden_states)
+
+        return embedding
+
+    def encode(
+        self,
+        audio_list: list[np.ndarray],
+        batch_size: int = 16,
+        sampling_rate: int = 16000,
+        show_progress: bool = True,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Encode a list of audio waveforms into embeddings."""
+        all_embeddings = []
+
+        for i in range(0, len(audio_list), batch_size):
+            batch_audio = audio_list[i : i + batch_size]
+            with torch.no_grad():
+                embeddings = self.forward(
+                    audio=batch_audio,
+                    sampling_rate=sampling_rate,
+                    **kwargs,
+                )
+            all_embeddings.append(embeddings.cpu())
+
+        return torch.cat(all_embeddings, dim=0)

omni_components/audio_io.py

@@ -0,0 +1,11 @@
+"""Optional file decoding; waveform inference does not require a codec package."""
+class _AudioFiles:
+    @staticmethod
+    def load(*args, **kwargs):
+        try:
+            from librosa import load
+        except ImportError as exc:
+            raise ImportError("Install librosa for direct audio-file loading, or pass mono waveforms to encode_audio") from exc
+        return load(*args, **kwargs)
+
+librosa = _AudioFiles()

omni_components/embedder.py

@@ -0,0 +1,394 @@
+"""Main multimodal embedder combining all modality encoders."""
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from PIL import Image
+from torch import nn
+from torch.nn import functional
+
+from .audio_encoder import AudioEncoder
+from .fusion import ModalityFusion
+from .image_encoder import ImageEncoder
+from .text_encoder import TextEncoder
+
+logger = logging.getLogger(__name__)
+
+
+class MultimodalEmbedder(nn.Module):
+    """
+    Multimodal embedding model combining text, image, and audio encoders.
+
+    Supports:
+    - Single modality encoding
+    - Multimodal fusion
+    - Adaptive layer exit (2DMSE)
+    - Dimension truncation (MRL)
+    """
+
+    def __init__(self, component_dir, legacy=False):
+        super().__init__()
+        root = Path(component_dir)
+        self.output_dim = 384
+        self.normalize = True
+        self.enable_layer_outputs = True
+        self.fusion_type = "transformer"
+        self.num_fusion_layers = 2
+        self.max_text_length = 128
+        self.text_encoder_revision = None
+        self.image_encoder_revision = None
+        self.audio_encoder_revision = None
+        self.text_encoder = TextEncoder(root / "text", output_dim=384, normalize=False,
+                                       max_length=128, enable_layer_outputs=True)
+        self.image_encoder = ImageEncoder(root / "image", output_dim=384, normalize=False,
+                                          enable_layer_outputs=True, legacy=legacy)
+        self.audio_encoder = AudioEncoder(root / "audio", output_dim=384, normalize=False,
+                                          enable_layer_outputs=True)
+        self.fusion = ModalityFusion(input_dim=384, hidden_dim=384, output_dim=384,
+                        fusion_type="transformer", num_fusion_layers=2, enable_layer_outputs=True)
+        self.encoder_layers = {"text": self.text_encoder.num_layers,
+                              "image": self.image_encoder.num_layers,
+                              "audio": self.audio_encoder.num_layers, "fusion": 2}
+
+
+    def encode_text(
+        self,
+        texts: str | list[str],
+        target_layer: int | None = None,
+        target_dim: int | None = None,
+    ) -> torch.Tensor:
+        """Encode text into embeddings."""
+        if isinstance(texts, str):
+            texts = [texts]
+
+        embedding = self.text_encoder(
+            texts=texts,
+            target_layer=target_layer,
+            target_dim=target_dim,
+        )
+
+        if self.normalize:
+            embedding = functional.normalize(embedding, p=2, dim=-1)
+
+        return embedding
+
+    def encode_image(
+        self,
+        images: Image.Image | list[Image.Image] | torch.Tensor,
+        target_layer: int | None = None,
+        target_dim: int | None = None,
+    ) -> torch.Tensor:
+        """Encode images into embeddings."""
+        if isinstance(images, Image.Image):
+            images = [images]
+
+        embedding = self.image_encoder(
+            images=images,
+            target_layer=target_layer,
+            target_dim=target_dim,
+        )
+
+        if self.normalize:
+            embedding = functional.normalize(embedding, p=2, dim=-1)
+
+        return embedding
+
+    def encode_audio(
+        self,
+        audio: np.ndarray | list[np.ndarray] | torch.Tensor,
+        sampling_rate: int = 16000,
+        target_layer: int | None = None,
+        target_dim: int | None = None,
+    ) -> torch.Tensor:
+        """Encode audio into embeddings."""
+        if isinstance(audio, np.ndarray) and audio.ndim == 1:
+            audio = [audio]
+
+        embedding = self.audio_encoder(
+            audio=audio,
+            sampling_rate=sampling_rate,
+            target_layer=target_layer,
+            target_dim=target_dim,
+        )
+
+        if self.normalize:
+            embedding = functional.normalize(embedding, p=2, dim=-1)
+
+        return embedding
+
+    def encode_multimodal(
+        self,
+        texts: str | list[str] | None = None,
+        images: Image.Image | list[Image.Image] | torch.Tensor | None = None,
+        audio: np.ndarray | list[np.ndarray] | torch.Tensor | None = None,
+        sampling_rate: int = 16000,
+        target_layer: int | None = None,
+        target_dim: int | None = None,
+        encoder_target_layers: dict[str, int] | None = None,
+    ) -> torch.Tensor:
+        """
+        Encode multimodal inputs into a fused embedding.
+
+        Args:
+            texts: Text input(s)
+            images: Image input(s)
+            audio: Audio input(s)
+            sampling_rate: Audio sampling rate
+            target_layer: Fusion layer to exit at (2DMSE)
+            target_dim: Output dimension to truncate to (MRL)
+            encoder_target_layers: Dict of encoder-specific layer exits
+        """
+        embeddings = {}
+        enc_layers = encoder_target_layers or {}
+
+        if texts is not None:
+            if isinstance(texts, str):
+                texts = [texts]
+            text_layer = enc_layers.get("text", None)
+            embeddings["text"] = self.text_encoder(
+                texts=texts,
+                target_layer=text_layer,
+            )
+
+        if images is not None:
+            if isinstance(images, Image.Image):
+                images = [images]
+            image_layer = enc_layers.get("image", None)
+            embeddings["image"] = self.image_encoder(
+                images=images,
+                target_layer=image_layer,
+            )
+
+        if audio is not None:
+            if isinstance(audio, np.ndarray) and audio.ndim == 1:
+                audio = [audio]
+            audio_layer = enc_layers.get("audio", None)
+            embeddings["audio"] = self.audio_encoder(
+                audio=audio,
+                sampling_rate=sampling_rate,
+                target_layer=audio_layer,
+            )
+
+        if not embeddings:
+            raise ValueError("At least one modality must be provided")
+
+        fused = self.fusion(
+            embeddings=embeddings,
+            target_layer=target_layer,
+            target_dim=target_dim,
+            normalize=self.normalize,
+        )
+        return fused
+
+    def forward(
+        self,
+        texts: str | list[str] | None = None,
+        images: Image.Image | list[Image.Image] | torch.Tensor | None = None,
+        audio: np.ndarray | list[np.ndarray] | torch.Tensor | None = None,
+        modality: str | None = None,
+        target_layer: int | None = None,
+        target_dim: int | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Forward pass - encode inputs into embeddings."""
+        provided = sum([texts is not None, images is not None, audio is not None])
+
+        if provided == 0:
+            raise ValueError("At least one input must be provided")
+
+        if provided == 1 or modality is not None:
+            if texts is not None or modality == "text":
+                return self.encode_text(texts, target_layer, target_dim)
+            elif images is not None or modality == "image":
+                return self.encode_image(images, target_layer, target_dim)
+            elif audio is not None or modality == "audio":
+                return self.encode_audio(
+                    audio, target_layer=target_layer, target_dim=target_dim, **kwargs
+                )
+
+        return self.encode_multimodal(
+            texts=texts,
+            images=images,
+            audio=audio,
+            target_layer=target_layer,
+            target_dim=target_dim,
+            **kwargs,
+        )
+
+    def get_layer_info(self) -> dict[str, int]:
+        """Get number of layers for each encoder."""
+        return self.encoder_layers.copy()
+
+    def _adaptive_embedding(
+        self,
+        encoder: nn.Module,
+        encoder_kwargs: dict[str, Any],
+        exit_layers: list[int],
+        confidence_threshold: float,
+    ) -> tuple[torch.Tensor, dict[str, int | float]]:
+        """Select the first candidate layer that satisfies the confidence gate."""
+        final_embedding, layer_embeddings = encoder(
+            **encoder_kwargs,
+            return_all_layers=True,
+        )
+        selected_layer = encoder.num_layers
+        selected_embedding = final_embedding
+        selected_confidence = 1.0
+        for layer in exit_layers:
+            if layer > len(layer_embeddings):
+                continue
+            candidate = layer_embeddings[layer - 1]
+            confidence = self._estimate_confidence(candidate)
+            if confidence >= confidence_threshold:
+                selected_layer = layer
+                selected_embedding = candidate
+                selected_confidence = confidence
+                break
+        return selected_embedding, {
+            "layer": selected_layer,
+            "confidence": selected_confidence,
+        }
+
+    def encode_adaptive(
+        self,
+        texts: str | list[str] | None = None,
+        images: Image.Image | list[Image.Image] | torch.Tensor | None = None,
+        audio: np.ndarray | list[np.ndarray] | torch.Tensor | None = None,
+        confidence_threshold: float = 0.8,
+        return_exit_info: bool = False,
+        **kwargs,
+    ):
+        """
+        Encode with confidence-based adaptive layer exit.
+
+        Each modality independently decides when to exit based on confidence.
+        Uses single forward pass with all layer outputs for efficiency.
+
+        Args:
+            texts: Text input(s)
+            images: Image input(s)
+            audio: Audio input(s)
+            confidence_threshold: Minimum confidence to exit early (0-1)
+            return_exit_info: If True, return (embedding, exit_info_dict)
+
+        Returns:
+            embedding or (embedding, exit_info) with layer/confidence per modality
+        """
+        embeddings = {}
+        exit_info = {}
+
+        text_exit_layers = [6, 11, 16, min(22, self.text_encoder.num_layers)]
+        image_exit_layers = [6, 13, 20, min(27, self.image_encoder.num_layers)]
+
+        if texts is not None:
+            if isinstance(texts, str):
+                texts = [texts]
+
+            embedding, info = self._adaptive_embedding(
+                self.text_encoder,
+                {"texts": texts},
+                text_exit_layers,
+                confidence_threshold,
+            )
+            embeddings["text"] = embedding
+            exit_info["text"] = info
+
+        if images is not None:
+            if isinstance(images, Image.Image):
+                images = [images]
+
+            embedding, info = self._adaptive_embedding(
+                self.image_encoder,
+                {"images": images},
+                image_exit_layers,
+                confidence_threshold,
+            )
+            embeddings["image"] = embedding
+            exit_info["image"] = info
+
+        if audio is not None:
+            emb = self.audio_encoder(audio=audio, **kwargs)
+            embeddings["audio"] = emb
+            exit_info["audio"] = {
+                "layer": self.audio_encoder.num_layers,
+                "confidence": 1.0,
+            }
+
+        # Combine embeddings
+        if len(embeddings) == 0:
+            raise ValueError("At least one modality must be provided")
+        elif len(embeddings) == 1:
+            result = next(iter(embeddings.values()))
+        else:
+            result = self.fusion(embeddings=embeddings, normalize=self.normalize)
+
+        if self.normalize and len(embeddings) == 1:
+            result = functional.normalize(result, p=2, dim=-1)
+
+        if return_exit_info:
+            return result, exit_info
+        return result
+
+    def encode_with_fixed_layers(
+        self,
+        texts: str | list[str] | None = None,
+        images: Image.Image | list[Image.Image] | torch.Tensor | None = None,
+        audio: np.ndarray | list[np.ndarray] | torch.Tensor | None = None,
+        text_layer: int | None = None,
+        image_layer: int | None = None,
+        audio_layer: int | None = None,
+        target_dim: int | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Encode with explicit layer selection per modality.
+
+        This is the most efficient method when you know which layers to use,
+        as it only computes up to the specified layer (true early exit).
+
+        Args:
+            text_layer: Exit layer for text encoder (None = full)
+            image_layer: Exit layer for image encoder (None = full)
+            audio_layer: Exit layer for audio encoder (None = full)
+            target_dim: Output dimension truncation (MRL)
+        """
+        return self.encode_multimodal(
+            texts=texts,
+            images=images,
+            audio=audio,
+            encoder_target_layers={
+                "text": text_layer,
+                "image": image_layer,
+                "audio": audio_layer,
+            },
+            target_dim=target_dim,
+            **kwargs,
+        )
+
+    def _estimate_confidence(self, embedding: torch.Tensor) -> float:
+        """
+        Estimate confidence based on embedding properties.
+
+        For production, replace with trained ConfidenceEstimator
+        that learns to predict when early layers are sufficient.
+        """
+        with torch.no_grad():
+            # Variance-based: lower variance = more confident/stable
+            variance = embedding.var(dim=-1).mean().item()
+            conf_from_var = 1.0 / (1.0 + variance)
+
+            # Norm-based: well-normalized = good embedding
+            norm = embedding.norm(dim=-1).mean().item()
+            conf_from_norm = min(1.0, norm)
+
+            return (conf_from_var + conf_from_norm) / 2
+
+
+
+
+# Legacy alias for backward compatibility
+DiffusionMultimodalEmbedder = MultimodalEmbedder  # Deprecated

omni_components/fusion.py

@@ -0,0 +1,344 @@
+"""Modality fusion module for combining embeddings from different modalities."""
+
+import logging
+
+import torch
+from torch import nn
+from torch.nn import functional
+
+logger = logging.getLogger(__name__)
+
+
+class ModalityProjector(nn.Module):
+    """Projects modality-specific embeddings to a common dimension."""
+
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        num_layers: int = 2,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+
+        layers = []
+        current_dim = input_dim
+
+        for _i in range(num_layers - 1):
+            layers.extend(
+                [
+                    nn.Linear(current_dim, output_dim),
+                    nn.LayerNorm(output_dim),
+                    nn.GELU(),
+                    nn.Dropout(dropout),
+                ]
+            )
+            current_dim = output_dim
+
+        # Final projection without activation
+        layers.append(nn.Linear(current_dim, output_dim))
+
+        self.projector = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.projector(x)
+
+
+class CrossModalAttention(nn.Module):
+    """Cross-modal attention for fusing different modality embeddings."""
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        num_heads: int = 8,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+
+        self.attention = nn.MultiheadAttention(
+            embed_dim=hidden_dim,
+            num_heads=num_heads,
+            dropout=dropout,
+            batch_first=True,
+        )
+        self.norm1 = nn.LayerNorm(hidden_dim)
+        self.norm2 = nn.LayerNorm(hidden_dim)
+
+        self.ffn = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim * 4),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim * 4, hidden_dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+            query: [batch, seq_q, dim]
+            key: [batch, seq_k, dim]
+            value: [batch, seq_v, dim]
+        """
+        # Cross attention
+        attn_output, _ = self.attention(query, key, value)
+        x = self.norm1(query + attn_output)
+
+        # FFN
+        ffn_output = self.ffn(x)
+        x = self.norm2(x + ffn_output)
+
+        return x
+
+
+class FusionTransformerBlock(nn.Module):
+    """Transformer block for multimodal fusion."""
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        num_heads: int = 8,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+
+        self.self_attention = nn.MultiheadAttention(
+            embed_dim=hidden_dim,
+            num_heads=num_heads,
+            dropout=dropout,
+            batch_first=True,
+        )
+        self.norm1 = nn.LayerNorm(hidden_dim)
+        self.norm2 = nn.LayerNorm(hidden_dim)
+
+        self.ffn = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim * 4),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim * 4, hidden_dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Self attention
+        attn_output, _ = self.self_attention(x, x, x)
+        x = self.norm1(x + attn_output)
+
+        # FFN
+        ffn_output = self.ffn(x)
+        x = self.norm2(x + ffn_output)
+
+        return x
+
+
+class ModalityFusion(nn.Module):
+    """
+    Fuses embeddings from multiple modalities into a unified representation.
+
+    Supports:
+    - Simple concatenation + projection
+    - Attention-based fusion
+    - Transformer-based fusion with early exit support
+    """
+
+    def __init__(
+        self,
+        input_dim: int = 1024,
+        hidden_dim: int = 1024,
+        output_dim: int = 1024,
+        num_modalities: int = 3,  # text, image, audio
+        fusion_type: str = "attention",  # simple, attention, transformer
+        num_fusion_layers: int = 4,
+        num_heads: int = 8,
+        dropout: float = 0.1,
+        enable_layer_outputs: bool = True,
+    ):
+        super().__init__()
+
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.output_dim = output_dim
+        self.num_modalities = num_modalities
+        self.fusion_type = fusion_type
+        self.num_fusion_layers = num_fusion_layers
+        self.enable_layer_outputs = enable_layer_outputs
+
+        # Modality type embeddings
+        self.modality_embeddings = nn.Embedding(num_modalities, hidden_dim)
+
+        # Input projections for each modality
+        self.input_projections = nn.ModuleDict(
+            {
+                "text": ModalityProjector(input_dim, hidden_dim),
+                "image": ModalityProjector(input_dim, hidden_dim),
+                "audio": ModalityProjector(input_dim, hidden_dim),
+            }
+        )
+
+        if fusion_type == "simple":
+            # Simple concatenation + MLP
+            self.fusion = nn.Sequential(
+                nn.Linear(hidden_dim * num_modalities, hidden_dim),
+                nn.LayerNorm(hidden_dim),
+                nn.GELU(),
+                nn.Linear(hidden_dim, output_dim),
+            )
+
+        elif fusion_type == "attention":
+            # Cross-modal attention
+            self.cross_attention = CrossModalAttention(
+                hidden_dim=hidden_dim,
+                num_heads=num_heads,
+                dropout=dropout,
+            )
+            self.output_projection = nn.Linear(hidden_dim, output_dim)
+
+        elif fusion_type == "transformer":
+            # Full transformer fusion
+            self.fusion_layers = nn.ModuleList(
+                [
+                    FusionTransformerBlock(
+                        hidden_dim=hidden_dim,
+                        num_heads=num_heads,
+                        dropout=dropout,
+                    )
+                    for _ in range(num_fusion_layers)
+                ]
+            )
+            self.output_projection = nn.Linear(hidden_dim, output_dim)
+
+            # Per-layer projections for 2DMSE
+            if enable_layer_outputs:
+                self.layer_projections = nn.ModuleList(
+                    [
+                        nn.Linear(hidden_dim, output_dim)
+                        for _ in range(num_fusion_layers)
+                    ]
+                )
+                self.layer_norms = nn.ModuleList(
+                    [nn.LayerNorm(hidden_dim) for _ in range(num_fusion_layers)]
+                )
+
+        else:
+            raise ValueError(f"Unknown fusion type: {fusion_type}")
+
+    def _project_modalities(
+        self,
+        embeddings: dict[str, torch.Tensor],
+    ) -> tuple[dict[str, torch.Tensor], int, torch.device]:
+        """Project available modalities and attach learned modality embeddings."""
+        available = {
+            name: value for name, value in embeddings.items() if value is not None
+        }
+        if not available:
+            raise ValueError("At least one modality embedding must be provided")
+
+        first_embedding = next(iter(available.values()))
+        batch_size = first_embedding.shape[0]
+        device = first_embedding.device
+        modality_idx = {"text": 0, "image": 1, "audio": 2}
+        projected = {}
+        for modality, embedding in available.items():
+            if embedding.shape[0] != batch_size:
+                raise ValueError(
+                    "All modality embeddings must have the same batch size"
+                )
+            projection = self.input_projections[modality](embedding)
+            type_embedding = self.modality_embeddings(
+                torch.tensor([modality_idx[modality]], device=device)
+            ).expand(batch_size, -1)
+            projected[modality] = projection + type_embedding
+        return projected, batch_size, device
+
+    def _simple_fusion(
+        self,
+        projected: dict[str, torch.Tensor],
+        batch_size: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        """Concatenate projected modalities, padding missing inputs with zeros."""
+        modality_list = [
+            projected.get(
+                modality,
+                torch.zeros(batch_size, self.hidden_dim, device=device),
+            )
+            for modality in ("text", "image", "audio")
+        ]
+        return self.fusion(torch.cat(modality_list, dim=-1))
+
+    def _attention_fusion(
+        self,
+        projected: dict[str, torch.Tensor],
+    ) -> torch.Tensor:
+        """Fuse available modalities with cross-modal attention."""
+        available = list(projected.values())
+        if len(available) == 1:
+            return self.output_projection(available[0])
+        stacked = torch.stack(available, dim=1)
+        query = stacked.mean(dim=1, keepdim=True)
+        fused = self.cross_attention(query, stacked, stacked)
+        return self.output_projection(fused.squeeze(1))
+
+    def _transformer_fusion(
+        self,
+        projected: dict[str, torch.Tensor],
+        target_layer: int | None,
+    ) -> torch.Tensor:
+        """Fuse available modalities through the transformer stack."""
+        available = list(projected.values())
+        x = (
+            available[0].unsqueeze(1)
+            if len(available) == 1
+            else torch.stack(available, dim=1)
+        )
+        for layer_idx, layer in enumerate(self.fusion_layers):
+            x = layer(x)
+            if target_layer is None or layer_idx != target_layer:
+                continue
+            pooled = x.mean(dim=1)
+            if self.enable_layer_outputs:
+                normalized = self.layer_norms[layer_idx](x)
+                return self.layer_projections[layer_idx](normalized.mean(dim=1))
+            return self.output_projection(pooled)
+        return self.output_projection(x.mean(dim=1))
+
+    def forward(
+        self,
+        embeddings: dict[str, torch.Tensor],
+        target_layer: int | None = None,  # For 2DMSE
+        target_dim: int | None = None,  # For MRL
+        normalize: bool = True,
+    ) -> torch.Tensor:
+        """
+        Fuse modality embeddings.
+
+        Args:
+            embeddings: Dict with keys 'text', 'image', 'audio' and embedding tensors
+            target_layer: Exit at this fusion layer (transformer mode only)
+            target_dim: Truncate output to this dimension
+            normalize: Whether to L2 normalize output
+
+        Returns:
+            fused_embedding: [batch_size, output_dim]
+        """
+        projected, batch_size, device = self._project_modalities(embeddings)
+        if self.fusion_type == "simple":
+            output = self._simple_fusion(projected, batch_size, device)
+        elif self.fusion_type == "attention":
+            output = self._attention_fusion(projected)
+        else:
+            output = self._transformer_fusion(projected, target_layer)
+
+        # Dimension truncation (MRL)
+        if target_dim is not None:
+            output = output[:, :target_dim]
+
+        # Normalize
+        if normalize:
+            output = functional.normalize(output, p=2, dim=-1)
+
+        return output

omni_components/image_encoder.py

@@ -0,0 +1,215 @@
+"""Image encoder based on SigLIP 2 / ViT architecture."""
+
+import logging
+from typing import Any
+
+import torch
+from PIL import Image
+from torch import nn
+from torch.nn import functional
+from transformers import AutoConfig, AutoModel, AutoProcessor
+
+logger = logging.getLogger(__name__)
+
+
+class ImageEncoder(nn.Module):
+    """
+    Image encoder using SigLIP 2 / ViT architecture.
+
+    Default: SigLIP 2 SO400M (384px, 400M params)
+    - SOTA vision encoder for multimodal learning
+    - Multilingual support
+    - Apache 2.0 license
+
+    Supports adaptive layer exit (2DMSE) and dimension truncation (MRL).
+    """
+
+    def __init__(self, model_name_or_path, revision=None, output_dim=384,
+                 pooling_mode="mean", normalize=True, enable_layer_outputs=True,
+                 legacy=False):
+        super().__init__()
+        from transformers import AutoImageProcessor, SiglipVisionModel
+        self.model_name = str(model_name_or_path)
+        self.revision = revision
+        self.output_dim = output_dim
+        self.pooling_mode = pooling_mode
+        self.normalize = normalize
+        self.enable_layer_outputs = enable_layer_outputs
+        self.is_siglip = True
+        self.is_siglip2 = False
+        self.config = AutoConfig.from_pretrained(model_name_or_path,
+                            local_files_only=True, trust_remote_code=False)
+        self.processor = AutoImageProcessor.from_pretrained(model_name_or_path,
+                            local_files_only=True, trust_remote_code=False, use_fast=False)
+        if legacy:
+            self.encoder = AutoModel.from_config(self.config, attn_implementation="sdpa")
+            self.vision_encoder = self.encoder.vision_model
+        else:
+            self.config.vision_config._attn_implementation = "sdpa"
+            self.vision_encoder = SiglipVisionModel(self.config.vision_config).vision_model
+        self.hidden_size = self.config.vision_config.hidden_size
+        self.num_layers = self.config.vision_config.num_hidden_layers
+        self.projection = nn.Linear(self.hidden_size, output_dim) if self.hidden_size != output_dim else nn.Identity()
+        self.layer_projections = nn.ModuleList([
+            nn.Linear(self.hidden_size, output_dim) for _ in range(self.num_layers)
+        ]) if enable_layer_outputs else None
+
+
+    def _pool(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Pool hidden states to get image embedding."""
+        if self.pooling_mode == "mean":
+            # Skip CLS token if present, average over patches
+            if hidden_states.shape[1] > 1:
+                return hidden_states[:, 1:].mean(dim=1)
+            return hidden_states.mean(dim=1)
+        elif self.pooling_mode == "cls":
+            return hidden_states[:, 0]
+        else:
+            raise ValueError(f"Unknown pooling mode: {self.pooling_mode}")
+
+    def preprocess(
+        self,
+        images: Image.Image | list[Image.Image] | torch.Tensor,
+    ) -> torch.Tensor:
+        """Preprocess images for the encoder."""
+        if isinstance(images, torch.Tensor):
+            return images
+
+        if isinstance(images, Image.Image):
+            images = [images]
+
+        processed = self.processor(images=images, return_tensors="pt")
+        return processed["pixel_values"]
+
+    def _prepare_pixel_values(
+        self,
+        pixel_values: torch.Tensor | None,
+        images: Image.Image | list[Image.Image] | None,
+    ) -> torch.Tensor | None:
+        """Preprocess raw images and move them to the encoder device."""
+        if pixel_values is not None or images is None:
+            return pixel_values
+        prepared = self.preprocess(images)
+        parameter = next(self.parameters(), None)
+        device = (
+            parameter.device
+            if parameter is not None
+            else torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        )
+        return prepared.to(device)
+
+    @staticmethod
+    def _hidden_states(outputs: Any) -> tuple[torch.Tensor, ...] | None:
+        """Return encoder-layer states without the input embedding state."""
+        hidden_states = getattr(outputs, "hidden_states", None)
+        return hidden_states[1:] if hidden_states is not None else None
+
+    def _project_layer(
+        self,
+        hidden_states: torch.Tensor,
+        layer_idx: int,
+    ) -> torch.Tensor:
+        """Pool and project one intermediate vision layer."""
+        pooled = self._pool(hidden_states)
+        if self.layer_projections is not None:
+            return self.layer_projections[layer_idx](pooled)
+        return self.projection(pooled)
+
+    def _select_embedding(
+        self,
+        outputs: Any,
+        all_hidden_states: tuple[torch.Tensor, ...] | None,
+        target_layer: int | None,
+    ) -> torch.Tensor:
+        """Select either an intermediate 2DMSE layer or the final output."""
+        if target_layer is not None and all_hidden_states is not None:
+            layer_idx = min(target_layer, len(all_hidden_states) - 1)
+            return self._project_layer(all_hidden_states[layer_idx], layer_idx)
+        pooled = getattr(outputs, "pooler_output", None)
+        if pooled is None:
+            pooled = self._pool(outputs.last_hidden_state)
+        return self.projection(pooled)
+
+    def _all_layer_embeddings(
+        self,
+        all_hidden_states: tuple[torch.Tensor, ...],
+    ) -> list[torch.Tensor]:
+        """Project every vision layer for confidence-based exit training."""
+        embeddings = []
+        for layer_idx, hidden_states in enumerate(all_hidden_states):
+            embedding = self._project_layer(hidden_states, layer_idx)
+            if self.normalize:
+                embedding = functional.normalize(embedding, p=2, dim=-1)
+            embeddings.append(embedding)
+        return embeddings
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor | None = None,
+        images: Image.Image | list[Image.Image] | None = None,
+        target_layer: int | None = None,
+        target_dim: int | None = None,
+        return_all_layers: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
+        """
+        Forward pass for image encoding.
+
+        Args:
+            pixel_values: Preprocessed image tensors
+            images: Raw PIL images
+            target_layer: Exit at this layer (2DMSE)
+            target_dim: Truncate to this dimension (MRL)
+            return_all_layers: Return embeddings from all layers
+        """
+        pixel_values = self._prepare_pixel_values(pixel_values, images)
+
+        # Forward through encoder
+        outputs = self.vision_encoder(
+            pixel_values=pixel_values,
+            output_hidden_states=self.enable_layer_outputs or return_all_layers,
+        )
+
+        all_hidden_states = self._hidden_states(outputs)
+        embedding = self._select_embedding(outputs, all_hidden_states, target_layer)
+
+        # Dimension truncation (MRL)
+        if target_dim is not None:
+            embedding = embedding[:, :target_dim]
+
+        # Normalize
+        if self.normalize:
+            embedding = functional.normalize(embedding, p=2, dim=-1)
+
+        # Return all layer embeddings if requested
+        if return_all_layers and all_hidden_states is not None:
+            return embedding, self._all_layer_embeddings(all_hidden_states)
+
+        return embedding
+
+    def encode(
+        self,
+        images: list[Image.Image],
+        batch_size: int = 32,
+        target_layer: int | None = None,
+        target_dim: int | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Encode images into embeddings."""
+        all_embeddings = []
+
+        for i in range(0, len(images), batch_size):
+            batch_images = images[i : i + batch_size]
+            with torch.no_grad():
+                embeddings = self.forward(
+                    images=batch_images,
+                    target_layer=target_layer,
+                    target_dim=target_dim,
+                    **kwargs,
+                )
+            all_embeddings.append(embeddings.cpu())
+
+        return torch.cat(all_embeddings, dim=0)
+
+
+# Legacy alias for backward compatibility
+DiffusionImageEncoder = ImageEncoder  # Deprecated

omni_components/mini.py

@@ -0,0 +1,232 @@
+from collections import defaultdict
+from collections.abc import Callable
+from typing import Any
+
+from .audio_io import librosa
+import torch
+from PIL import Image
+from torch import nn
+from torch.nn import functional
+from torch.nn.utils.rnn import pad_sequence
+from transformers import AutoModel, AutoProcessor, WhisperFeatureExtractor, WhisperModel
+
+from .records import PairItem
+
+
+class MultiModalSentenceEmbedder(nn.Module):
+    def __init__(self, component_dir, legacy=False):
+        super().__init__()
+        from pathlib import Path
+        from transformers import AutoConfig, AutoImageProcessor, SiglipVisionModel, Siglip2VisionModel, WhisperConfig
+        from transformers.models.whisper.modeling_whisper import WhisperEncoder
+        from .text_backbone import TextBackbone
+        root = Path(component_dir)
+        self.text_model = TextBackbone(root / "text", max_length=32768)
+        image_config = AutoConfig.from_pretrained(root / "image", local_files_only=True,
+                                                  trust_remote_code=False)
+        if legacy:
+            self.image_model = AutoModel.from_config(image_config, attn_implementation="sdpa")
+        else:
+            image_config.vision_config._attn_implementation = "sdpa"
+            vision_class = {"siglip": SiglipVisionModel, "siglip2": Siglip2VisionModel}[image_config.model_type]
+            self.image_model = vision_class(image_config.vision_config)
+        self.image_processor = AutoImageProcessor.from_pretrained(root / "image",
+                        local_files_only=True, trust_remote_code=False, use_fast=False)
+        audio_config = WhisperConfig.from_pretrained(root / "audio", local_files_only=True)
+        audio_config._attn_implementation = "sdpa"
+        self.audio_model = WhisperEncoder(audio_config)
+        self.audio_processor = WhisperFeatureExtractor.from_pretrained(root / "audio",
+                                                                        local_files_only=True)
+        self.text_proj = nn.Identity()
+        self.image_proj = nn.Linear(image_config.vision_config.hidden_size, 768)
+        self.audio_proj = nn.Linear(audio_config.d_model, 768)
+
+
+    @staticmethod
+    def _get_vision_dim(model: nn.Module) -> int:
+        if hasattr(model, "vision_model") and hasattr(model.config, "vision_config"):
+            return int(model.config.vision_config.hidden_size)
+        if hasattr(model.config, "hidden_size"):
+            return int(model.config.hidden_size)
+        raise ValueError("Could not infer image hidden size")
+
+    def _encode_text(self, texts: list[Any]) -> torch.Tensor:
+        device = next(self.parameters()).device
+        normalized: list[torch.Tensor | None] = [None] * len(texts)
+
+        dict_positions = [
+            idx for idx, item in enumerate(texts) if isinstance(item, dict)
+        ]
+        if dict_positions:
+            pad_values = {
+                "input_ids": 0,
+                "attention_mask": 0,
+                "token_type_ids": 0,
+            }
+            dict_items = [texts[idx] for idx in dict_positions]
+            features = {
+                key: pad_sequence(
+                    [item[key].detach().cpu() for item in dict_items],
+                    batch_first=True,
+                    padding_value=pad_values.get(key, 0),
+                ).to(device)
+                for key in dict_items[0]
+            }
+            out = self.text_model(features)
+            emb = functional.normalize(
+                self.text_proj(out["sentence_embedding"]), p=2, dim=-1
+            )
+            for loc, row in zip(dict_positions, emb, strict=False):
+                normalized[loc] = row
+
+        raw_positions = [
+            idx for idx, item in enumerate(texts) if not isinstance(item, dict)
+        ]
+        if raw_positions:
+            raw_texts = [texts[idx] for idx in raw_positions]
+            features = self.text_model.tokenize(raw_texts)
+            features = {
+                k: (v.to(device) if hasattr(v, "to") else v)
+                for k, v in features.items()
+            }
+            out = self.text_model(features)
+            emb = functional.normalize(
+                self.text_proj(out["sentence_embedding"]), p=2, dim=-1
+            )
+            for loc, row in zip(raw_positions, emb, strict=False):
+                normalized[loc] = row
+
+        return torch.stack([row for row in normalized if row is not None], dim=0)
+
+    def _encode_image_paths(self, paths: list[str]) -> torch.Tensor:
+        images = [Image.open(path).convert("RGB") for path in paths]
+        proc = self.image_processor(images=images, return_tensors="pt")
+        device = next(self.parameters()).device
+        proc = {k: v.to(device) for k, v in proc.items()}
+        return self._encode_image_pixel_values(proc["pixel_values"])
+
+    def _encode_image_pixel_values(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        device = next(self.parameters()).device
+        proc = {"pixel_values": pixel_values.to(device)}
+        if hasattr(self.image_model, "vision_model"):
+            out = self.image_model.vision_model(**proc, output_hidden_states=False)
+            hidden = out.last_hidden_state
+        else:
+            out = self.image_model(**proc, output_hidden_states=False)
+            hidden = out.last_hidden_state
+        pooled = (
+            hidden[:, 1:].mean(dim=1) if hidden.shape[1] > 1 else hidden.mean(dim=1)
+        )
+        emb = self.image_proj(pooled)
+        return functional.normalize(emb, p=2, dim=-1)
+
+    def _encode_audio_paths(self, paths: list[str]) -> torch.Tensor:
+        waves = [librosa.load(path, sr=16000, mono=True)[0] for path in paths]
+        proc = self.audio_processor(waves, sampling_rate=16000, return_tensors="pt")
+        return self._encode_audio_features(proc["input_features"])
+
+    def _encode_audio_features(self, input_features: torch.Tensor) -> torch.Tensor:
+        device = next(self.parameters()).device
+        input_features = input_features.to(device)
+        input_features = input_features.to(self.audio_model.conv1.weight.dtype)
+        out = self.audio_model(
+            input_features=input_features, output_hidden_states=False
+        )
+        pooled = out.last_hidden_state.mean(dim=1)
+        emb = self.audio_proj(pooled)
+        return functional.normalize(emb, p=2, dim=-1)
+
+    @staticmethod
+    def _stack_tensor_values(values: list[Any]) -> torch.Tensor:
+        tensors = []
+        for value in values:
+            if not torch.is_tensor(value):
+                raise TypeError("Expected tensor payload in cached item")
+            tensor = value.detach().cpu()
+            if tensor.dim() > 0 and tensor.shape[0] == 1:
+                tensor = tensor.squeeze(0)
+            tensors.append(tensor)
+        return torch.stack(tensors, dim=0)
+
+    @staticmethod
+    def _assign_embeddings(
+        output: list[torch.Tensor | None],
+        positions: tuple[int, ...],
+        embeddings: torch.Tensor,
+    ) -> None:
+        for position, embedding in zip(positions, embeddings, strict=True):
+            output[position] = embedding
+
+    @staticmethod
+    def _partition_payloads(
+        pairs: list[tuple[int, Any]],
+    ) -> tuple[list[tuple[int, Any]], list[tuple[int, Any]]]:
+        tensor_pairs = [(idx, value) for idx, value in pairs if torch.is_tensor(value)]
+        path_pairs = [
+            (idx, value) for idx, value in pairs if not torch.is_tensor(value)
+        ]
+        return tensor_pairs, path_pairs
+
+    def _encode_text_group(
+        self, output: list[torch.Tensor | None], pairs: list[tuple[int, Any]]
+    ) -> None:
+        positions, values = zip(*pairs, strict=True)
+        self._assign_embeddings(output, positions, self._encode_text(list(values)))
+
+    def _encode_media_group(
+        self,
+        output: list[torch.Tensor | None],
+        pairs: list[tuple[int, Any]],
+        path_encoder: Callable[[list[Any]], torch.Tensor],
+        tensor_encoder: Callable[[torch.Tensor], torch.Tensor],
+    ) -> None:
+        tensor_pairs, path_pairs = self._partition_payloads(pairs)
+        if path_pairs:
+            positions, values = zip(*path_pairs, strict=True)
+            self._assign_embeddings(output, positions, path_encoder(list(values)))
+        if tensor_pairs:
+            positions, values = zip(*tensor_pairs, strict=True)
+            tensors = self._stack_tensor_values(list(values))
+            self._assign_embeddings(output, positions, tensor_encoder(tensors))
+
+    def encode_items(self, items: list[PairItem]) -> torch.Tensor:
+        grouped: dict[str, list[tuple[int, Any]]] = defaultdict(list)
+        for idx, item in enumerate(items):
+            grouped[item.modality].append((idx, item.value))
+
+        device = next(self.parameters()).device
+        output: list[torch.Tensor | None] = [None] * len(items)
+
+        if grouped["text"]:
+            self._encode_text_group(output, grouped["text"])
+
+        if grouped["image"]:
+            self._encode_media_group(
+                output,
+                grouped["image"],
+                self._encode_image_paths,
+                self._encode_image_pixel_values,
+            )
+
+        if grouped["audio"]:
+            self._encode_media_group(
+                output,
+                grouped["audio"],
+                self._encode_audio_paths,
+                self._encode_audio_features,
+            )
+
+        if any(embedding is None for embedding in output):
+            raise ValueError("Every item must use a supported modality")
+        stacked = torch.stack(output, dim=0).to(device=device, dtype=torch.float32)
+        return functional.normalize(stacked, p=2, dim=-1)
+
+
+def multiple_negatives_ranking_loss(
+    anchor: torch.Tensor, positive: torch.Tensor, scale: float = 20.0
+) -> torch.Tensor:
+    scores = torch.matmul(anchor, positive.T) * scale
+    labels = torch.arange(scores.shape[0], device=scores.device)
+    loss_a = torch.nn.functional.cross_entropy(scores, labels)
+    loss_b = torch.nn.functional.cross_entropy(scores.T, labels)
+    return (loss_a + loss_b) * 0.5

omni_components/records.py

@@ -0,0 +1,7 @@
+from dataclasses import dataclass
+from typing import Any
+
+@dataclass
+class PairItem:
+    modality: str
+    value: Any

omni_components/text_backbone.py

@@ -0,0 +1,60 @@
+"""Native encoder and mean pooling with the Sentence Transformers state layout."""
+import json
+from pathlib import Path
+
+import torch
+from torch import nn
+from transformers import AutoConfig, AutoModel, AutoTokenizer
+
+class Transformer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.model = AutoModel.from_config(config, attn_implementation="sdpa")
+
+class TextBackbone(nn.Sequential):
+    def __init__(self, directory, max_length):
+        directory = Path(directory)
+        raw = json.loads((directory / "config.json").read_text())
+        config = AutoConfig.from_pretrained(directory, local_files_only=True,
+                                           trust_remote_code=False)
+        if config.model_type != "modernbert":
+            raise ValueError("Mini requires a ModernBERT text encoder")
+        # Saved current foundation uses native 4.57 fields. Refuse ambiguous v5 RoPE.
+        if raw.get("rope_parameters"):
+            if raw.get("rope_scaling"):
+                raise ValueError("Ambiguous saved rotary configuration")
+            for layer, field in (("full_attention", "global_rope_theta"),
+                                 ("sliding_attention", "local_rope_theta")):
+                block = raw["rope_parameters"][layer]
+                if block.get("rope_type", "default") != "default":
+                    raise ValueError("Unsupported nested rotary configuration")
+                theta = float(block["rope_theta"])
+                if not torch.isfinite(torch.tensor(theta)) or theta <= 0:
+                    raise ValueError("Invalid rotary theta")
+                if field in raw and float(raw[field]) != theta:
+                    raise ValueError("Conflicting rotary theta")
+                setattr(config, field, theta)
+        config.reference_compile = False
+        pool = json.loads((directory / "1_Pooling/config.json").read_text())
+        enabled = [key for key, value in pool.items() if key.startswith("pooling_mode_") and value]
+        if enabled != ["pooling_mode_mean_tokens"] or pool.get("include_prompt", True) is not True:
+            raise ValueError("Unsupported saved pooling contract")
+        super().__init__(Transformer(config))
+        self.tokenizer = AutoTokenizer.from_pretrained(directory, local_files_only=True,
+                                                       trust_remote_code=False)
+        self.max_seq_length = max_length
+        sentence_config = json.loads((directory / "sentence_bert_config.json").read_text())
+        self.do_lower_case = sentence_config.get("do_lower_case", False)
+
+    def tokenize(self, texts):
+        texts = [str(text).strip() for text in texts]
+        if self.do_lower_case:
+            texts = [text.lower() for text in texts]
+        return self.tokenizer(texts, padding=True, truncation="longest_first",
+                              max_length=self.max_seq_length, return_tensors="pt")
+
+    def forward(self, features):
+        outputs = self[0].model(**features)
+        tokens = outputs.last_hidden_state
+        mask = features["attention_mask"].unsqueeze(-1).expand(tokens.size()).to(tokens.dtype)
+        return {"sentence_embedding": torch.sum(tokens * mask, 1) / torch.clamp(mask.sum(1), min=1e-9)}

omni_components/text_encoder.py

@@ -0,0 +1,249 @@
+"""Text encoder supporting 2D Matryoshka (layer exit + dimension truncation)."""
+
+import logging
+from typing import Any
+
+import torch
+from torch import nn
+from torch.nn import functional
+from transformers import AutoConfig, AutoModel, AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+
+class TextEncoder(nn.Module):
+    """
+    Text encoder with 2D Matryoshka support.
+
+    Default model: mmbert-embed-32k-2d-matryoshka
+    - 32K context length
+    - 1800+ languages
+    - Built-in layer exit (2DMSE)
+    - Built-in dimension truncation (MRL)
+    """
+
+    def __init__(self, model_name_or_path, revision=None, output_dim=384,
+                 pooling_mode="mean", normalize=True, max_length=128,
+                 enable_layer_outputs=True):
+        super().__init__()
+        self.model_name = str(model_name_or_path)
+        self.revision = revision
+        self.output_dim = output_dim
+        self.pooling_mode = pooling_mode
+        self.normalize = normalize
+        self.max_length = max_length
+        self.enable_layer_outputs = enable_layer_outputs
+        self.is_mmbert = False
+        self.config = AutoConfig.from_pretrained(model_name_or_path, local_files_only=True,
+                                                 trust_remote_code=False)
+        self.encoder = AutoModel.from_config(self.config, attn_implementation="sdpa")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path,
+                          local_files_only=True, trust_remote_code=False)
+        self.hidden_size = self.config.hidden_size
+        self.num_layers = self.config.num_hidden_layers
+        self.layer_projections = None
+        if self.hidden_size != output_dim:
+            self.projection = nn.Linear(self.hidden_size, output_dim)
+            if enable_layer_outputs:
+                self.layer_projections = nn.ModuleList([
+                    nn.Linear(self.hidden_size, output_dim) for _ in range(self.num_layers)])
+        else:
+            self.projection = nn.Identity()
+
+
+    def _pool(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Pool hidden states to get sentence embedding."""
+        if self.pooling_mode == "mean":
+            if attention_mask is not None:
+                mask = attention_mask.unsqueeze(-1).float()
+                hidden_states = hidden_states * mask
+                return hidden_states.sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
+            return hidden_states.mean(dim=1)
+        elif self.pooling_mode == "cls":
+            return hidden_states[:, 0]
+        elif self.pooling_mode == "last":
+            if attention_mask is not None:
+                seq_lens = attention_mask.sum(dim=1) - 1
+                batch_size = hidden_states.shape[0]
+                return hidden_states[
+                    torch.arange(batch_size, device=hidden_states.device), seq_lens
+                ]
+            return hidden_states[:, -1]
+        else:
+            raise ValueError(f"Unknown pooling mode: {self.pooling_mode}")
+
+    def _apply_final_norm(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Apply final layer norm if available (for mmBERT layer exit)."""
+        if hasattr(self.encoder, "final_norm"):
+            return self.encoder.final_norm(hidden_states)
+        elif hasattr(self.encoder, "norm"):
+            return self.encoder.norm(hidden_states)
+        return hidden_states
+
+    def _prepare_inputs(
+        self,
+        input_ids: torch.Tensor | None,
+        attention_mask: torch.Tensor | None,
+        texts: list[str] | None,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        """Tokenize raw text and move tensors to the encoder device."""
+        if input_ids is not None or texts is None:
+            return input_ids, attention_mask
+        encoded = self.tokenizer(
+            texts,
+            padding=True,
+            truncation=True,
+            max_length=self.max_length,
+            return_tensors="pt",
+        )
+        parameter = next(self.parameters(), None)
+        device = (
+            parameter.device
+            if parameter is not None
+            else torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        )
+        return encoded["input_ids"].to(device), encoded["attention_mask"].to(device)
+
+    @staticmethod
+    def _hidden_states(outputs: Any) -> tuple[torch.Tensor, ...] | None:
+        """Return encoder-layer states without the input embedding state."""
+        hidden_states = getattr(outputs, "hidden_states", None)
+        return hidden_states[1:] if hidden_states else None
+
+    def _project_layer(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        layer_idx: int | None = None,
+    ) -> torch.Tensor:
+        """Normalize, pool, and project one encoder layer."""
+        normalized = self._apply_final_norm(hidden_states)
+        pooled = self._pool(normalized, attention_mask)
+        if layer_idx is not None and self.layer_projections is not None:
+            return self.layer_projections[layer_idx](pooled)
+        return self.projection(pooled)
+
+    def _select_embedding(
+        self,
+        outputs: Any,
+        all_hidden_states: tuple[torch.Tensor, ...] | None,
+        attention_mask: torch.Tensor | None,
+        target_layer: int | None,
+    ) -> torch.Tensor:
+        """Select either an intermediate 2DMSE layer or the final layer."""
+        if target_layer is not None and all_hidden_states is not None:
+            layer_idx = min(target_layer, len(all_hidden_states) - 1)
+            return self._project_layer(
+                all_hidden_states[layer_idx],
+                attention_mask,
+                layer_idx,
+            )
+        return self.projection(self._pool(outputs.last_hidden_state, attention_mask))
+
+    def _all_layer_embeddings(
+        self,
+        all_hidden_states: tuple[torch.Tensor, ...],
+        attention_mask: torch.Tensor | None,
+    ) -> list[torch.Tensor]:
+        """Project every encoder layer for confidence-based exit training."""
+        embeddings = []
+        for layer_idx, hidden_states in enumerate(all_hidden_states):
+            embedding = self._project_layer(
+                hidden_states,
+                attention_mask,
+                layer_idx,
+            )
+            if self.normalize:
+                embedding = functional.normalize(embedding, p=2, dim=-1)
+            embeddings.append(embedding)
+        return embeddings
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        texts: list[str] | None = None,
+        target_layer: int | None = None,
+        target_dim: int | None = None,
+        return_all_layers: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
+        """
+        Forward pass for text encoding.
+
+        Args:
+            input_ids: Tokenized input IDs
+            attention_mask: Attention mask
+            texts: Raw text strings (tokenized if input_ids not provided)
+            target_layer: Exit at this layer (2DMSE)
+            target_dim: Truncate to this dimension (MRL)
+            return_all_layers: Return embeddings from all layers
+        """
+        input_ids, attention_mask = self._prepare_inputs(
+            input_ids,
+            attention_mask,
+            texts,
+        )
+
+        # Forward through encoder
+        outputs = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+        )
+
+        all_hidden_states = self._hidden_states(outputs)
+        embedding = self._select_embedding(
+            outputs,
+            all_hidden_states,
+            attention_mask,
+            target_layer,
+        )
+
+        # Dimension truncation (MRL)
+        if target_dim is not None:
+            embedding = embedding[:, :target_dim]
+
+        # Normalize
+        if self.normalize:
+            embedding = functional.normalize(embedding, p=2, dim=-1)
+
+        # Return all layer embeddings if requested
+        if return_all_layers and all_hidden_states is not None:
+            return embedding, self._all_layer_embeddings(
+                all_hidden_states,
+                attention_mask,
+            )
+
+        return embedding
+
+    def encode(
+        self,
+        texts: list[str],
+        batch_size: int = 32,
+        target_layer: int | None = None,
+        target_dim: int | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Encode texts into embeddings."""
+        all_embeddings = []
+
+        for i in range(0, len(texts), batch_size):
+            batch_texts = texts[i : i + batch_size]
+            with torch.no_grad():
+                embeddings = self.forward(
+                    texts=batch_texts,
+                    target_layer=target_layer,
+                    target_dim=target_dim,
+                    **kwargs,
+                )
+            all_embeddings.append(embeddings.cpu())
+
+        return torch.cat(all_embeddings, dim=0)
+
+
+# Legacy alias for backward compatibility
+DiffusionTextEncoder = TextEncoder  # Deprecated

scores.json

@@ -0,0 +1,420 @@
+{
+  "schema_version": 1,
+  "model": "llm-semantic-router/Vela-1.0-Omni-134M-Nano",
+  "revision_semantics": "self denotes the Hugging Face snapshot containing this file; the native artifact fingerprint binds its evaluated model files.",
+  "score_scale": "0-1; multiply by 100 for percentages",
+  "matrix": {
+    "embedding": {
+      "model": "llm-semantic-router/Vela-1.0-Encoder-307M-Embedding",
+      "revision": "972c180aecd2aa3fca97159098ab00d25d53fffb",
+      "artifact_fingerprint_sha256": "2b4747393b1e3e477fb1391a5d6817917362ab68fc5ae2436c769a7357a541ce",
+      "metrics": {
+        "text.banking77.accuracy": {
+          "numerator": 2464,
+          "denominator": 3080,
+          "value": 0.8
+        },
+        "text.massive-en.accuracy": {
+          "numerator": 2248,
+          "denominator": 2972,
+          "value": 0.756393001345895
+        }
+      },
+      "not_applicable": [
+        "image",
+        "audio"
+      ]
+    },
+    "small": {
+      "model": "llm-semantic-router/multi-modal-embed-small",
+      "revision": "fdf8e01b7b0f3a69ac1ac8e2a64dcb1ede177ba4",
+      "artifact_fingerprint_sha256": "684fe0c48a4ddd3944167f6e5dc75e97ada24e5b120bd9a2921f00dc03112eeb",
+      "metrics": {
+        "text.banking77.accuracy": {
+          "numerator": 2169,
+          "denominator": 3080,
+          "value": 0.7042207792207792
+        },
+        "text.massive-en.accuracy": {
+          "numerator": 1960,
+          "denominator": 2972,
+          "value": 0.6594885598923284
+        },
+        "image.image_to_text.recall@1": {
+          "numerator": 336,
+          "denominator": 823,
+          "value": 0.4082624544349939
+        },
+        "image.image_to_text.recall@5": {
+          "numerator": 553,
+          "denominator": 823,
+          "value": 0.6719319562575942
+        },
+        "image.image_to_text.recall@10": {
+          "numerator": 646,
+          "denominator": 823,
+          "value": 0.7849331713244229
+        },
+        "image.text_to_image.recall@1": {
+          "numerator": 1242,
+          "denominator": 4115,
+          "value": 0.3018226002430134
+        },
+        "image.text_to_image.recall@5": {
+          "numerator": 2445,
+          "denominator": 4115,
+          "value": 0.5941676792223572
+        },
+        "image.text_to_image.recall@10": {
+          "numerator": 3057,
+          "denominator": 4115,
+          "value": 0.7428918590522479
+        },
+        "audio.audio_to_text.recall@1": {
+          "numerator": 110,
+          "denominator": 2611,
+          "value": 0.04212945231711988
+        },
+        "audio.audio_to_text.recall@5": {
+          "numerator": 313,
+          "denominator": 2611,
+          "value": 0.11987744159325929
+        },
+        "audio.audio_to_text.recall@10": {
+          "numerator": 497,
+          "denominator": 2611,
+          "value": 0.1903485254691689
+        },
+        "audio.text_to_audio.recall@1": {
+          "numerator": 250,
+          "denominator": 2610,
+          "value": 0.09578544061302682
+        },
+        "audio.text_to_audio.recall@5": {
+          "numerator": 588,
+          "denominator": 2610,
+          "value": 0.22528735632183908
+        },
+        "audio.text_to_audio.recall@10": {
+          "numerator": 801,
+          "denominator": 2610,
+          "value": 0.30689655172413793
+        }
+      }
+    },
+    "large": {
+      "model": "llm-semantic-router/multi-modal-embed-large",
+      "revision": "e21cde3ccc414c56f504b322662f42c603a939ee",
+      "artifact_fingerprint_sha256": "8439b0c46c36d13a6371081cf7e40cf2c1cc7df68a2bf69e41d514452f65d549",
+      "metrics": {
+        "text.banking77.accuracy": {
+          "numerator": 2334,
+          "denominator": 3080,
+          "value": 0.7577922077922078
+        },
+        "text.massive-en.accuracy": {
+          "numerator": 2149,
+          "denominator": 2972,
+          "value": 0.7230820995962315
+        },
+        "image.image_to_text.recall@1": {
+          "numerator": 350,
+          "denominator": 823,
+          "value": 0.425273390036452
+        },
+        "image.image_to_text.recall@5": {
+          "numerator": 619,
+          "denominator": 823,
+          "value": 0.7521263669501823
+        },
+        "image.image_to_text.recall@10": {
+          "numerator": 721,
+          "denominator": 823,
+          "value": 0.8760631834750912
+        },
+        "image.text_to_image.recall@1": {
+          "numerator": 1442,
+          "denominator": 4115,
+          "value": 0.35042527339003643
+        },
+        "image.text_to_image.recall@5": {
+          "numerator": 2884,
+          "denominator": 4115,
+          "value": 0.7008505467800729
+        },
+        "image.text_to_image.recall@10": {
+          "numerator": 3453,
+          "denominator": 4115,
+          "value": 0.8391251518833536
+        },
+        "audio.audio_to_text.recall@1": {
+          "numerator": 1488,
+          "denominator": 2611,
+          "value": 0.5698965913443125
+        },
+        "audio.audio_to_text.recall@5": {
+          "numerator": 2137,
+          "denominator": 2611,
+          "value": 0.8184603600153199
+        },
+        "audio.audio_to_text.recall@10": {
+          "numerator": 2296,
+          "denominator": 2611,
+          "value": 0.8793565683646113
+        },
+        "audio.text_to_audio.recall@1": {
+          "numerator": 2051,
+          "denominator": 2610,
+          "value": 0.785823754789272
+        },
+        "audio.text_to_audio.recall@5": {
+          "numerator": 2454,
+          "denominator": 2610,
+          "value": 0.9402298850574713
+        },
+        "audio.text_to_audio.recall@10": {
+          "numerator": 2532,
+          "denominator": 2610,
+          "value": 0.9701149425287356
+        }
+      }
+    },
+    "candidate": {
+      "model": "llm-semantic-router/Vela-1.0-Omni-134M-Nano",
+      "revision": "self",
+      "artifact_fingerprint_sha256": "32f410c3febdbd02e3f790b907db482e4ae1d493e6a1c4eaab375ebada234f8b",
+      "metrics": {
+        "text.banking77.accuracy": {
+          "numerator": 2169,
+          "denominator": 3080,
+          "value": 0.7042207792207792
+        },
+        "text.massive-en.accuracy": {
+          "numerator": 1960,
+          "denominator": 2972,
+          "value": 0.6594885598923284
+        },
+        "image.image_to_text.recall@1": {
+          "numerator": 364,
+          "denominator": 823,
+          "value": 0.4422843256379101
+        },
+        "image.image_to_text.recall@5": {
+          "numerator": 601,
+          "denominator": 823,
+          "value": 0.7302551640340219
+        },
+        "image.image_to_text.recall@10": {
+          "numerator": 708,
+          "denominator": 823,
+          "value": 0.8602673147023087
+        },
+        "image.text_to_image.recall@1": {
+          "numerator": 1486,
+          "denominator": 4115,
+          "value": 0.3611178614823815
+        },
+        "image.text_to_image.recall@5": {
+          "numerator": 2828,
+          "denominator": 4115,
+          "value": 0.6872417982989064
+        },
+        "image.text_to_image.recall@10": {
+          "numerator": 3375,
+          "denominator": 4115,
+          "value": 0.8201701093560145
+        },
+        "audio.audio_to_text.recall@1": {
+          "numerator": 138,
+          "denominator": 2611,
+          "value": 0.05285331290693221
+        },
+        "audio.audio_to_text.recall@5": {
+          "numerator": 413,
+          "denominator": 2611,
+          "value": 0.1581769436997319
+        },
+        "audio.audio_to_text.recall@10": {
+          "numerator": 624,
+          "denominator": 2611,
+          "value": 0.23898889314438912
+        },
+        "audio.text_to_audio.recall@1": {
+          "numerator": 303,
+          "denominator": 2610,
+          "value": 0.11609195402298851
+        },
+        "audio.text_to_audio.recall@5": {
+          "numerator": 693,
+          "denominator": 2610,
+          "value": 0.2655172413793103
+        },
+        "audio.text_to_audio.recall@10": {
+          "numerator": 925,
+          "denominator": 2610,
+          "value": 0.3544061302681992
+        }
+      }
+    }
+  },
+  "datasets": {
+    "banking77": {
+      "source": "https://github.com/PolyAI-LDN/task-specific-datasets",
+      "revision": "57ec275d8078af65b7731c2a98be812d844a6d6b",
+      "license": "CC-BY-4.0",
+      "split": "test",
+      "queries": 3080,
+      "classes": 77
+    },
+    "massive-en": {
+      "source": "https://huggingface.co/datasets/AmazonScience/massive",
+      "revision": "ff6bd8e4b27c3543e4f8fe2108f32bb95a6f8740",
+      "license": "CC-BY-4.0",
+      "configuration": "en-US",
+      "split": "test",
+      "queries": 2972,
+      "prototype_classes": 60
+    },
+    "image": {
+      "source": "https://cocodataset.org/",
+      "caption_split_source": "https://cs.stanford.edu/people/karpathy/deepimagesent/",
+      "split": "Karpathy test, CC-BY 2.0 image subset",
+      "images": 823,
+      "captions": 4115,
+      "image_license": "CC-BY-2.0",
+      "caption_license": "CC-BY-4.0"
+    },
+    "audio": {
+      "source": "https://www.openslr.org/12",
+      "split": "test-clean",
+      "archive_sha256": "39fde525e59672dc6d1551919b1478f724438a95aa55f874b576be21967e6c23",
+      "license": "CC-BY-4.0",
+      "clips": 2611,
+      "unique_transcripts": 2610,
+      "maximum_seconds": 30
+    }
+  },
+  "protocol": {
+    "embedding_dtype": "float32",
+    "attention": "SDPA",
+    "autocast": false,
+    "tf32": false,
+    "transformers_version": "4.57.6",
+    "torch_version_family": "2.12",
+    "text_max_tokens": 128,
+    "truncation": "right, native tokenizer including special tokens",
+    "batch_sizes": {
+      "text": 64,
+      "image": 8,
+      "audio": 8
+    },
+    "pooling": "native full-dimension model pooling",
+    "similarity": "FP32 cosine of L2-normalized embeddings",
+    "ties": "lexicographic candidate ID",
+    "retrieval": "complete candidate pools; all annotated positives",
+    "classification": "nearest normalized mean class prototype; fixed prototype split",
+    "aggregation": "Each family is the mean of its two primary metrics; utility is the mean of all six primaries. Retrieval primaries are R@1; R@5 and R@10 are reported separately.",
+    "evaluation_scope": "fixed test-split evaluation pools"
+  },
+  "paired_uncertainty": {
+    "comparison": {
+      "candidate": "llm-semantic-router/Vela-1.0-Omni-134M-Nano",
+      "baseline": "llm-semantic-router/multi-modal-embed-small"
+    },
+    "seed": 20260918,
+    "replicates": 2000,
+    "method": "paired stratified group percentile; PCG64; fixed candidate pools",
+    "intervals": {
+      "audio.audio_to_text.recall@1": {
+        "point_delta": 0.010723860589812333,
+        "lower": 0.005401297417205939,
+        "upper": 0.016930329895273797
+      },
+      "audio.audio_to_text.recall@10": {
+        "point_delta": 0.048640367675220224,
+        "lower": 0.03862647576747038,
+        "upper": 0.057686079679882395
+      },
+      "audio.audio_to_text.recall@5": {
+        "point_delta": 0.038299502106472615,
+        "lower": 0.028966882934863172,
+        "upper": 0.04718931610578753
+      },
+      "audio.text_to_audio.recall@1": {
+        "point_delta": 0.020306513409961684,
+        "lower": 0.01207735759631001,
+        "upper": 0.029131333167300467
+      },
+      "audio.text_to_audio.recall@10": {
+        "point_delta": 0.0475095785440613,
+        "lower": 0.036695309471507946,
+        "upper": 0.05838249564916254
+      },
+      "audio.text_to_audio.recall@5": {
+        "point_delta": 0.040229885057471264,
+        "lower": 0.027551627678505008,
+        "upper": 0.05248271039052227
+      },
+      "family.audio": {
+        "point_delta": 0.015515186999887009,
+        "lower": 0.010126625706639377,
+        "upper": 0.021365285246569197
+      },
+      "family.image": {
+        "point_delta": 0.04665856622114216,
+        "lower": 0.031831713244228434,
+        "upper": 0.061849939246658554
+      },
+      "family.text": {
+        "point_delta": 0.0,
+        "lower": 0.0,
+        "upper": 0.0
+      },
+      "image.image_to_text.recall@1": {
+        "point_delta": 0.03402187120291616,
+        "lower": 0.009720534629404616,
+        "upper": 0.05953827460510328
+      },
+      "image.image_to_text.recall@10": {
+        "point_delta": 0.07533414337788578,
+        "lower": 0.05103280680437424,
+        "upper": 0.10085054678007291
+      },
+      "image.image_to_text.recall@5": {
+        "point_delta": 0.0583232077764277,
+        "lower": 0.031591737545565005,
+        "upper": 0.08383961117861483
+      },
+      "image.text_to_image.recall@1": {
+        "point_delta": 0.059295261239368166,
+        "lower": 0.04325637910085055,
+        "upper": 0.07630619684082625
+      },
+      "image.text_to_image.recall@10": {
+        "point_delta": 0.07727825030376671,
+        "lower": 0.061482381530984204,
+        "upper": 0.0928311057108141
+      },
+      "image.text_to_image.recall@5": {
+        "point_delta": 0.09307411907654921,
+        "lower": 0.07606318347509113,
+        "upper": 0.11009113001215064
+      },
+      "text.banking77.accuracy": {
+        "point_delta": 0.0,
+        "lower": 0.0,
+        "upper": 0.0
+      },
+      "text.massive-en.accuracy": {
+        "point_delta": 0.0,
+        "lower": 0.0,
+        "upper": 0.0
+      },
+      "utility": {
+        "point_delta": 0.020724584407009722,
+        "lower": 0.015571272805177982,
+        "upper": 0.026256156724631264
+      }
+    }
+  }
+}

vela_omni.py

@@ -0,0 +1,101 @@
+"""Self-contained Vela Omni text, image, and audio embeddings."""
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from safetensors.torch import load_file
+from torch import nn
+from torch.nn import functional as F
+
+from omni_components.embedder import MultimodalEmbedder
+from omni_components.mini import MultiModalSentenceEmbedder
+
+
+class VelaOmni(nn.Module):
+    def __init__(self, root, variant, legacy=False, audio_projection="identity"):
+        super().__init__()
+        self.variant = variant
+        self.embedding_dim = 384 if variant == "nano" else 768
+        self.max_text_length = 128 if variant == "nano" else 32768
+        if variant == "nano":
+            self.model = MultimodalEmbedder(Path(root) / "components", legacy=legacy)
+            self.tokenizer = self.model.text_encoder.tokenizer
+            if audio_projection == "linear":
+                self.model.audio_encoder.projection = nn.Linear(384, 384)
+            elif audio_projection != "identity":
+                raise ValueError("Unsupported Nano audio projection")
+        elif variant == "mini":
+            if audio_projection != "identity":
+                raise ValueError("Mini uses its native audio projection")
+            self.model = MultiModalSentenceEmbedder(Path(root) / "components", legacy=legacy)
+            self.tokenizer = self.model.text_model.tokenizer
+        else:
+            raise ValueError("Unknown Vela Omni variant")
+
+    @classmethod
+    def from_pretrained(cls, repo_id, revision=None, device="cpu", dtype=torch.float32):
+        root = Path(repo_id)
+        if not root.is_dir():
+            from huggingface_hub import snapshot_download
+            root = Path(snapshot_download(repo_id, revision=revision,
+                    allow_patterns=["config.json", "model.safetensors", "components/*"]))
+        config = json.loads((root / "config.json").read_text())
+        if config.get("format_version") not in (1, 2):
+            raise ValueError("Unsupported Vela Omni artifact format")
+        instance = cls(root, config["variant"], audio_projection=config.get("audio_projection", "identity"))
+        weights = load_file(str(root / "model.safetensors"), device="cpu")
+        result = instance.model.load_state_dict(weights, strict=True, assign=True)
+        if result.missing_keys or result.unexpected_keys:
+            raise ValueError("Incomplete Vela Omni weights")
+        if any(p.is_meta for p in instance.parameters()):
+            raise ValueError("Uninitialized model parameters")
+        if sum(p.numel() for p in instance.model.parameters()) != config["parameter_count"]:
+            raise ValueError("Parameter count does not match the artifact")
+        instance.to(device=device, dtype=dtype).eval()
+        return instance
+
+    @torch.inference_mode()
+    def encode_text(self, texts):
+        if isinstance(texts, str):
+            texts = [texts]
+        if not texts or any(not isinstance(text, str) for text in texts):
+            raise ValueError("Provide at least one text string")
+        lengths = self.tokenizer(texts, truncation=False, padding=False, return_length=True)["length"]
+        if max(lengths) > self.max_text_length:
+            raise ValueError(f"Text exceeds {self.max_text_length} tokens; shorten or explicitly chunk it")
+        if self.variant == "nano":
+            values = self.model.encode_text(texts)
+        else:
+            values = self.model._encode_text(list(texts))
+        return F.normalize(values.float(), dim=-1)
+
+    @torch.inference_mode()
+    def encode_image(self, images):
+        images = list(images)
+        if not images:
+            raise ValueError("Provide at least one image")
+        if self.variant == "nano":
+            values = self.model.encode_image(images)
+        else:
+            pixels = self.model.image_processor(images=images, return_tensors="pt")["pixel_values"]
+            pixels = pixels.to(device=next(self.parameters()).device, dtype=next(self.model.image_model.parameters()).dtype)
+            values = self.model._encode_image_pixel_values(pixels)
+        return F.normalize(values.float(), dim=-1)
+
+    @torch.inference_mode()
+    def encode_audio(self, waveforms, sampling_rate=16000):
+        if sampling_rate != 16000:
+            raise ValueError("Resample audio to 16000 Hz before encoding")
+        waveforms = [np.asarray(w, dtype=np.float32) for w in waveforms]
+        if not waveforms or any(w.ndim != 1 or not w.size or not np.isfinite(w).all() for w in waveforms):
+            raise ValueError("Audio must contain finite nonempty mono waveforms")
+        if any(len(w) > 30 * sampling_rate for w in waveforms):
+            raise ValueError("Audio inputs must be at most 30 seconds")
+        if self.variant == "nano":
+            values = self.model.encode_audio(waveforms, sampling_rate=sampling_rate)
+        else:
+            features = self.model.audio_processor(waveforms, sampling_rate=sampling_rate,
+                                    return_tensors="pt")["input_features"]
+            values = self.model._encode_audio_features(features)
+        return F.normalize(values.float(), dim=-1)