Upload 3 files

README.md

@@ -1,3 +1,121 @@
----
-license: apache-2.0
----
+---
+language:
+- zh
+base_model: OpenSearch-AI/Ops-MoA-Yuan-embedding-1.0
+model-index:
+- name: Ops-MoA-Yuan-embedding-1.0
+  results:
+  - task:
+      type: Retrieval
+    dataset:
+      type: C-MTEB/CmedqaRetrieval
+      name: MTEB CmedqaRetrieval
+      config: default
+      split: dev
+      revision: cd540c506dae1cf9e9a59c3e06f42030d54e7301
+    metrics:
+    - type: ndcg_at_10
+      value: 51.46
+  - task:
+      type: Retrieval
+    dataset:
+      type: C-MTEB/CovidRetrieval
+      name: MTEB CovidRetrieval
+      config: default
+      split: dev
+      revision: 1271c7809071a13532e05f25fb53511ffce77117
+    metrics:
+    - type: ndcg_at_10
+      value: 93.2
+  - task:
+      type: Retrieval
+    dataset:
+      type: C-MTEB/DuRetrieval
+      name: MTEB DuRetrieval
+      config: default
+      split: dev
+      revision: a1a333e290fe30b10f3f56498e3a0d911a693ced
+    metrics:
+    - type: ndcg_at_10
+      value: 89.84
+  - task:
+      type: Retrieval
+    dataset:
+      type: C-MTEB/EcomRetrieval
+      name: MTEB EcomRetrieval
+      config: default
+      split: dev
+      revision: 687de13dc7294d6fd9be10c6945f9e8fec8166b9
+    metrics:
+    - type: ndcg_at_10
+      value: 71.08
+  - task:
+      type: Retrieval
+    dataset:
+      type: C-MTEB/MMarcoRetrieval
+      name: MTEB MMarcoRetrieval
+      config: default
+      split: dev
+      revision: 539bbde593d947e2a124ba72651aafc09eb33fc2
+    metrics:
+    - type: ndcg_at_10
+      value: 79.27
+  - task:
+      type: Retrieval
+    dataset:
+      type: C-MTEB/MedicalRetrieval
+      name: MTEB MedicalRetrieval
+      config: default
+      split: dev
+      revision: 2039188fb5800a9803ba5048df7b76e6fb151fc6
+    metrics:
+    - type: ndcg_at_10
+      value: 74.84
+  - task:
+      type: Retrieval
+    dataset:
+      type: C-MTEB/T2Retrieval
+      name: MTEB T2Retrieval
+      config: default
+      split: dev
+      revision: 8731a845f1bf500a4f111cf1070785c793d10e64
+    metrics:
+    - type: ndcg_at_10
+      value: 85.78
+  - task:
+      type: Retrieval
+    dataset:
+      type: C-MTEB/VideoRetrieval
+      name: MTEB VideoRetrieval
+      config: default
+      split: dev
+      revision: 58c2597a5943a2ba48f4668c3b90d796283c5639
+    metrics:
+    - type: ndcg_at_10
+      value: 79.51
+pipeline_tag: feature-extraction
+tags:
+- mteb
+- sentence-transformers
+library_name: transformers
+---
+
+```python
+import torch.nn as nn
+from sentence_transformers import SentenceTransformer
+from modeling_adaptor import MixtureOfAdaptors
+class CustomSentenceTransformer(nn.Module):
+    def __init__(self, output_dim=1536):
+        super(CustomSentenceTransformer, self).__init__()
+        self.model = SentenceTransformer('IEITYuan/Yuan-embedding-1.0', trust_remote_code=True)
+        adaptor = MixtureOfAdaptors(5, 1792)
+        adaptor.load_state_dict(torch.load(f"yuan-adaptors.pth"))
+        self.model.add_module('adaptor', adaptor)
+        self.output_dim = output_dim
+    
+    def encode(self, sentences, **kwargs):
+        embeddings = self.model.encode(sentences, **kwargs)
+        return embeddings[:, :self.output_dim]
+
+model = CustomSentenceTransformer(output_dim=1536)
+model.encode(['text'])

modeling_adaptor.py

@@ -0,0 +1,91 @@
+import torch
+import torch.nn as nn
+import math
+import torch.nn.functional as F
+import joblib
+
+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
+
+

yuan-adaptors.pth

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