Integrate with Sentence Transformers

1_Pooling/config.json

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

README.md

@@ -45,7 +45,25 @@ Cost comparison of FastText and PEARL. The estimated memory is calculated by the
 
 ## Usage
 
-Below is an example of entity retrieval, and we reuse the code from E5.
+### Sentence Transformers
+PEARL is integrated with the Sentence Transformers library, and can be used like so:
+
+```python
+from sentence_transformers import SentenceTransformer, util
+
+query_texts = ["The New York Times"]
+doc_texts = [ "NYTimes", "New York Post", "New York"]
+input_texts = query_texts + doc_texts
+
+model = SentenceTransformer("Lihuchen/pearl_base")
+embeddings = model.encode(input_texts)
+scores = util.cos_sim(embeddings[0], embeddings[1:]) * 100
+print(scores.tolist())
+# [[85.61601257324219, 73.65623474121094, 70.36174774169922]]
+```
+
+### Transformers
+You can also use `transformers` to use PEARL. Below is an example of entity retrieval, and we reuse the code from E5.
 
 ```python
 import torch.nn.functional as F

config_sentence_transformers.json

@@ -0,0 +1,7 @@
+{
+  "__version__": {
+    "sentence_transformers": "2.3.1",
+    "transformers": "4.37.0",
+    "pytorch": "2.1.0+cu121"
+  }
+}

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

sentence_bert_config.json

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