Repository: embeddings-benchmark/mteb
Pull request #2439: add ops_moa_models
Author: login: ahxgw
Created: 2025-03-26T14:11:09Z | Merged: 2025-03-29T15:47:56Z

<!-- If you are submitting a dataset or a model for the model registry please use the corresponding checklists below otherwise feel free to remove them. -->

<!-- add additional description, question etc. related to the new dataset -->


### Code Quality
<!-- Please do not delete this -->
- [x] **Code Formatted**: Format the code using `make lint` to maintain consistent style.

### Documentation
<!-- Please do not delete this -->
- [ ] **Updated Documentation**: Add or update documentation to reflect the changes introduced in this PR.

### Testing
<!-- Please do not delete this -->
- [ ] **New Tests Added**: Write tests to cover new functionality. Validate with `make test-with-coverage`.
- [ ] **Tests Passed**: Run tests locally using `make test` or `make test-with-coverage` to ensure no existing functionality is broken.


### Adding datasets checklist
<!-- see also https://github.com/embeddings-benchmark/mteb/blob/main/docs/adding_a_dataset.md -->

**Reason for dataset addition**: ... <!-- Add reason for adding dataset here. E.g. it covers task/language/domain previously not covered -->

- [ ] I have run the following models on the task (adding the results to the pr). These can be run using the `mteb -m {model_name} -t {task_name}` command.
  - [ ] `sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2`
  - [ ] `intfloat/multilingual-e5-small`
- [ ] I have checked that the performance is neither trivial (both models gain close to perfect scores) nor random (both models gain close to random scores).
- [ ] If the dataset is too big (e.g. >2048 examples), considering using `self.stratified_subsampling() under dataset_transform()`
- [ ] I have filled out the metadata object in the dataset file (find documentation on it [here](https://github.com/embeddings-benchmark/mteb/blob/main/docs/adding_a_dataset.md#2-creating-the-metadata-object)).
- [ ] Run tests locally to make sure nothing is broken using `make test`.
- [ ] Run the formatter to format the code using `make lint`.


### Adding a model checklist
<!--
When adding a model to the model registry
see also https://github.com/embeddings-benchmark/mteb/blob/main/docs/reproducible_workflow.md
-->

 - [x] I have filled out the ModelMeta object to the extent possible
 - [x] I have ensured that my model can be loaded using
   - [x] `mteb.get_model(model_name, revision)` and
   - [x] `mteb.get_model_meta(model_name, revision)`
 - [x] I have tested the implementation works on a representative set of tasks.



Diff:
diff --git a/mteb/models/ops_moa_models.py b/mteb/models/ops_moa_models.py
new file mode 100644
index 0000000000..965b5800a2
--- /dev/null
+++ b/mteb/models/ops_moa_models.py
@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+from mteb.model_meta import ModelMeta
+from mteb.models.wrapper import Wrapper
+from functools import partial
+from sentence_transformers import SentenceTransformer
+import torch
+import torch.nn as nn
+from huggingface_hub import snapshot_download
+
+
+class CustomWrapper(Wrapper):
+    def __init__(self, model_name, revision):
+        super().__init__()
+        self.model = SentenceTransformer(
+            model_name, revision=revision, trust_remote_code=True
+        )
+        self.output_dim = 1536
+
+    def encode(self, sentences: list[str], **kwargs) -> np.ndarray:
+        embeddings = self.model.encode(sentences, **kwargs)
+        return embeddings[:, : self.output_dim]
+
+
+ops_moa_conan_embedding = ModelMeta(
+    name="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+    revision="46dcd58753f3daa920c66f89e47086a534089350",
+    release_date="2025-03-26",
+    languages=["zho_Hans"],
+    loader=partial(
+        CustomWrapper,
+        "OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+        "46dcd58753f3daa920c66f89e47086a534089350",
+    ),
+    n_parameters=343 * 1e6,
+    memory_usage_mb=2e3,
+    max_tokens=512,
+    embed_dim=1536,
+    license="cc-by-nc-4.0",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=["PyTorch", "Sentence Transformers"],
+    reference="https://huggingface.co/OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+    similarity_fn_name="cosine",
+    use_instructions=False,
+    training_datasets={
+        "T2Retrieval": ["train"],
+        "MMarcoRetrieval": ["train"],
+        "DuRetrieval": ["train"],
+        "CovidRetrieval": ["train"],
+        "CmedqaRetrieval": ["train"],
+        "EcomRetrieval": ["train"],
+        "MedicalRetrieval": ["train"],
+        "VideoRetrieval": ["train"],
+    },
+    superseded_by=None,
+)
+
+ops_moa_yuan_embedding = ModelMeta(
+    name="OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0",
+    revision="23712d0766417b0eb88a2513c6e212a58b543268",
+    release_date="2025-03-26",
+    languages=["zho_Hans"],
+    loader=partial(
+        CustomWrapper,
+        "OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0",
+        "23712d0766417b0eb88a2513c6e212a58b543268",
+    ),
+    n_parameters=343 * 1e6,
+    memory_usage_mb=2e3,
+    max_tokens=512,
+    embed_dim=1536,
+    license="cc-by-nc-4.0",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=["PyTorch", "Sentence Transformers"],
+    reference="https://huggingface.co/OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0",
+    similarity_fn_name="cosine",
+    use_instructions=False,
+    training_datasets={
+        "T2Retrieval": ["train"],
+        "MMarcoRetrieval": ["train"],
+        "DuRetrieval": ["train"],
+        "CovidRetrieval": ["train"],
+        "CmedqaRetrieval": ["train"],
+        "EcomRetrieval": ["train"],
+        "MedicalRetrieval": ["train"],
+        "VideoRetrieval": ["train"],
+    },
+    superseded_by=None,
+)
diff --git a/mteb/models/overview.py b/mteb/models/overview.py
index a715ab2730..a2c4267da6 100644
--- a/mteb/models/overview.py
+++ b/mteb/models/overview.py
@@ -57,6 +57,7 @@
     nvidia_models,
     openai_models,
     openclip_models,
+    ops_moa_models,
     piccolo_models,
     promptriever_models,
     qodo_models,
@@ -124,6 +125,7 @@
     nvidia_models,
     openai_models,
     openclip_models,
+    ops_moa_models,
     piccolo_models,
     gme_v_models,
     promptriever_models,


Reviews:
user: login: Samoed
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-26T14:15:03Z
user: login: Samoed
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-28T07:42:13Z
user: login: ahxgw
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-28T08:06:50Z
user: login: Samoed
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-28T08:14:41Z
user: login: ahxgw
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-28T11:00:17Z
user: login: Samoed
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-28T14:22:06Z
user: login: ahxgw
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-28T15:13:57Z
user: login: Samoed
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-28T18:29:59Z
user: login: ahxgw
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-29T01:24:32Z
user: login: Samoed
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-29T10:35:52Z
user: login: Samoed
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-29T15:31:02Z

Review comments:
diff_hunk: @@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from mteb.model_meta import ModelMeta
+
+ops_moa_conan_embedding = ModelMeta(
+    name="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+    revision="cd42de6d61c103047b7bcd780ef0dbaa9a9d0472",
+    release_date="2025-03-26",
+    languages=["zho_Hans"],
+    loader=None,
path: mteb/models/ops_moa_models.py
user: login: Samoed
body: Can you add implementation of your model? Because it seems your model required some more additions
created_at: 2025-03-26T14:13:29Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from mteb.model_meta import ModelMeta
+
+ops_moa_conan_embedding = ModelMeta(
+    name="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+    revision="cd42de6d61c103047b7bcd780ef0dbaa9a9d0472",
+    release_date="2025-03-26",
+    languages=["zho_Hans"],
+    loader=None,
+    n_parameters=343 * 1e6,
+    memory_usage_mb=None,
path: mteb/models/ops_moa_models.py
user: login: Samoed
body: Can you add it?
created_at: 2025-03-26T14:13:37Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from mteb.model_meta import ModelMeta
+
+ops_moa_conan_embedding = ModelMeta(
+    name="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+    revision="cd42de6d61c103047b7bcd780ef0dbaa9a9d0472",
+    release_date="2025-03-26",
+    languages=["zho_Hans"],
+    loader=None,
+    n_parameters=343 * 1e6,
+    memory_usage_mb=None,
+    max_tokens=512,
+    embed_dim=1536,
+    license=None
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
path: mteb/models/ops_moa_models.py
user: login: Samoed
body: Can you add datasets?
created_at: 2025-03-26T14:14:08Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from mteb.model_meta import ModelMeta
+
+ops_moa_conan_embedding = ModelMeta(
+    name="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+    revision="cd42de6d61c103047b7bcd780ef0dbaa9a9d0472",
+    release_date="2025-03-26",
+    languages=["zho_Hans"],
+    loader=None,
+    n_parameters=343 * 1e6,
+    memory_usage_mb=None,
+    max_tokens=512,
+    embed_dim=1536,
+    license=None
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=["PyTorch", "Sentence Transformers"],
+    reference="https://huggingface.co/OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+    similarity_fn_name="cosine",
+    use_instructions=None,
path: mteb/models/ops_moa_models.py
user: login: Samoed
body: ```suggestion
    use_instructions=False,
```
created_at: 2025-03-26T14:14:24Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from mteb.model_meta import ModelMeta
+
+ops_moa_conan_embedding = ModelMeta(
+    name="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+    revision="cd42de6d61c103047b7bcd780ef0dbaa9a9d0472",
+    release_date="2025-03-26",
+    languages=["zho_Hans"],
+    loader=None,
path: mteb/models/ops_moa_models.py
user: login: Samoed
body: I think it would be better to add custom implementations directly in your model
created_at: 2025-03-26T14:14:59Z
author_association: MEMBER
diff_hunk: @@ -8,143 +8,64 @@
 import torch.nn as nn
 from huggingface_hub import snapshot_download
 
-class MoAGate(nn.Module):
-    def __init__(self, num_adaptors, hidden_dim):
-        super().__init__()
-        self.routing_vectors = nn.Parameter(
-                torch.empty(num_adaptors, hidden_dim, dtype=torch.float32),
-                requires_grad=False
-            )
-    def forward(self, hidden_states):
-        if self.routing_vectors.device == torch.device('cpu'):
-            self.routing_vectors = self.routing_vectors.to(hidden_states.device)
-        hidden_states = hidden_states.unsqueeze(1)
-        batch_size, seq_len, hidden_dim = hidden_states.shape
-
-        hidden_states = hidden_states.view(-1, hidden_dim)
-        distances = torch.cdist(hidden_states, self.routing_vectors)
-
-        _, cluster_indices = torch.min(distances, dim=1)
-        cluster_indices = cluster_indices.view(-1, 1)
-
-        topk_indices = cluster_indices
-        topk_indices = torch.zeros_like(topk_indices, device=hidden_states.device)
-        topk_weights = torch.ones_like(topk_indices, device=hidden_states.device)
-
-        return topk_indices, topk_weights
-
-class LinearLayer(nn.Module):
-    def __init__(self, input_dim, output_dim):
-        super().__init__()
-        self.linear = nn.Linear(input_dim, output_dim)
-
-    def forward(self, x):
-        return self.linear(x)
-
-class MixtureOfAdaptors(nn.Module):
-    def __init__(self, num_adaptors, hidden_dim):
-        super().__init__()
-        self.adaptors = nn.ModuleList([
-            LinearLayer(input_dim=hidden_dim, output_dim=hidden_dim)
-            for _ in range(num_adaptors)
-        ])
-        self.gate = MoAGate(num_adaptors, hidden_dim)
-
-    def forward(self, inputs):
-        if isinstance(inputs, dict):
-            hidden_states = inputs['sentence_embedding']
-        else:
-            hidden_states = inputs
-
-        residual = hidden_states
-        original_shape = hidden_states.shape
-        topk_indices, topk_weights = self.gate(hidden_states)
-
-        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-        flat_topk_indices = topk_indices.view(-1)
-        output = self.moa_inference(hidden_states, flat_topk_indices, topk_weights.view(-1, 1)).view(*original_shape)
-
-        if isinstance(inputs, dict):
-            inputs['sentence_embedding'] = output
-            return inputs
-        return output
-
-    @torch.no_grad()
-    def moa_inference(self, x, flat_adaptor_indices, flat_adaptor_weights):
-        adaptor_cache = torch.zeros_like(x)
-        sorted_indices = flat_adaptor_indices.argsort()
-        tokens_per_adaptor = flat_adaptor_indices.bincount().cpu().numpy().cumsum(0)
-        token_indices = sorted_indices
-        for i, end_idx in enumerate(tokens_per_adaptor):
-            start_idx = 0 if i == 0 else tokens_per_adaptor[i-1]
-            if start_idx == end_idx:
-                continue
-            adaptor = self.adaptors[i]
-            adaptor_token_indices = token_indices[start_idx:end_idx]
-            adaptor_tokens = x[adaptor_token_indices]
-            adaptor_output = adaptor(adaptor_tokens)
-            adaptor_output.mul_(flat_adaptor_weights[sorted_indices[start_idx:end_idx]])
-            adaptor_cache.scatter_reduce_(
-                0,
-                adaptor_token_indices.view(-1, 1).repeat(1, x.shape[-1]),
-                adaptor_output,
-                reduce='sum'
-            )
-        return adaptor_cache
 
 class YuanWrapper(Wrapper):
     def __init__(self):
         super().__init__()
-        self.model = SentenceTransformer('IEITYuan/Yuan-embedding-1.0', trust_remote_code=True)
-        adaptor = MixtureOfAdaptors(5, 1792)
+        self.model = SentenceTransformer(
+            "IEITYuan/Yuan-embedding-1.0", trust_remote_code=True
+        )
 
         model_path = snapshot_download(
             repo_id="OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0",
-            local_dir="./Ops-MoA-Yuan-embedding-1.0"
+            local_dir="./Ops_MoA_Yuan/",
         )
-        adaptor.load_state_dict(torch.load(f"Ops-MoA-Yuan-embedding-1.0/yuan-adaptors.pth"))
+        from Ops_MoA_Yuan.modeling_adaptor import MixtureOfAdaptors
+
+        adaptor = MixtureOfAdaptors(5, 1792)
 
-        self.model.add_module('adaptor', adaptor)
+        adaptor.load_state_dict(torch.load(f"Ops_MoA_Yuan/yuan-adaptors.pth"))
+
+        self.model.add_module("adaptor", adaptor)
         self.output_dim = 1536
 
-    def encode(
-         self,
-         sentences: list[str],
-         **kwargs
-    ) -> np.ndarray:
+    def encode(self, sentences: list[str], **kwargs) -> np.ndarray:
         embeddings = self.model.encode(sentences, **kwargs)
-        return embeddings[:, :self.output_dim]
+        return embeddings[:, : self.output_dim]
+
 
 class ConanWrapper(Wrapper):
     def __init__(self):
         super().__init__()
-        self.model = SentenceTransformer('TencentBAC/Conan-embedding-v1', trust_remote_code=True)
-        adaptor = MixtureOfAdaptors(5, 1792)
+        self.model = SentenceTransformer(
+            "TencentBAC/Conan-embedding-v1", trust_remote_code=True
+        )
 
         model_path = snapshot_download(
             repo_id="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
-            local_dir="./Ops-MoA-Conan-embedding-v1"
+            local_dir="./Ops_MoA_Conan",
         )
-        adaptor.load_state_dict(torch.load(f"Ops-MoA-Conan-embedding-v1/conan-adaptors.pth"))
+        from Ops_MoA_Conan.modeling_adaptor import MixtureOfAdaptors
 
-        self.model.add_module('adaptor', adaptor)
+        adaptor = MixtureOfAdaptors(5, 1792)
+
+        adaptor.load_state_dict(torch.load(f"Ops_MoA_Conan/conan-adaptors.pth"))
+
+        self.model.add_module("adaptor", adaptor)
         self.output_dim = 1536
 
-    def encode(
-         self,
-         sentences: list[str],
-         **kwargs
-    ) -> np.ndarray:
+    def encode(self, sentences: list[str], **kwargs) -> np.ndarray:
         embeddings = self.model.encode(sentences, **kwargs)
-        return embeddings[:, :self.output_dim]
+        return embeddings[:, : self.output_dim]
+
 
 ops_moa_conan_embedding = ModelMeta(
     name="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
path: mteb/models/ops_moa_models.py
user: login: Samoed
body: Can you upload full model `TencentBAC/Conan-embedding-v1` to your repository to integrate your model with SentenceTransformers directly
created_at: 2025-03-28T07:42:13Z
author_association: MEMBER
diff_hunk: @@ -8,143 +8,64 @@
 import torch.nn as nn
 from huggingface_hub import snapshot_download
 
-class MoAGate(nn.Module):
-    def __init__(self, num_adaptors, hidden_dim):
-        super().__init__()
-        self.routing_vectors = nn.Parameter(
-                torch.empty(num_adaptors, hidden_dim, dtype=torch.float32),
-                requires_grad=False
-            )
-    def forward(self, hidden_states):
-        if self.routing_vectors.device == torch.device('cpu'):
-            self.routing_vectors = self.routing_vectors.to(hidden_states.device)
-        hidden_states = hidden_states.unsqueeze(1)
-        batch_size, seq_len, hidden_dim = hidden_states.shape
-
-        hidden_states = hidden_states.view(-1, hidden_dim)
-        distances = torch.cdist(hidden_states, self.routing_vectors)
-
-        _, cluster_indices = torch.min(distances, dim=1)
-        cluster_indices = cluster_indices.view(-1, 1)
-
-        topk_indices = cluster_indices
-        topk_indices = torch.zeros_like(topk_indices, device=hidden_states.device)
-        topk_weights = torch.ones_like(topk_indices, device=hidden_states.device)
-
-        return topk_indices, topk_weights
-
-class LinearLayer(nn.Module):
-    def __init__(self, input_dim, output_dim):
-        super().__init__()
-        self.linear = nn.Linear(input_dim, output_dim)
-
-    def forward(self, x):
-        return self.linear(x)
-
-class MixtureOfAdaptors(nn.Module):
-    def __init__(self, num_adaptors, hidden_dim):
-        super().__init__()
-        self.adaptors = nn.ModuleList([
-            LinearLayer(input_dim=hidden_dim, output_dim=hidden_dim)
-            for _ in range(num_adaptors)
-        ])
-        self.gate = MoAGate(num_adaptors, hidden_dim)
-
-    def forward(self, inputs):
-        if isinstance(inputs, dict):
-            hidden_states = inputs['sentence_embedding']
-        else:
-            hidden_states = inputs
-
-        residual = hidden_states
-        original_shape = hidden_states.shape
-        topk_indices, topk_weights = self.gate(hidden_states)
-
-        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-        flat_topk_indices = topk_indices.view(-1)
-        output = self.moa_inference(hidden_states, flat_topk_indices, topk_weights.view(-1, 1)).view(*original_shape)
-
-        if isinstance(inputs, dict):
-            inputs['sentence_embedding'] = output
-            return inputs
-        return output
-
-    @torch.no_grad()
-    def moa_inference(self, x, flat_adaptor_indices, flat_adaptor_weights):
-        adaptor_cache = torch.zeros_like(x)
-        sorted_indices = flat_adaptor_indices.argsort()
-        tokens_per_adaptor = flat_adaptor_indices.bincount().cpu().numpy().cumsum(0)
-        token_indices = sorted_indices
-        for i, end_idx in enumerate(tokens_per_adaptor):
-            start_idx = 0 if i == 0 else tokens_per_adaptor[i-1]
-            if start_idx == end_idx:
-                continue
-            adaptor = self.adaptors[i]
-            adaptor_token_indices = token_indices[start_idx:end_idx]
-            adaptor_tokens = x[adaptor_token_indices]
-            adaptor_output = adaptor(adaptor_tokens)
-            adaptor_output.mul_(flat_adaptor_weights[sorted_indices[start_idx:end_idx]])
-            adaptor_cache.scatter_reduce_(
-                0,
-                adaptor_token_indices.view(-1, 1).repeat(1, x.shape[-1]),
-                adaptor_output,
-                reduce='sum'
-            )
-        return adaptor_cache
 
 class YuanWrapper(Wrapper):
     def __init__(self):
         super().__init__()
-        self.model = SentenceTransformer('IEITYuan/Yuan-embedding-1.0', trust_remote_code=True)
-        adaptor = MixtureOfAdaptors(5, 1792)
+        self.model = SentenceTransformer(
+            "IEITYuan/Yuan-embedding-1.0", trust_remote_code=True
+        )
 
         model_path = snapshot_download(
             repo_id="OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0",
-            local_dir="./Ops-MoA-Yuan-embedding-1.0"
+            local_dir="./Ops_MoA_Yuan/",
         )
-        adaptor.load_state_dict(torch.load(f"Ops-MoA-Yuan-embedding-1.0/yuan-adaptors.pth"))
+        from Ops_MoA_Yuan.modeling_adaptor import MixtureOfAdaptors
+
+        adaptor = MixtureOfAdaptors(5, 1792)
 
-        self.model.add_module('adaptor', adaptor)
+        adaptor.load_state_dict(torch.load(f"Ops_MoA_Yuan/yuan-adaptors.pth"))
+
+        self.model.add_module("adaptor", adaptor)
         self.output_dim = 1536
 
-    def encode(
-         self,
-         sentences: list[str],
-         **kwargs
-    ) -> np.ndarray:
+    def encode(self, sentences: list[str], **kwargs) -> np.ndarray:
         embeddings = self.model.encode(sentences, **kwargs)
-        return embeddings[:, :self.output_dim]
+        return embeddings[:, : self.output_dim]
+
 
 class ConanWrapper(Wrapper):
     def __init__(self):
         super().__init__()
-        self.model = SentenceTransformer('TencentBAC/Conan-embedding-v1', trust_remote_code=True)
-        adaptor = MixtureOfAdaptors(5, 1792)
+        self.model = SentenceTransformer(
+            "TencentBAC/Conan-embedding-v1", trust_remote_code=True
+        )
 
         model_path = snapshot_download(
             repo_id="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
-            local_dir="./Ops-MoA-Conan-embedding-v1"
+            local_dir="./Ops_MoA_Conan",
         )
-        adaptor.load_state_dict(torch.load(f"Ops-MoA-Conan-embedding-v1/conan-adaptors.pth"))
+        from Ops_MoA_Conan.modeling_adaptor import MixtureOfAdaptors
 
-        self.model.add_module('adaptor', adaptor)
+        adaptor = MixtureOfAdaptors(5, 1792)
+
+        adaptor.load_state_dict(torch.load(f"Ops_MoA_Conan/conan-adaptors.pth"))
+
+        self.model.add_module("adaptor", adaptor)
         self.output_dim = 1536
 
-    def encode(
-         self,
-         sentences: list[str],
-         **kwargs
-    ) -> np.ndarray:
+    def encode(self, sentences: list[str], **kwargs) -> np.ndarray:
         embeddings = self.model.encode(sentences, **kwargs)
-        return embeddings[:, :self.output_dim]
+        return embeddings[:, : self.output_dim]
+
 
 ops_moa_conan_embedding = ModelMeta(
     name="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
path: mteb/models/ops_moa_models.py
user: login: ahxgw
body: Sorry, due to the special adaptor code of our model, we cannot directly support SentenceTransformers yet.
created_at: 2025-03-28T08:06:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -8,143 +8,64 @@
 import torch.nn as nn
 from huggingface_hub import snapshot_download
 
-class MoAGate(nn.Module):
-    def __init__(self, num_adaptors, hidden_dim):
-        super().__init__()
-        self.routing_vectors = nn.Parameter(
-                torch.empty(num_adaptors, hidden_dim, dtype=torch.float32),
-                requires_grad=False
-            )
-    def forward(self, hidden_states):
-        if self.routing_vectors.device == torch.device('cpu'):
-            self.routing_vectors = self.routing_vectors.to(hidden_states.device)
-        hidden_states = hidden_states.unsqueeze(1)
-        batch_size, seq_len, hidden_dim = hidden_states.shape
-
-        hidden_states = hidden_states.view(-1, hidden_dim)
-        distances = torch.cdist(hidden_states, self.routing_vectors)
-
-        _, cluster_indices = torch.min(distances, dim=1)
-        cluster_indices = cluster_indices.view(-1, 1)
-
-        topk_indices = cluster_indices
-        topk_indices = torch.zeros_like(topk_indices, device=hidden_states.device)
-        topk_weights = torch.ones_like(topk_indices, device=hidden_states.device)
-
-        return topk_indices, topk_weights
-
-class LinearLayer(nn.Module):
-    def __init__(self, input_dim, output_dim):
-        super().__init__()
-        self.linear = nn.Linear(input_dim, output_dim)
-
-    def forward(self, x):
-        return self.linear(x)
-
-class MixtureOfAdaptors(nn.Module):
-    def __init__(self, num_adaptors, hidden_dim):
-        super().__init__()
-        self.adaptors = nn.ModuleList([
-            LinearLayer(input_dim=hidden_dim, output_dim=hidden_dim)
-            for _ in range(num_adaptors)
-        ])
-        self.gate = MoAGate(num_adaptors, hidden_dim)
-
-    def forward(self, inputs):
-        if isinstance(inputs, dict):
-            hidden_states = inputs['sentence_embedding']
-        else:
-            hidden_states = inputs
-
-        residual = hidden_states
-        original_shape = hidden_states.shape
-        topk_indices, topk_weights = self.gate(hidden_states)
-
-        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-        flat_topk_indices = topk_indices.view(-1)
-        output = self.moa_inference(hidden_states, flat_topk_indices, topk_weights.view(-1, 1)).view(*original_shape)
-
-        if isinstance(inputs, dict):
-            inputs['sentence_embedding'] = output
-            return inputs
-        return output
-
-    @torch.no_grad()
-    def moa_inference(self, x, flat_adaptor_indices, flat_adaptor_weights):
-        adaptor_cache = torch.zeros_like(x)
-        sorted_indices = flat_adaptor_indices.argsort()
-        tokens_per_adaptor = flat_adaptor_indices.bincount().cpu().numpy().cumsum(0)
-        token_indices = sorted_indices
-        for i, end_idx in enumerate(tokens_per_adaptor):
-            start_idx = 0 if i == 0 else tokens_per_adaptor[i-1]
-            if start_idx == end_idx:
-                continue
-            adaptor = self.adaptors[i]
-            adaptor_token_indices = token_indices[start_idx:end_idx]
-            adaptor_tokens = x[adaptor_token_indices]
-            adaptor_output = adaptor(adaptor_tokens)
-            adaptor_output.mul_(flat_adaptor_weights[sorted_indices[start_idx:end_idx]])
-            adaptor_cache.scatter_reduce_(
-                0,
-                adaptor_token_indices.view(-1, 1).repeat(1, x.shape[-1]),
-                adaptor_output,
-                reduce='sum'
-            )
-        return adaptor_cache
 
 class YuanWrapper(Wrapper):
     def __init__(self):
         super().__init__()
-        self.model = SentenceTransformer('IEITYuan/Yuan-embedding-1.0', trust_remote_code=True)
-        adaptor = MixtureOfAdaptors(5, 1792)
+        self.model = SentenceTransformer(
+            "IEITYuan/Yuan-embedding-1.0", trust_remote_code=True
+        )
 
         model_path = snapshot_download(
             repo_id="OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0",
-            local_dir="./Ops-MoA-Yuan-embedding-1.0"
+            local_dir="./Ops_MoA_Yuan/",
         )
-        adaptor.load_state_dict(torch.load(f"Ops-MoA-Yuan-embedding-1.0/yuan-adaptors.pth"))
+        from Ops_MoA_Yuan.modeling_adaptor import MixtureOfAdaptors
+
+        adaptor = MixtureOfAdaptors(5, 1792)
 
-        self.model.add_module('adaptor', adaptor)
+        adaptor.load_state_dict(torch.load(f"Ops_MoA_Yuan/yuan-adaptors.pth"))
+
+        self.model.add_module("adaptor", adaptor)
         self.output_dim = 1536
 
-    def encode(
-         self,
-         sentences: list[str],
-         **kwargs
-    ) -> np.ndarray:
+    def encode(self, sentences: list[str], **kwargs) -> np.ndarray:
         embeddings = self.model.encode(sentences, **kwargs)
-        return embeddings[:, :self.output_dim]
+        return embeddings[:, : self.output_dim]
+
 
 class ConanWrapper(Wrapper):
     def __init__(self):
         super().__init__()
-        self.model = SentenceTransformer('TencentBAC/Conan-embedding-v1', trust_remote_code=True)
-        adaptor = MixtureOfAdaptors(5, 1792)
+        self.model = SentenceTransformer(
+            "TencentBAC/Conan-embedding-v1", trust_remote_code=True
+        )
 
         model_path = snapshot_download(
             repo_id="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
-            local_dir="./Ops-MoA-Conan-embedding-v1"
+            local_dir="./Ops_MoA_Conan",
         )
-        adaptor.load_state_dict(torch.load(f"Ops-MoA-Conan-embedding-v1/conan-adaptors.pth"))
+        from Ops_MoA_Conan.modeling_adaptor import MixtureOfAdaptors
 
-        self.model.add_module('adaptor', adaptor)
+        adaptor = MixtureOfAdaptors(5, 1792)
+
+        adaptor.load_state_dict(torch.load(f"Ops_MoA_Conan/conan-adaptors.pth"))
+
+        self.model.add_module("adaptor", adaptor)
         self.output_dim = 1536
 
-    def encode(
-         self,
-         sentences: list[str],
-         **kwargs
-    ) -> np.ndarray:
+    def encode(self, sentences: list[str], **kwargs) -> np.ndarray:
         embeddings = self.model.encode(sentences, **kwargs)
-        return embeddings[:, :self.output_dim]
+        return embeddings[:, : self.output_dim]
+
 
 ops_moa_conan_embedding = ModelMeta(
     name="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
path: mteb/models/ops_moa_models.py
user: login: Samoed
body: If you would upload `TencentBAC/Conan-embedding-v1` I will easily add support for SentenceTransformers
created_at: 2025-03-28T08:14:41Z
author_association: MEMBER
diff_hunk: @@ -8,143 +8,64 @@
 import torch.nn as nn
 from huggingface_hub import snapshot_download
 
-class MoAGate(nn.Module):
-    def __init__(self, num_adaptors, hidden_dim):
-        super().__init__()
-        self.routing_vectors = nn.Parameter(
-                torch.empty(num_adaptors, hidden_dim, dtype=torch.float32),
-                requires_grad=False
-            )
-    def forward(self, hidden_states):
-        if self.routing_vectors.device == torch.device('cpu'):
-            self.routing_vectors = self.routing_vectors.to(hidden_states.device)
-        hidden_states = hidden_states.unsqueeze(1)
-        batch_size, seq_len, hidden_dim = hidden_states.shape
-
-        hidden_states = hidden_states.view(-1, hidden_dim)
-        distances = torch.cdist(hidden_states, self.routing_vectors)
-
-        _, cluster_indices = torch.min(distances, dim=1)
-        cluster_indices = cluster_indices.view(-1, 1)
-
-        topk_indices = cluster_indices
-        topk_indices = torch.zeros_like(topk_indices, device=hidden_states.device)
-        topk_weights = torch.ones_like(topk_indices, device=hidden_states.device)
-
-        return topk_indices, topk_weights
-
-class LinearLayer(nn.Module):
-    def __init__(self, input_dim, output_dim):
-        super().__init__()
-        self.linear = nn.Linear(input_dim, output_dim)
-
-    def forward(self, x):
-        return self.linear(x)
-
-class MixtureOfAdaptors(nn.Module):
-    def __init__(self, num_adaptors, hidden_dim):
-        super().__init__()
-        self.adaptors = nn.ModuleList([
-            LinearLayer(input_dim=hidden_dim, output_dim=hidden_dim)
-            for _ in range(num_adaptors)
-        ])
-        self.gate = MoAGate(num_adaptors, hidden_dim)
-
-    def forward(self, inputs):
-        if isinstance(inputs, dict):
-            hidden_states = inputs['sentence_embedding']
-        else:
-            hidden_states = inputs
-
-        residual = hidden_states
-        original_shape = hidden_states.shape
-        topk_indices, topk_weights = self.gate(hidden_states)
-
-        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-        flat_topk_indices = topk_indices.view(-1)
-        output = self.moa_inference(hidden_states, flat_topk_indices, topk_weights.view(-1, 1)).view(*original_shape)
-
-        if isinstance(inputs, dict):
-            inputs['sentence_embedding'] = output
-            return inputs
-        return output
-
-    @torch.no_grad()
-    def moa_inference(self, x, flat_adaptor_indices, flat_adaptor_weights):
-        adaptor_cache = torch.zeros_like(x)
-        sorted_indices = flat_adaptor_indices.argsort()
-        tokens_per_adaptor = flat_adaptor_indices.bincount().cpu().numpy().cumsum(0)
-        token_indices = sorted_indices
-        for i, end_idx in enumerate(tokens_per_adaptor):
-            start_idx = 0 if i == 0 else tokens_per_adaptor[i-1]
-            if start_idx == end_idx:
-                continue
-            adaptor = self.adaptors[i]
-            adaptor_token_indices = token_indices[start_idx:end_idx]
-            adaptor_tokens = x[adaptor_token_indices]
-            adaptor_output = adaptor(adaptor_tokens)
-            adaptor_output.mul_(flat_adaptor_weights[sorted_indices[start_idx:end_idx]])
-            adaptor_cache.scatter_reduce_(
-                0,
-                adaptor_token_indices.view(-1, 1).repeat(1, x.shape[-1]),
-                adaptor_output,
-                reduce='sum'
-            )
-        return adaptor_cache
 
 class YuanWrapper(Wrapper):
     def __init__(self):
         super().__init__()
-        self.model = SentenceTransformer('IEITYuan/Yuan-embedding-1.0', trust_remote_code=True)
-        adaptor = MixtureOfAdaptors(5, 1792)
+        self.model = SentenceTransformer(
+            "IEITYuan/Yuan-embedding-1.0", trust_remote_code=True
+        )
 
         model_path = snapshot_download(
             repo_id="OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0",
-            local_dir="./Ops-MoA-Yuan-embedding-1.0"
+            local_dir="./Ops_MoA_Yuan/",
         )
-        adaptor.load_state_dict(torch.load(f"Ops-MoA-Yuan-embedding-1.0/yuan-adaptors.pth"))
+        from Ops_MoA_Yuan.modeling_adaptor import MixtureOfAdaptors
+
+        adaptor = MixtureOfAdaptors(5, 1792)
 
-        self.model.add_module('adaptor', adaptor)
+        adaptor.load_state_dict(torch.load(f"Ops_MoA_Yuan/yuan-adaptors.pth"))
+
+        self.model.add_module("adaptor", adaptor)
         self.output_dim = 1536
 
-    def encode(
-         self,
-         sentences: list[str],
-         **kwargs
-    ) -> np.ndarray:
+    def encode(self, sentences: list[str], **kwargs) -> np.ndarray:
         embeddings = self.model.encode(sentences, **kwargs)
-        return embeddings[:, :self.output_dim]
+        return embeddings[:, : self.output_dim]
+
 
 class ConanWrapper(Wrapper):
     def __init__(self):
         super().__init__()
-        self.model = SentenceTransformer('TencentBAC/Conan-embedding-v1', trust_remote_code=True)
-        adaptor = MixtureOfAdaptors(5, 1792)
+        self.model = SentenceTransformer(
+            "TencentBAC/Conan-embedding-v1", trust_remote_code=True
+        )
 
         model_path = snapshot_download(
             repo_id="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
-            local_dir="./Ops-MoA-Conan-embedding-v1"
+            local_dir="./Ops_MoA_Conan",
         )
-        adaptor.load_state_dict(torch.load(f"Ops-MoA-Conan-embedding-v1/conan-adaptors.pth"))
+        from Ops_MoA_Conan.modeling_adaptor import MixtureOfAdaptors
 
-        self.model.add_module('adaptor', adaptor)
+        adaptor = MixtureOfAdaptors(5, 1792)
+
+        adaptor.load_state_dict(torch.load(f"Ops_MoA_Conan/conan-adaptors.pth"))
+
+        self.model.add_module("adaptor", adaptor)
         self.output_dim = 1536
 
-    def encode(
-         self,
-         sentences: list[str],
-         **kwargs
-    ) -> np.ndarray:
+    def encode(self, sentences: list[str], **kwargs) -> np.ndarray:
         embeddings = self.model.encode(sentences, **kwargs)
-        return embeddings[:, :self.output_dim]
+        return embeddings[:, : self.output_dim]
+
 
 ops_moa_conan_embedding = ModelMeta(
     name="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
path: mteb/models/ops_moa_models.py
user: login: ahxgw
body: We have uploaded two full models to OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0 and OpenSearch-AI/Ops-MoA-Conan-embedding-v1.
created_at: 2025-03-28T11:00:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -8,143 +8,64 @@
 import torch.nn as nn
 from huggingface_hub import snapshot_download
 
-class MoAGate(nn.Module):
-    def __init__(self, num_adaptors, hidden_dim):
-        super().__init__()
-        self.routing_vectors = nn.Parameter(
-                torch.empty(num_adaptors, hidden_dim, dtype=torch.float32),
-                requires_grad=False
-            )
-    def forward(self, hidden_states):
-        if self.routing_vectors.device == torch.device('cpu'):
-            self.routing_vectors = self.routing_vectors.to(hidden_states.device)
-        hidden_states = hidden_states.unsqueeze(1)
-        batch_size, seq_len, hidden_dim = hidden_states.shape
-
-        hidden_states = hidden_states.view(-1, hidden_dim)
-        distances = torch.cdist(hidden_states, self.routing_vectors)
-
-        _, cluster_indices = torch.min(distances, dim=1)
-        cluster_indices = cluster_indices.view(-1, 1)
-
-        topk_indices = cluster_indices
-        topk_indices = torch.zeros_like(topk_indices, device=hidden_states.device)
-        topk_weights = torch.ones_like(topk_indices, device=hidden_states.device)
-
-        return topk_indices, topk_weights
-
-class LinearLayer(nn.Module):
-    def __init__(self, input_dim, output_dim):
-        super().__init__()
-        self.linear = nn.Linear(input_dim, output_dim)
-
-    def forward(self, x):
-        return self.linear(x)
-
-class MixtureOfAdaptors(nn.Module):
-    def __init__(self, num_adaptors, hidden_dim):
-        super().__init__()
-        self.adaptors = nn.ModuleList([
-            LinearLayer(input_dim=hidden_dim, output_dim=hidden_dim)
-            for _ in range(num_adaptors)
-        ])
-        self.gate = MoAGate(num_adaptors, hidden_dim)
-
-    def forward(self, inputs):
-        if isinstance(inputs, dict):
-            hidden_states = inputs['sentence_embedding']
-        else:
-            hidden_states = inputs
-
-        residual = hidden_states
-        original_shape = hidden_states.shape
-        topk_indices, topk_weights = self.gate(hidden_states)
-
-        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-        flat_topk_indices = topk_indices.view(-1)
-        output = self.moa_inference(hidden_states, flat_topk_indices, topk_weights.view(-1, 1)).view(*original_shape)
-
-        if isinstance(inputs, dict):
-            inputs['sentence_embedding'] = output
-            return inputs
-        return output
-
-    @torch.no_grad()
-    def moa_inference(self, x, flat_adaptor_indices, flat_adaptor_weights):
-        adaptor_cache = torch.zeros_like(x)
-        sorted_indices = flat_adaptor_indices.argsort()
-        tokens_per_adaptor = flat_adaptor_indices.bincount().cpu().numpy().cumsum(0)
-        token_indices = sorted_indices
-        for i, end_idx in enumerate(tokens_per_adaptor):
-            start_idx = 0 if i == 0 else tokens_per_adaptor[i-1]
-            if start_idx == end_idx:
-                continue
-            adaptor = self.adaptors[i]
-            adaptor_token_indices = token_indices[start_idx:end_idx]
-            adaptor_tokens = x[adaptor_token_indices]
-            adaptor_output = adaptor(adaptor_tokens)
-            adaptor_output.mul_(flat_adaptor_weights[sorted_indices[start_idx:end_idx]])
-            adaptor_cache.scatter_reduce_(
-                0,
-                adaptor_token_indices.view(-1, 1).repeat(1, x.shape[-1]),
-                adaptor_output,
-                reduce='sum'
-            )
-        return adaptor_cache
 
 class YuanWrapper(Wrapper):
     def __init__(self):
         super().__init__()
-        self.model = SentenceTransformer('IEITYuan/Yuan-embedding-1.0', trust_remote_code=True)
-        adaptor = MixtureOfAdaptors(5, 1792)
+        self.model = SentenceTransformer(
+            "IEITYuan/Yuan-embedding-1.0", trust_remote_code=True
+        )
 
         model_path = snapshot_download(
             repo_id="OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0",
-            local_dir="./Ops-MoA-Yuan-embedding-1.0"
+            local_dir="./Ops_MoA_Yuan/",
         )
-        adaptor.load_state_dict(torch.load(f"Ops-MoA-Yuan-embedding-1.0/yuan-adaptors.pth"))
+        from Ops_MoA_Yuan.modeling_adaptor import MixtureOfAdaptors
+
+        adaptor = MixtureOfAdaptors(5, 1792)
 
-        self.model.add_module('adaptor', adaptor)
+        adaptor.load_state_dict(torch.load(f"Ops_MoA_Yuan/yuan-adaptors.pth"))
+
+        self.model.add_module("adaptor", adaptor)
         self.output_dim = 1536
 
-    def encode(
-         self,
-         sentences: list[str],
-         **kwargs
-    ) -> np.ndarray:
+    def encode(self, sentences: list[str], **kwargs) -> np.ndarray:
         embeddings = self.model.encode(sentences, **kwargs)
-        return embeddings[:, :self.output_dim]
+        return embeddings[:, : self.output_dim]
+
 
 class ConanWrapper(Wrapper):
     def __init__(self):
         super().__init__()
-        self.model = SentenceTransformer('TencentBAC/Conan-embedding-v1', trust_remote_code=True)
-        adaptor = MixtureOfAdaptors(5, 1792)
+        self.model = SentenceTransformer(
+            "TencentBAC/Conan-embedding-v1", trust_remote_code=True
+        )
 
         model_path = snapshot_download(
             repo_id="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
-            local_dir="./Ops-MoA-Conan-embedding-v1"
+            local_dir="./Ops_MoA_Conan",
         )
-        adaptor.load_state_dict(torch.load(f"Ops-MoA-Conan-embedding-v1/conan-adaptors.pth"))
+        from Ops_MoA_Conan.modeling_adaptor import MixtureOfAdaptors
 
-        self.model.add_module('adaptor', adaptor)
+        adaptor = MixtureOfAdaptors(5, 1792)
+
+        adaptor.load_state_dict(torch.load(f"Ops_MoA_Conan/conan-adaptors.pth"))
+
+        self.model.add_module("adaptor", adaptor)
         self.output_dim = 1536
 
-    def encode(
-         self,
-         sentences: list[str],
-         **kwargs
-    ) -> np.ndarray:
+    def encode(self, sentences: list[str], **kwargs) -> np.ndarray:
         embeddings = self.model.encode(sentences, **kwargs)
-        return embeddings[:, :self.output_dim]
+        return embeddings[:, : self.output_dim]
+
 
 ops_moa_conan_embedding = ModelMeta(
     name="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
path: mteb/models/ops_moa_models.py
user: login: Samoed
body: @ahxgw Added integration for `OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0` https://huggingface.co/OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0/discussions/1 
created_at: 2025-03-28T14:22:06Z
author_association: MEMBER
diff_hunk: @@ -8,143 +8,64 @@
 import torch.nn as nn
 from huggingface_hub import snapshot_download
 
-class MoAGate(nn.Module):
-    def __init__(self, num_adaptors, hidden_dim):
-        super().__init__()
-        self.routing_vectors = nn.Parameter(
-                torch.empty(num_adaptors, hidden_dim, dtype=torch.float32),
-                requires_grad=False
-            )
-    def forward(self, hidden_states):
-        if self.routing_vectors.device == torch.device('cpu'):
-            self.routing_vectors = self.routing_vectors.to(hidden_states.device)
-        hidden_states = hidden_states.unsqueeze(1)
-        batch_size, seq_len, hidden_dim = hidden_states.shape
-
-        hidden_states = hidden_states.view(-1, hidden_dim)
-        distances = torch.cdist(hidden_states, self.routing_vectors)
-
-        _, cluster_indices = torch.min(distances, dim=1)
-        cluster_indices = cluster_indices.view(-1, 1)
-
-        topk_indices = cluster_indices
-        topk_indices = torch.zeros_like(topk_indices, device=hidden_states.device)
-        topk_weights = torch.ones_like(topk_indices, device=hidden_states.device)
-
-        return topk_indices, topk_weights
-
-class LinearLayer(nn.Module):
-    def __init__(self, input_dim, output_dim):
-        super().__init__()
-        self.linear = nn.Linear(input_dim, output_dim)
-
-    def forward(self, x):
-        return self.linear(x)
-
-class MixtureOfAdaptors(nn.Module):
-    def __init__(self, num_adaptors, hidden_dim):
-        super().__init__()
-        self.adaptors = nn.ModuleList([
-            LinearLayer(input_dim=hidden_dim, output_dim=hidden_dim)
-            for _ in range(num_adaptors)
-        ])
-        self.gate = MoAGate(num_adaptors, hidden_dim)
-
-    def forward(self, inputs):
-        if isinstance(inputs, dict):
-            hidden_states = inputs['sentence_embedding']
-        else:
-            hidden_states = inputs
-
-        residual = hidden_states
-        original_shape = hidden_states.shape
-        topk_indices, topk_weights = self.gate(hidden_states)
-
-        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-        flat_topk_indices = topk_indices.view(-1)
-        output = self.moa_inference(hidden_states, flat_topk_indices, topk_weights.view(-1, 1)).view(*original_shape)
-
-        if isinstance(inputs, dict):
-            inputs['sentence_embedding'] = output
-            return inputs
-        return output
-
-    @torch.no_grad()
-    def moa_inference(self, x, flat_adaptor_indices, flat_adaptor_weights):
-        adaptor_cache = torch.zeros_like(x)
-        sorted_indices = flat_adaptor_indices.argsort()
-        tokens_per_adaptor = flat_adaptor_indices.bincount().cpu().numpy().cumsum(0)
-        token_indices = sorted_indices
-        for i, end_idx in enumerate(tokens_per_adaptor):
-            start_idx = 0 if i == 0 else tokens_per_adaptor[i-1]
-            if start_idx == end_idx:
-                continue
-            adaptor = self.adaptors[i]
-            adaptor_token_indices = token_indices[start_idx:end_idx]
-            adaptor_tokens = x[adaptor_token_indices]
-            adaptor_output = adaptor(adaptor_tokens)
-            adaptor_output.mul_(flat_adaptor_weights[sorted_indices[start_idx:end_idx]])
-            adaptor_cache.scatter_reduce_(
-                0,
-                adaptor_token_indices.view(-1, 1).repeat(1, x.shape[-1]),
-                adaptor_output,
-                reduce='sum'
-            )
-        return adaptor_cache
 
 class YuanWrapper(Wrapper):
     def __init__(self):
         super().__init__()
-        self.model = SentenceTransformer('IEITYuan/Yuan-embedding-1.0', trust_remote_code=True)
-        adaptor = MixtureOfAdaptors(5, 1792)
+        self.model = SentenceTransformer(
+            "IEITYuan/Yuan-embedding-1.0", trust_remote_code=True
+        )
 
         model_path = snapshot_download(
             repo_id="OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0",
-            local_dir="./Ops-MoA-Yuan-embedding-1.0"
+            local_dir="./Ops_MoA_Yuan/",
         )
-        adaptor.load_state_dict(torch.load(f"Ops-MoA-Yuan-embedding-1.0/yuan-adaptors.pth"))
+        from Ops_MoA_Yuan.modeling_adaptor import MixtureOfAdaptors
+
+        adaptor = MixtureOfAdaptors(5, 1792)
 
-        self.model.add_module('adaptor', adaptor)
+        adaptor.load_state_dict(torch.load(f"Ops_MoA_Yuan/yuan-adaptors.pth"))
+
+        self.model.add_module("adaptor", adaptor)
         self.output_dim = 1536
 
-    def encode(
-         self,
-         sentences: list[str],
-         **kwargs
-    ) -> np.ndarray:
+    def encode(self, sentences: list[str], **kwargs) -> np.ndarray:
         embeddings = self.model.encode(sentences, **kwargs)
-        return embeddings[:, :self.output_dim]
+        return embeddings[:, : self.output_dim]
+
 
 class ConanWrapper(Wrapper):
     def __init__(self):
         super().__init__()
-        self.model = SentenceTransformer('TencentBAC/Conan-embedding-v1', trust_remote_code=True)
-        adaptor = MixtureOfAdaptors(5, 1792)
+        self.model = SentenceTransformer(
+            "TencentBAC/Conan-embedding-v1", trust_remote_code=True
+        )
 
         model_path = snapshot_download(
             repo_id="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
-            local_dir="./Ops-MoA-Conan-embedding-v1"
+            local_dir="./Ops_MoA_Conan",
         )
-        adaptor.load_state_dict(torch.load(f"Ops-MoA-Conan-embedding-v1/conan-adaptors.pth"))
+        from Ops_MoA_Conan.modeling_adaptor import MixtureOfAdaptors
 
-        self.model.add_module('adaptor', adaptor)
+        adaptor = MixtureOfAdaptors(5, 1792)
+
+        adaptor.load_state_dict(torch.load(f"Ops_MoA_Conan/conan-adaptors.pth"))
+
+        self.model.add_module("adaptor", adaptor)
         self.output_dim = 1536
 
-    def encode(
-         self,
-         sentences: list[str],
-         **kwargs
-    ) -> np.ndarray:
+    def encode(self, sentences: list[str], **kwargs) -> np.ndarray:
         embeddings = self.model.encode(sentences, **kwargs)
-        return embeddings[:, :self.output_dim]
+        return embeddings[:, : self.output_dim]
+
 
 ops_moa_conan_embedding = ModelMeta(
     name="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
path: mteb/models/ops_moa_models.py
user: login: ahxgw
body: Thanks very much. We have merged the PR and update another model.
created_at: 2025-03-28T15:13:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from mteb.model_meta import ModelMeta
+from mteb.models.wrapper import Wrapper
+from functools import partial
+from sentence_transformers import SentenceTransformer
+import torch
+import torch.nn as nn
+from huggingface_hub import snapshot_download
+
+
+class CustomWrapper(Wrapper):
+    def __init__(self, model_name, revision):
+        super().__init__()
+        self.model = SentenceTransformer(
+            model_name, revision=revision, trust_remote_code=True
+        )
+        self.output_dim = 1536
+
+    def encode(self, sentences: list[str], **kwargs) -> np.ndarray:
+        embeddings = self.model.encode(sentences, **kwargs)
+        return embeddings[:, : self.output_dim]
+
+
+ops_moa_conan_embedding = ModelMeta(
+    name="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+    revision="46dcd58753f3daa920c66f89e47086a534089350",
+    release_date="2025-03-26",
+    languages=["zho_Hans"],
+    loader=partial(
+        CustomWrapper,
+        "OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+        "46dcd58753f3daa920c66f89e47086a534089350",
+    ),
+    n_parameters=343 * 1e6,
+    memory_usage_mb=2e3,
+    max_tokens=512,
+    embed_dim=1536,
+    license="cc-by-nc-4.0",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=["PyTorch"],
+    reference="https://huggingface.co/OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+    similarity_fn_name="cosine",
+    use_instructions=False,
+    training_datasets=None,
+    superseded_by=None,
+)
+
+ops_moa_yuan_embedding = ModelMeta(
+    name="OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0",
+    revision="23712d0766417b0eb88a2513c6e212a58b543268",
+    release_date="2025-03-26",
+    languages=["zho_Hans"],
+    loader=partial(
+        CustomWrapper,
+        "OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0",
+        "23712d0766417b0eb88a2513c6e212a58b543268",
+    ),
+    n_parameters=343 * 1e6,
+    memory_usage_mb=2e3,
+    max_tokens=512,
+    embed_dim=1536,
+    license="cc-by-nc-4.0",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=["PyTorch"],
+    reference="https://huggingface.co/OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0",
+    similarity_fn_name="cosine",
+    use_instructions=False,
+    training_datasets=None,
path: mteb/models/ops_moa_models.py
user: login: Samoed
body: Can you add traning datasets of your model and base model?
created_at: 2025-03-28T18:29:59Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from mteb.model_meta import ModelMeta
+from mteb.models.wrapper import Wrapper
+from functools import partial
+from sentence_transformers import SentenceTransformer
+import torch
+import torch.nn as nn
+from huggingface_hub import snapshot_download
+
+
+class CustomWrapper(Wrapper):
+    def __init__(self, model_name, revision):
+        super().__init__()
+        self.model = SentenceTransformer(
+            model_name, revision=revision, trust_remote_code=True
+        )
+        self.output_dim = 1536
+
+    def encode(self, sentences: list[str], **kwargs) -> np.ndarray:
+        embeddings = self.model.encode(sentences, **kwargs)
+        return embeddings[:, : self.output_dim]
+
+
+ops_moa_conan_embedding = ModelMeta(
+    name="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+    revision="46dcd58753f3daa920c66f89e47086a534089350",
+    release_date="2025-03-26",
+    languages=["zho_Hans"],
+    loader=partial(
+        CustomWrapper,
+        "OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+        "46dcd58753f3daa920c66f89e47086a534089350",
+    ),
+    n_parameters=343 * 1e6,
+    memory_usage_mb=2e3,
+    max_tokens=512,
+    embed_dim=1536,
+    license="cc-by-nc-4.0",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=["PyTorch"],
+    reference="https://huggingface.co/OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+    similarity_fn_name="cosine",
+    use_instructions=False,
+    training_datasets=None,
+    superseded_by=None,
+)
+
+ops_moa_yuan_embedding = ModelMeta(
+    name="OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0",
+    revision="23712d0766417b0eb88a2513c6e212a58b543268",
+    release_date="2025-03-26",
+    languages=["zho_Hans"],
+    loader=partial(
+        CustomWrapper,
+        "OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0",
+        "23712d0766417b0eb88a2513c6e212a58b543268",
+    ),
+    n_parameters=343 * 1e6,
+    memory_usage_mb=2e3,
+    max_tokens=512,
+    embed_dim=1536,
+    license="cc-by-nc-4.0",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=["PyTorch"],
+    reference="https://huggingface.co/OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0",
+    similarity_fn_name="cosine",
+    use_instructions=False,
+    training_datasets=None,
path: mteb/models/ops_moa_models.py
user: login: ahxgw
body: OK, we have added.
created_at: 2025-03-29T01:24:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from mteb.model_meta import ModelMeta
+from mteb.models.wrapper import Wrapper
+from functools import partial
+from sentence_transformers import SentenceTransformer
+import torch
+import torch.nn as nn
+from huggingface_hub import snapshot_download
+
+
+class CustomWrapper(Wrapper):
+    def __init__(self, model_name, revision):
+        super().__init__()
+        self.model = SentenceTransformer(
+            model_name, revision=revision, trust_remote_code=True
+        )
+        self.output_dim = 1536
+
+    def encode(self, sentences: list[str], **kwargs) -> np.ndarray:
+        embeddings = self.model.encode(sentences, **kwargs)
+        return embeddings[:, : self.output_dim]
+
+
+ops_moa_conan_embedding = ModelMeta(
+    name="OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+    revision="46dcd58753f3daa920c66f89e47086a534089350",
+    release_date="2025-03-26",
+    languages=["zho_Hans"],
+    loader=partial(
+        CustomWrapper,
+        "OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+        "46dcd58753f3daa920c66f89e47086a534089350",
+    ),
+    n_parameters=343 * 1e6,
+    memory_usage_mb=2e3,
+    max_tokens=512,
+    embed_dim=1536,
+    license="cc-by-nc-4.0",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=["PyTorch", "Sentence Transformers"],
+    reference="https://huggingface.co/OpenSearch-AI/Ops-MoA-Conan-embedding-v1",
+    similarity_fn_name="cosine",
+    use_instructions=False,
+    training_datasets={"cmteb_retrieval": ["train"]},
path: mteb/models/ops_moa_models.py
user: login: Samoed
body: ```suggestion
    training_datasets={
                "T2Retrieval": ["train"],
                "MMarcoRetrieval": ["train"],
                "DuRetrieval": ["train"],
                "CovidRetrieval": ["train"],
                "CmedqaRetrieval": ["train"],
                "EcomRetrieval": ["train"],
                "MedicalRetrieval": ["train"],
                "VideoRetrieval": ["train"],
    },
```
created_at: 2025-03-29T10:35:52Z
author_association: MEMBER

PR comments:
user: login: ahxgw
created_at: 2025-03-27T02:29:34Z
body: We have added the model implementation and other information, except the dataset (which is not yet public), thanks
author_association: CONTRIBUTOR
user: login: Samoed
created_at: 2025-03-27T05:05:40Z
body: I don't like to have much code for this model to run. I'll add custom implementation to your model. I'll create PR on HF
author_association: MEMBER
user: login: ahxgw
created_at: 2025-03-27T06:13:44Z
body: OK, thx
author_association: CONTRIBUTOR