first commit

1_Pooling/config.json

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

README.md

@@ -1,3 +1,247 @@
-# gte-multilingual-base_a13715200061861888459946
+---
+tasks:
+- sentence-embedding
+widgets:
+- version: v1.1.0
+  task: sentence-embedding
+  inputs:
+      - type: text
+        name: source_sentence
+      - type: text-list
+        name: sentences_to_compare
+  examples:
+      - name: 示例1
+        inputs:
+          - data:
+            - 功和功率的区别
+          - data: 
+            - 功反映做功多少，功率反映做功快慢。
+            - 什么是有功功率和无功功率?无功功率有什么用什么是有功功率和无功功率?无功功率有什么用电力系统中的电源是由发电机产生的三相正弦交流电,在交>流电路中,由电源供给负载的电功率有两种;一种是有功功率,一种是无功功率.
+            - 优质解答在物理学中,用电功率表示消耗电能的快慢．电功率用P表示,它的单位是瓦特（Watt）,简称瓦（Wa）符号是W.电流在单位时间内做的功叫做电功率 以灯泡为例,电功率越大,灯泡越亮.灯泡的亮暗由电功率（实际功率）决定,不由通过的电流、电压、电能决定!
+      - name: 示例2
+        inputs:
+          - data:
+            - 什么是桥
+          - data:
+            - 由全国科学技术名词审定委员会审定的科技名词“桥”的定义为：跨越河流、山谷、障碍物或其他交通线而修建的架空通道
+            - 桥是一种连接两岸的建筑
+            - 转向桥，是指承担转向任务的车桥。一般汽车的前桥是转向桥。四轮转向汽车的前后桥，都是转向桥。
+            - 桥梁艺术研究桥梁美学效果的艺术。起源于人类修建桥梁的实践。作为一种建筑，由于功能不同、使用材料有差别，桥梁表现为不同的结构和形式。基本的有拱桥、梁桥和吊桥。
+      - name: 示例3
+        inputs:
+          - data:
+            - 福鼎在哪个市
+          - data:
+            - 福鼎是福建省宁德市福鼎市。
+            - 福鼎位于福建省东北部，东南濒东海，水系发达，岛屿星罗棋布。除此之外，福鼎地貌丰富，著名的太姥山就在福鼎辖内，以峰险、石奇、洞幽、雾幻四绝著称于世。
+            - 福鼎市人民政府真诚的欢迎国内外朋友前往迷人的太姥山观光、旅游。
+            - 福建省福鼎市桐山街道,地处福鼎城区,地理环境优越,三面环山,东临大海,境内水陆交通方便,是福鼎政治、经济、文化、交通中心区。
+      - name: 示例4
+        inputs:
+          - data:
+            - 吃完海鲜可以喝牛奶吗?
+          - data: 
+            - 不可以，早晨喝牛奶不科学
+            - 吃了海鲜后是不能再喝牛奶的，因为牛奶中含得有维生素C，如果海鲜喝牛奶一起服用会对人体造成一定的伤害
+            - 吃海鲜是不能同时喝牛奶吃水果，这个至少间隔6小时以上才可以。
+            - 吃海鲜是不可以吃柠檬的因为其中的维生素C会和海鲜中的矿物质形成砷
+model-type:
+- bert
 
-gte-multilingual-base
+domain:
+- nlp
+
+frameworks:
+- pytorch
+
+backbone:
+- transformer
+
+metrics:
+- mrr@10
+- recall@1000
+- ndcg@10
+
+license: Apache License 2.0
+
+language:
+- multilingual
+
+tags:
+- text representation
+- text retrieval
+- passage retrieval
+- Transformer
+- GTE
+- 多语言
+- 长文本
+- 文本相关性
+- 文本相似度
+
+---
+
+# GTE多语言通用文本表示模型
+
+文本表示是自然语言处理(NLP)领域的核心问题, 其在很多NLP、信息检索的下游任务中发挥着非常重要的作用。近几年, 随着深度学习的发展，尤其是预训练语言模型的出现极大的推动了文本表示技术的效果, 基于预训练语言模型的文本表示模型在学术研究数据、工业实际应用中都明显优于传统的基于统计模型或者浅层神经网络的文本表示模型。这里, 我们主要关注基于预训练语言模型的文本表示。
+
+文本表示示例, 输入一个句子, 输入一个固定维度的连续向量:
+
+- 输入: 吃完海鲜可以喝牛奶吗?
+- 输出: [0.27162,-0.66159,0.33031,0.24121,0.46122,...]
+
+文本的向量表示通常可以用于文本聚类、文本相似度计算、文本向量召回等下游任务中。
+
+## GTE多语言通用文本表示模型
+
+gte-multilingual-base模型是GTE(通用文本向量，General Text Embedding) 系列模型中的最新的多语言模型，该模型具有几个关键特性：
+
+ - 高性能: 在多语言检索任务和多任务表示模型评估中，相比同规模模型，达到了最先进的（SOTA）结果。
+ - 高推理效率：采用Encoder-Only架构进行训练，使模型体积更小(300M)。LLM架构的向量模型(例如，gte-qwen2-7b-instruct)不同，该模型在推理时对硬件的要求更低，推理速度提高了10倍
+ - 长上下文：支持最多8192个token的文本长度
+ - 多语言能力：支持70多种语言
+ - 弹性向量维度：支持弹性的向量表示，同时保持下游任务的有效性，显著降低存储成本并提高了检索效率。
+ - 稀疏向量：除了连续向量表示外，还可以生成稀疏向量表示
+
+模型训练更多细节可以参考论文: [mGTE: Generalized Long-Context Text Representation and Reranking Models for Multilingual Text Retrieval](https://arxiv.org/abs/2407.19669)
+
+### 如何使用
+```
+import torch.nn.functional as F
+from modelscope import AutoModel, AutoTokenizer
+
+input_texts = [
+    "what is the capital of China?",
+    "how to implement quick sort in python?",
+    "北京",
+    "快排算法介绍"
+]
+
+model_name_or_path = 'iic/gte_sentence-embedding_multilingual-base'
+tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+model = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=True)
+
+# Tokenize the input texts
+batch_dict = tokenizer(input_texts, max_length=8192, padding=True, truncation=True, return_tensors='pt')
+
+outputs = model(**batch_dict)
+
+dimension=768 # The output dimension of the output embedding, should be in [128, 768]
+embeddings = outputs.last_hidden_state[:, 0][:dimension]
+
+embeddings = F.normalize(embeddings, p=2, dim=1)
+scores = (embeddings[:1] @ embeddings[1:].T)
+print(scores.tolist())
+
+# [[0.3016996383666992, 0.7503870129585266, 0.3203084468841553]]
+```
+
+- 使用sentence-transformers进行推理
+```
+# Requires sentences-transformers>=3.0.0
+
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.util import cos_sim
+import numpy as np
+
+input_texts = [
+    "what is the capital of China?",
+    "how to implement quick sort in python?",
+    "北京",
+    "快排算法介绍"
+]
+
+model_name_or_path="Alibaba-NLP/gte-multilingual-base"
+model = SentenceTransformer(', trust_remote_code=True)
+embeddings = model.encode(input_texts) # embeddings.shape (4, 768)
+
+# normalized embeddings
+norms = np.linalg.norm(embeddings, ord=2, axis=1, keepdims=True)
+norms[norms == 0] = 1
+embeddings = embeddings / norms
+
+# sim scores
+scores = (embeddings[:1] @ embeddings[1:].T)
+
+print(scores.tolist())
+# [[0.301699697971344, 0.7503870129585266, 0.32030850648880005]]
+```
+
+- 使用自定义GTE-Embedding脚本获取连续向量表示和离散向量表示
+```gte_embedding.py```文件在```scripts/gte_embedding.py```中
+
+```
+from gte_embedding import GTEEmbeddidng
+
+model_name_or_path = 'gte_sentence-embedding_multilingual-base'
+model = GTEEmbeddidng(model_name_or_path)
+query = "中国的首都在哪儿"
+
+docs = [
+    "what is the capital of China?",
+    "how to implement quick sort in python?",
+    "北京",
+    "快排算法介绍"
+]
+
+embs = model.encode(docs, return_dense=True,return_sparse=True)
+print('dense_embeddings vecs', embs['dense_embeddings'])
+print('token_weights', embs['token_weights'])
+pairs = [(query, doc) for doc in docs]
+dense_scores = model.compute_scores(pairs, dense_weight=1.0, sparse_weight=0.0)
+sparse_scores = model.compute_scores(pairs, dense_weight=0.0, sparse_weight=1.0)
+hybird_scores = model.compute_scores(pairs, dense_weight=1.0, sparse_weight=0.3)
+
+print('dense_scores', dense_scores)
+print('sparse_scores', sparse_scores)
+print('hybird_scores', hybird_scores)
+
+#dense_scores [0.85302734375, 0.257568359375, 0.76953125, 0.325439453125]
+#sparse_scores [0.0, 0.0, 4.600879669189453, 1.570279598236084]
+#hybird_scores [0.85302734375, 0.257568359375, 2.1497951507568356, 0.7965233325958252]
+```
+
+## 模型效果评估
+
+我们在多个下游任务上验证了GTE多语言表征模型的效果，包括多语言检索、跨语言检索、长文本检索，以及在[MTEB](https://huggingface.co/spaces/mteb/leaderboard)多任务文本表征评测
+
+- **检索任务**
+
+Retrieval results on [MIRACL](https://arxiv.org/abs/2210.09984) and [MLDR](https://arxiv.org/abs/2402.03216) (multilingual), [MKQA](https://arxiv.org/abs/2007.15207) (crosslingual), [BEIR](https://arxiv.org/abs/2104.08663) and [LoCo](https://arxiv.org/abs/2402.07440) (English)
+
+![image](./images/mgte-retrieval.png)
+
+- **MTEB多语言**
+
+Results on MTEB English, Chinese, French, Polish
+
+![image](./images/mgte-mteb.png)
+
+更多实验结果可以参考论文[paper](https://arxiv.org/pdf/2407.19669).
+
+## API云服务
+
+除了开源的 GTE 系列模型，GTE系列模型同时在阿里云上提供商用API服务:
+
+- [文本Embedding模型](https://help.aliyun.com/zh/model-studio/developer-reference/general-text-embedding/): 提供三种版本的文本嵌入模型：text-embedding-v1/v2/v3，其中v3是最新版本的模型服务
+- [文本ReRank模型](https://help.aliyun.com/zh/model-studio/developer-reference/general-text-sorting-model/): 提供gte-rerank模型服务, 模型持续迭代中
+
+Note: API服务背后的模型与开源模型并不完全相同
+
+
+
+## 引用
+
+```BibTeX
+@article{zhang2024mgte,
+  title={mGTE: Generalized Long-Context Text Representation and Reranking Models for Multilingual Text Retrieval},
+  author={Zhang, Xin and Zhang, Yanzhao and Long, Dingkun and Xie, Wen and Dai, Ziqi and Tang, Jialong and Lin, Huan and Yang, Baosong and Xie, Pengjun and Huang, Fei and others},
+  journal={arXiv preprint arXiv:2407.19669},
+  year={2024}
+}
+@article{li2023towards,
+  title={Towards general text embeddings with multi-stage contrastive learning},
+  author={Li, Zehan and Zhang, Xin and Zhang, Yanzhao and Long, Dingkun and Xie, Pengjun and Zhang, Meishan},
+  journal={arXiv preprint arXiv:2308.03281},
+  year={2023}
+}
+```

config.json

@@ -0,0 +1,43 @@
+{
+  "architectures": [
+    "NewModel",
+    "NewForTokenClassification"
+  ],
+  "attention_probs_dropout_prob": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration.NewConfig",
+    "AutoModelForMaskedLM": "modeling.NewForMaskedLM",
+    "AutoModel": "modeling.NewModel",
+    "AutoModelForMultipleChoice": "modeling.NewForMultipleChoice",
+    "AutoModelForQuestionAnswering": "modeling.NewForQuestionAnswering",
+    "AutoModelForSequenceClassification": "modeling.NewForSequenceClassification",
+    "AutoModelForTokenClassification": "modeling.NewForTokenClassification"
+  },
+  "classifier_dropout": 0.0,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-12,
+  "layer_norm_type": "layer_norm",
+  "max_position_embeddings": 8192,
+  "model_type": "new",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "num_labels": 1,
+  "pack_qkv": true,
+  "pad_token_id": 1,
+  "position_embedding_type": "rope",
+  "rope_scaling": {
+    "factor": 8.0,
+    "type": "ntk"
+  },
+  "rope_theta": 20000,
+  "torch_dtype": "float16",
+  "transformers_version": "4.39.1",
+  "type_vocab_size": 1,
+  "unpad_inputs": false,
+  "use_memory_efficient_attention": false,
+  "vocab_size": 250048
+}

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"sentence-embedding"}

configuration.py

@@ -0,0 +1,145 @@
+# coding=utf-8
+# Copyright 2024 The GTE Team Authors and Alibaba Group.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+""" NEW model configuration"""
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class NewConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`NewModel`] or a [`TFNewModel`]. It is used to
+    instantiate a NEW model according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the NEW
+    [izhx/new-base-en](https://huggingface.co/izhx/new-base-en) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 30522):
+            Vocabulary size of the NEW model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`NewModel`] or [`TFNewModel`].
+        hidden_size (`int`, *optional*, defaults to 768):
+            Dimensionality of the encoder layers and the pooler layer.
+        num_hidden_layers (`int`, *optional*, defaults to 12):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 12):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        intermediate_size (`int`, *optional*, defaults to 3072):
+            Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
+        hidden_act (`str` or `Callable`, *optional*, defaults to `"gelu"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"silu"` and `"gelu_new"` are supported.
+        hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
+            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
+        attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
+            The dropout ratio for the attention probabilities.
+        max_position_embeddings (`int`, *optional*, defaults to 512):
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        type_vocab_size (`int`, *optional*, defaults to 2):
+            The vocabulary size of the `token_type_ids` passed when calling [`NewModel`] or [`TFNewModel`].
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-12):
+            The epsilon used by the layer normalization layers.
+        position_embedding_type (`str`, *optional*, defaults to `"rope"`):
+            Type of position embedding. Choose one of `"absolute"`, `"rope"`.
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
+            strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
+            `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
+            `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
+            these scaling strategies behave:
+            https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
+            experimental feature, subject to breaking API changes in future versions.
+        classifier_dropout (`float`, *optional*):
+            The dropout ratio for the classification head.
+
+    Examples:
+
+    ```python
+    >>> from transformers import NewConfig, NewModel
+
+    >>> # Initializing a NEW izhx/new-base-en style configuration
+    >>> configuration = NewConfig()
+
+    >>> # Initializing a model (with random weights) from the izhx/new-base-en style configuration
+    >>> model = NewModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "new"
+
+    def __init__(
+        self,
+        vocab_size=30528,
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        intermediate_size=3072,
+        hidden_act="gelu",
+        hidden_dropout_prob=0.1,
+        attention_probs_dropout_prob=0.0,
+        max_position_embeddings=2048,
+        type_vocab_size=1,
+        initializer_range=0.02,
+        layer_norm_type='layer_norm',
+        layer_norm_eps=1e-12,
+        # pad_token_id=0,
+        position_embedding_type="rope",
+        rope_theta=10000.0,
+        rope_scaling=None,
+        classifier_dropout=None,
+        pack_qkv=True,
+        unpad_inputs=False,
+        use_memory_efficient_attention=False,
+        logn_attention_scale=False,
+        logn_attention_clip1=False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.type_vocab_size = type_vocab_size
+        self.initializer_range = initializer_range
+        self.layer_norm_type = layer_norm_type
+        self.layer_norm_eps = layer_norm_eps
+        self.position_embedding_type = position_embedding_type
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.classifier_dropout = classifier_dropout
+
+        self.pack_qkv = pack_qkv
+        self.unpad_inputs = unpad_inputs
+        self.use_memory_efficient_attention = use_memory_efficient_attention
+        self.logn_attention_scale = logn_attention_scale
+        self.logn_attention_clip1 = logn_attention_clip1

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"
+  }
+]

scripts/gte_embedding.py

@@ -0,0 +1,154 @@
+# coding=utf-8
+# Copyright 2024 The GTE Team Authors and Alibaba Group.
+# Licensed under the Apache License, Version 2.0 (the "License");
+
+from collections import defaultdict
+from typing import Dict, List, Tuple
+
+import numpy as np
+import torch
+from transformers import AutoModelForTokenClassification, AutoTokenizer
+from transformers.utils import is_torch_npu_available
+
+
+class GTEEmbeddidng(torch.nn.Module):
+    def __init__(self,
+                 model_name: str = None,
+                 normalized: bool = True,
+                 use_fp16: bool = True,
+                 device: str = None
+                ):
+        super().__init__()
+        self.normalized = normalized
+        if device:
+            self.device = torch.device(device)
+        else:
+            if torch.cuda.is_available():
+                self.device = torch.device("cuda")
+            elif torch.backends.mps.is_available():
+                self.device = torch.device("mps")
+            elif is_torch_npu_available():
+                self.device = torch.device("npu")
+            else:
+                self.device = torch.device("cpu")
+                use_fp16 = False
+        self.use_fp16 = use_fp16
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForTokenClassification.from_pretrained(
+            model_name, trust_remote_code=True, torch_dtype=torch.float16 if self.use_fp16 else None
+        )
+        self.vocab_size = self.model.config.vocab_size
+        self.model.to(self.device)
+
+    def _process_token_weights(self, token_weights: np.ndarray, input_ids: list):
+        # conver to dict
+        result = defaultdict(int)
+        unused_tokens = set([self.tokenizer.cls_token_id, self.tokenizer.eos_token_id, self.tokenizer.pad_token_id,
+                             self.tokenizer.unk_token_id])
+        # token_weights = np.ceil(token_weights * 100)
+        for w, idx in zip(token_weights, input_ids):
+            if idx not in unused_tokens and w > 0:
+                token = self.tokenizer.decode([int(idx)])
+                if w > result[token]:
+                    result[token] = w
+        return result
+
+    @torch.no_grad()
+    def encode(self,
+               texts: None,
+               dimension: int = None,
+               max_length: int = 8192,
+               batch_size: int = 16,
+               return_dense: bool = True,
+               return_sparse: bool = False):
+        if dimension is None:
+            dimension = self.model.config.hidden_size
+        if isinstance(texts, str):
+            texts = [texts]
+        num_texts = len(texts)
+        all_dense_vecs = []
+        all_token_weights = []
+        for n, i in enumerate(range(0, num_texts, batch_size)):
+            batch = texts[i: i + batch_size]
+            resulst = self._encode(batch, dimension, max_length, batch_size, return_dense, return_sparse)
+            if return_dense:
+                all_dense_vecs.append(resulst['dense_embeddings'])
+            if return_sparse:
+                all_token_weights.extend(resulst['token_weights'])
+        all_dense_vecs = torch.cat(all_dense_vecs, dim=0)
+        return {
+            "dense_embeddings": all_dense_vecs,
+            "token_weights": all_token_weights 
+        }
+
+    @torch.no_grad()
+    def _encode(self,
+                texts: Dict[str, torch.Tensor] = None,
+                dimension: int = None,
+                max_length: int = 1024,
+                batch_size: int = 16,
+                return_dense: bool = True,
+                return_sparse: bool = False):
+
+        text_input = self.tokenizer(texts, padding=True, truncation=True, return_tensors='pt', max_length=max_length)
+        text_input = {k: v.to(self.model.device) for k,v in text_input.items()}
+        model_out = self.model(**text_input, return_dict=True)
+
+        output = {}
+        if return_dense:
+            dense_vecs = model_out.last_hidden_state[:, 0, :dimension]
+            if self.normalized:
+                dense_vecs = torch.nn.functional.normalize(dense_vecs, dim=-1)
+            output['dense_embeddings'] = dense_vecs
+        if return_sparse:
+            token_weights = torch.relu(model_out.logits).squeeze(-1)
+            token_weights = list(map(self._process_token_weights, token_weights.detach().cpu().numpy().tolist(),
+                                                    text_input['input_ids'].cpu().numpy().tolist()))
+            output['token_weights'] = token_weights
+
+        return output
+
+    def _compute_sparse_scores(self, embs1, embs2):
+        scores = 0
+        for token, weight in embs1.items():
+            if token in embs2:
+                scores += weight * embs2[token]
+        return scores
+
+    def compute_sparse_scores(self, embs1, embs2):
+        scores = [self._compute_sparse_scores(emb1, emb2) for emb1, emb2 in zip(embs1, embs2)]
+        return np.array(scores)
+
+    def compute_dense_scores(self, embs1, embs2):
+        scores = torch.sum(embs1*embs2, dim=-1).cpu().detach().numpy()
+        return scores
+
+    @torch.no_grad()
+    def compute_scores(self, 
+        text_pairs: List[Tuple[str, str]], 
+        dimension: int = None,
+        max_length: int = 1024,
+        batch_size: int = 16,
+        dense_weight=1.0,
+        sparse_weight=0.1):
+        text1_list = [text_pair[0] for text_pair in text_pairs]
+        text2_list = [text_pair[1] for text_pair in text_pairs]
+        embs1 = self.encode(text1_list, dimension, max_length, batch_size, return_dense=True, return_sparse=True)
+        embs2 = self.encode(text2_list, dimension, max_length, batch_size, return_dense=True, return_sparse=True)
+        scores = self.compute_dense_scores(embs1['dense_embeddings'], embs2['dense_embeddings']) * dense_weight + \
+            self.compute_sparse_scores(embs1['token_weights'], embs2['token_weights']) * sparse_weight
+        scores = scores.tolist()
+        return scores
+
+
+if __name__ == '__main__':
+    gte = GTEEmbeddidng('Alibaba-NLP/gte-multilingual-base')
+    docs =  [
+        "黑龙江离俄罗斯很近",
+        "哈尔滨是中国黑龙江省的省会，位于中国东北",
+        "you are the hero"
+    ]
+    print('docs', docs)
+    embs = gte.encode(docs, return_dense=True,return_sparse=True)
+    print('dense vecs', embs['dense_embeddings'])
+    print('sparse vecs', embs['token_weights'])

sentence_bert_config.json

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

special_tokens_map.json

@@ -0,0 +1,51 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "cls_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "<mask>",
+    "lstrip": true,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,54 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "250001": {
+      "content": "<mask>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "mask_token": "<mask>",
+  "model_max_length": 32768,
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "tokenizer_class": "XLMRobertaTokenizer",
+  "unk_token": "<unk>"
+}