Upload 8 files

README.md

@@ -0,0 +1,93 @@
+---
+language:
+  - zh
+base_model: junnyu/roformer_chinese_base
+tags:
+  - transformers
+---
+
+## <u>INF</u> <u>W</u>ord-level <u>S</u>parse <u>E</u>mbedding v2 (INF-WSE-v2)
+
+**INF-WSE-v2** is the latest version of the word-level sparse embedding model developed by [INF TECH](https://www.infly.cn/en). 
+
+Compared to [INF-WSE-v1](https://huggingface.co/infly/inf-wse-v1-base-zh), INF-WSE-v2 continues to be pretrained on the [Wudao](https://huggingface.co/datasets/p208p2002/wudao) corpus (from [roformer_chinese_base](https://huggingface.co/junnyu/roformer_chinese_base)) and introduces enhanced token rewriting capabilities. These advancements improve the model's ability to generate more accurate, adaptable, and contextually relevant text embeddings, with a particular focus on Chinese language processing.
+
+### Key Features:
+
+- **Optimized for Retrieval**: INF-WSE-v2 is specifically designed for information retrieval tasks. By leveraging sparse embeddings, the model ensures efficient matching between queries and documents, making it ideal for semantic search, ranking, and other retrieval scenarios where both speed and accuracy are essential.
+
+- **Token Rewriting Capability**: A new token rewriting feature allows INF-WSE-v2 to dynamically modify tokens during the embedding process. This improves the model’s ability to produce more accurate and contextually relevant representations, especially when dealing with complex linguistic structures and nuances in Chinese text.
+
+- **Sparse Representation for Efficiency**: Unlike traditional dense embeddings, which have a fixed dimensionality, INF-WSE-v2 uses sparse embeddings where most dimensions are set to zero. This results in embeddings where only the most significant dimensions are non-zero, reducing computational load while maintaining high accuracy for retrieval tasks.
+
+## Usage
+
+### Transformers
+
+#### Infer embeddings
+```python
+import torch
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+queries = ['电脑一体机由什么构成？', '什么是掌上电脑？']
+documents = [
+    '电脑一体机，是由一台显示器、一个电脑键盘和一个鼠标组成的电脑。',
+    '掌上电脑是一种运行在嵌入式操作系统和内嵌式应用软件之上的、小巧、轻便、易带、实用、价廉的手持式计算设备。',
+]
+input_texts = queries + documents
+
+tokenizer = AutoTokenizer.from_pretrained("infly/inf-wse-v2-base-zh", trust_remote_code=True, use_fast=False)  # Fast tokenizer has not been supported yet
+model = AutoModelForMaskedLM.from_pretrained("infly/inf-wse-v2-base-zh", trust_remote_code=True)
+model.eval()
+
+max_length = 512
+
+input_batch = tokenizer(input_texts, padding=True, max_length=max_length, truncation=True, return_tensors="pt")
+with torch.no_grad():
+    embeddings = model(input_batch['input_ids'], input_batch['attention_mask'], return_sparse=False)  # if return_sparse=True, return sparse tensor, else return dense tensor
+
+scores = embeddings[:2] @ embeddings[2:].T
+print(scores.tolist())
+# [[25.137710571289062, 9.891149520874023], [11.703001976013184, 30.97362518310547]]
+```
+
+#### Convert embeddings to lexical weights
+```python
+from collections import OrderedDict
+def convert_embeddings_to_weights(embeddings, tokenizer):
+    values, indices = torch.sort(embeddings, dim=-1, descending=True)
+    
+    token2weight = []
+    for i in range(embeddings.size(0)):
+        token2weight.append(OrderedDict())
+
+        non_zero_mask = values[i] != 0
+        tokens = tokenizer.convert_ids_to_tokens(indices[i][non_zero_mask])
+        weights = values[i][non_zero_mask].tolist()
+
+        for token, weight in zip(tokens, weights):
+            token2weight[i][token] = weight
+
+    return token2weight
+
+token2weight = convert_embeddings_to_weights(embeddings, tokenizer)
+print(token2weight[1])
+# OrderedDict([('掌上', 1.9666814804077148), ('电脑', 1.4205719232559204), ('掌中', 1.2688857316970825), ('全称', 1.2548470497131348), ('to', 1.041936993598938), ('台式机', 0.9435897469520569), ('编程语言', 0.8740423917770386), ('pad', 0.8506593108177185), ('手持', 0.835372269153595), ('point', 0.8245767951011658), ('计算机', 0.8100651502609253), ('叫法', 0.8098558187484741), ('手部', 0.7246338725090027), ('手机', 0.6195603013038635), ('micro', 0.5971686244010925), ('电子产品', 0.5647062063217163), ('软件', 0.561561107635498), ('手指', 0.494046688079834), ('technology', 0.47637590765953064), ('pen', 0.4651668071746826), ('virtual', 0.4590775668621063), ('掌心', 0.4538556635379791), ('智能', 0.40049654245376587), ('智慧', 0.3949573338031769), ('touch', 0.38361087441444397), ('指向', 0.3723030686378479), ('移动', 0.3585004508495331), ('事物', 0.34118232131004333), ('电子元件', 0.3282782733440399), ('笔记本', 0.3156297206878662), ('原名', 0.3028894364833832), ('鼠标', 0.28492796421051025), ('android', 0.25649091601371765), ('指', 0.1655425727367401), ('掌握', 0.16021089255809784), ('chi', 0.15045176446437836), ('前臂', 0.11981695145368576), ('book', 0.09273456782102585), ('手掌', 0.07757095992565155), ('按键', 0.06321503221988678), ('小型', 0.05425526574254036), ('一体机', 0.04848058149218559), ('my', 0.03250341862440109), ('psp', 0.01875465363264084), ('跨平台', 0.01767222210764885), ('电脑游戏', 0.005152992904186249)])
+```
+
+## Evaluation
+
+### C-MTEB Retrieval task
+
+([Chinese Massive Text Embedding Benchmark](https://github.com/FlagOpen/FlagEmbedding/tree/master/C_MTEB))
+
+Metric: nDCG@10
+
+|                              Model Name                               | Max Length |  Average  |  Cmedqa   |   Covid   |    Du     |   Ecom    |  Medical  |  MMarco   |    T2     |   Video   | 
+|:---------------------------------------------------------------------:|:----------:|:---------:|:---------:|:---------:|:---------:|:---------:|:---------:|:---------:|:---------:|:---------:|
+|           [BM25-zh](https://github.com/castorini/pyserini)            |     -      |   50.37   |   13.70   | **86.58** |   57.13   |   44.04   |   32.08   |   48.31   |   60.48   |   60.64   |
+|          [bge-m3-sparse](https://huggingface.co/BAAI/bge-m3)          |    512     |   57.00   |   24.50   |   76.09   |   71.51   |   50.49   |   43.93   |   59.28   |   71.76   |   58.43   |
+| [inf-wse-v1-base-zh](https://huggingface.co/infly/inf-wse-v1-base-zh) |    512     |   61.16   |   20.51   |   76.41   |   79.84   |   56.78   |   46.24   |   66.40   |   76.50   |   68.57   |
+|                        **inf-wse-v2-base-zh**                         |    512     | **69.15** | **30.64** |   79.38   | **87.12** | **64.95** | **56.54** | **78.80** | **83.05** | **72.69** |
+
+All results, except for BM25, are measured by building the sparse index via [Qdrant](https://github.com/qdrant/qdrant).

config.json

@@ -0,0 +1,27 @@
+{
+  "architectures": [
+    "RoFormerForMaskedLM"
+  ],
+    "auto_map": {
+    "AutoModelForMaskedLM": "modeling_sparse.RoFormerForSparseEmbeddingV2"
+  },
+  "attention_probs_dropout_prob": 0.1,
+  "embedding_size": 768,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 1536,
+  "model_type": "roformer",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "rotary_value": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.39.3",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 50000
+}

model.safetensors

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

modeling_sparse.py

@@ -0,0 +1,18 @@
+import torch
+from transformers import RoFormerForMaskedLM
+
+
+class RoFormerForSparseEmbeddingV2(RoFormerForMaskedLM):
+    def forward(self, input_ids, attention_mask, return_sparse=False):
+        logits = super().forward(input_ids, attention_mask)['logits']  # [B,L,V]
+        token_mask = (1 - attention_mask.unsqueeze(-1)) * -1e4  # [B,L,1]
+        token_mask[:, 0, :] = -1e4
+        last_ind = torch.sum(attention_mask, -1, keepdim=True).unsqueeze(-1) - 1  # [B,1,1]
+        token_mask = torch.scatter(token_mask, -2, last_ind, -1e4)
+        logits = logits + token_mask
+        emb = torch.log(1 + torch.max(torch.relu(logits), dim=-2).values)  # [B,V]
+
+        if return_sparse:
+            emb = emb.to_sparse()
+
+        return emb

special_tokens_map.json

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

tokenization_sparse.py

@@ -0,0 +1,112 @@
+from transformers.models.roformer.tokenization_roformer import (WordpieceTokenizer, whitespace_tokenize,
+                                                                RoFormerTokenizer)
+
+
+# Copied from transformers.models.roformer.tokenization_roformer.BasicTokenizer._is_chinese_char
+def _is_chinese_char(cp):
+    """Checks whether CP is the codepoint of a CJK character."""
+    # This defines a "chinese character" as anything in the CJK Unicode block:
+    #   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
+    #
+    # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
+    # despite its name. The modern Korean Hangul alphabet is a different block,
+    # as is Japanese Hiragana and Katakana. Those alphabets are used to write
+    # space-separated words, so they are not treated specially and handled
+    # like the all of the other languages.
+    if (
+            (cp >= 0x4E00 and cp <= 0x9FFF)
+            or (cp >= 0x3400 and cp <= 0x4DBF)  #
+            or (cp >= 0x20000 and cp <= 0x2A6DF)  #
+            or (cp >= 0x2A700 and cp <= 0x2B73F)  #
+            or (cp >= 0x2B740 and cp <= 0x2B81F)  #
+            or (cp >= 0x2B820 and cp <= 0x2CEAF)  #
+            or (cp >= 0xF900 and cp <= 0xFAFF)
+            or (cp >= 0x2F800 and cp <= 0x2FA1F)  #
+    ):  #
+        return True
+
+    return False
+
+
+# Modified from transformers.models.roformer.tokenization_roformer.WordpieceTokenizer
+class ChineseWordpieceTokenizer(WordpieceTokenizer):
+    def tokenize(self, text):
+        """
+        Tokenizes a piece of text into its word pieces. This uses a greedy longest-match-first algorithm to perform
+        tokenization using the given vocabulary.
+
+        For example, `input = "unaffable"` wil return as output `["un", "##aff", "##able"]`.
+
+        Args:
+            text: A single token or whitespace separated tokens. This should have
+                already been passed through *BasicTokenizer*.
+
+        Returns:
+            A list of wordpiece tokens.
+        """
+
+        output_tokens = []
+        for token in whitespace_tokenize(text):
+            chars = list(token)
+            if len(chars) > self.max_input_chars_per_word:
+                output_tokens.append(self.unk_token)
+                continue
+
+            is_bad = False
+            start = 0
+            sub_tokens = []
+            while start < len(chars):
+                end = len(chars)
+                cur_substr = None
+                while start < end:
+                    substr = "".join(chars[start:end])
+                    if start > 0 and not _is_chinese_char(ord(substr[0])):  # only add ## when not Chinese character
+                        substr = "##" + substr
+                    if substr in self.vocab:
+                        cur_substr = substr
+                        break
+                    end -= 1
+                if cur_substr is None:
+                    is_bad = True
+                    break
+                sub_tokens.append(cur_substr)
+                start = end
+
+            if is_bad:
+                output_tokens.append(self.unk_token)
+            else:
+                output_tokens.extend(sub_tokens)
+        return output_tokens
+
+
+class ChineseRoFormerTokenizer(RoFormerTokenizer):
+    def __init__(
+            self,
+            vocab_file,
+            do_lower_case=True,
+            do_basic_tokenize=True,
+            never_split=None,
+            unk_token="[UNK]",
+            sep_token="[SEP]",
+            pad_token="[PAD]",
+            cls_token="[CLS]",
+            mask_token="[MASK]",
+            tokenize_chinese_chars=False,
+            strip_accents=None,
+            **kwargs,
+    ):
+        super().__init__(
+            vocab_file=vocab_file,
+            do_lower_case=do_lower_case,
+            do_basic_tokenize=do_basic_tokenize,
+            never_split=never_split,
+            unk_token=unk_token,
+            sep_token=sep_token,
+            pad_token=pad_token,
+            cls_token=cls_token,
+            mask_token=mask_token,
+            tokenize_chinese_chars=tokenize_chinese_chars,
+            strip_accents=strip_accents,
+            **kwargs,
+        )
+        self.wordpiece_tokenizer = ChineseWordpieceTokenizer(vocab=self.vocab, unk_token=str(unk_token))

tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_sparse.ChineseRoFormerTokenizer",
+      ""
+    ]
+  },
+  "tokenizer_class": "ChineseRoFormerTokenizer",
+  "tokenize_chinese_chars": false
+}