fit for Sentence Transformer

README.md

@@ -85,40 +85,53 @@ flash-attn>2.3.5
 
 ### 示例脚本 Demo
 
+#### Huggingface Transformers
+
 ```python
-from transformers import AutoModel, AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoModel, LlamaTokenizer, AutoModelForSequenceClassification
 import torch
 import numpy as np
 
+# from https://github.com/huggingface/transformers/blob/v4.44.2/src/transformers/models/xlm_roberta/tokenization_xlm_roberta.py
+class MiniCPMRerankerLLamaTokenizer(LlamaTokenizer):
+    def build_inputs_with_special_tokens(
+            self, token_ids_0, token_ids_1 = None
+        ):
+            """
+            - single sequence: `<s> X </s>`
+            - pair of sequences: `<s> A </s> B`
+
+            Args:
+                token_ids_0 (`List[int]`):
+                    List of IDs to which the special tokens will be added.
+                token_ids_1 (`List[int]`, *optional*):
+                    Optional second list of IDs for sequence pairs.
+
+            Returns:
+                `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
+            """
+
+            if token_ids_1 is None:
+                return super().build_inputs_with_special_tokens(token_ids_0)
+            bos = [self.bos_token_id]
+            sep = [self.eos_token_id]
+            return bos + token_ids_0 + sep + token_ids_1
+
 model_name = "openbmb/MiniCPM-Reranker"
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+tokenizer = MiniCPMRerankerLLamaTokenizer.from_pretrained(model_name, trust_remote_code=True)
 tokenizer.padding_side = "right"
+
 model = AutoModelForSequenceClassification.from_pretrained(model_name, trust_remote_code=True,attn_implementation="flash_attention_2", torch_dtype=torch.float16).to("cuda")
 model.eval()
-max_len_q, max_len_d = 512, 512
-
-def tokenize_our(query,doc):
-    input_id_query = tokenizer.encode(query, add_special_tokens=False, max_length=max_len_q, truncation=True)
-    input_id_doc = tokenizer.encode(doc, add_special_tokens=False, max_length=max_len_d, truncation=True)
-    pad_input = {"input_ids": [tokenizer.bos_token_id] + input_id_query + [tokenizer.eos_token_id] + input_id_doc}
-    return tokenizer.pad(
-        pad_input,
-        padding="max_length",
-        max_length=max_len_q + max_len_d + 2,
-        return_tensors="pt",
-    )
 
 @torch.no_grad()
 def rerank(input_query, input_docs):
-    tokenized_inputs = [tokenize_our(input_query, input_doc).to("cuda") for input_doc in input_docs]
-    input_ids = {
-      "input_ids": [tokenized_input["input_ids"] for tokenized_input in tokenized_inputs],
-      "attention_mask": [tokenized_input["attention_mask"] for tokenized_input in tokenized_inputs]
-    }
-
-    for k in input_ids:
-      input_ids[k] = torch.stack(input_ids[k]).to("cuda")
-    outputs = model(**input_ids)
+    tokenized_inputs = tokenizer([[input_query, input_doc] for input_doc in input_docs], return_tensors="pt", padding=True, truncation=True, max_length=1024) 
+
+    for k in tokenized_inputs:
+      tokenized_inputs [k] = tokenized_inputs[k].to("cuda")
+
+    outputs = model(**tokenized_inputs)
     score = outputs.logits
     return score.float().detach().cpu().numpy()
 
@@ -136,6 +149,62 @@ for i in range(len(queries)):
 print(np.array(scores))  # [[[-4.7460938][-8.8515625]]]
 ```
 
+#### Sentence Transformer
+
+```python
+from sentence_transformers import CrossEncoder
+from transformers import LlamaTokenizer
+import torch
+
+# from https://github.com/huggingface/transformers/blob/v4.44.2/src/transformers/models/xlm_roberta/tokenization_xlm_roberta.py
+class MiniCPMRerankerLLamaTokenizer(LlamaTokenizer):
+    def build_inputs_with_special_tokens(
+            self, token_ids_0, token_ids_1 = None
+        ):
+            """
+            - single sequence: `<s> X </s>`
+            - pair of sequences: `<s> A </s> B`
+
+            Args:
+                token_ids_0 (`List[int]`):
+                    List of IDs to which the special tokens will be added.
+                token_ids_1 (`List[int]`, *optional*):
+                    Optional second list of IDs for sequence pairs.
+
+            Returns:
+                `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
+            """
+
+            if token_ids_1 is None:
+                return super().build_inputs_with_special_tokens(token_ids_0)
+            bos = [self.bos_token_id]
+            sep = [self.eos_token_id]
+            return bos + token_ids_0 + sep + token_ids_1
+
+model_name = "openbmb/MiniCPM-Reranker"
+model = CrossEncoder(model_name,max_length=1024,trust_remote_code=True, automodel_args={"attn_implementation":"flash_attention_2","torch_dtype": torch.float16})
+model.tokenizer = MiniCPMRerankerLLamaTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model.tokenizer.padding_side = "right"
+
+query = "中国的首都是哪里？"
+passages = [["beijing", "shanghai"]]
+
+INSTRUCTION = "Query: "
+query = INSTRUCTION + query
+
+sentence_pairs = [[query, doc] for doc in passages]
+
+scores = model.predict(sentence_pairs, convert_to_tensor=True).tolist()
+rankings = model.rank(query, passages, return_documents=True, convert_to_tensor=True)
+
+print(scores) # [0.0087432861328125, 0.00020503997802734375]
+for ranking in rankings:
+    print(f"Score: {ranking['score']:.4f}, Corpus: {ranking['text']}")
+  
+# ID: 0, Score: 0.0087, Text: beijing
+# ID: 1, Score: 0.0002, Text: shanghai
+```
+
 ## 实验结果 Evaluation Results
 
 ### 中文与英文重排序结果 CN/EN Re-ranking Results

config.json

@@ -1,5 +1,5 @@
 {
-  "_name_or_path": "openbmb/RankCPM-R",
+  "_name_or_path": "openbmb/MiniCPM-Reranker",
   "architectures": [
       "MiniCPM"
   ],