Update README.md

README.md

@@ -48,15 +48,10 @@ SentenceTransformer(
 
 ## Usage
 
-### Direct Usage (Sentence Transformers)
+### Usage (Sentence Transformers)
 
-First install the Sentence Transformers library:
+You can perform inference using SentenceTransformers with the following code:
 
-```bash
-pip install -U sentence-transformers
-```
-
-Then you can load this model and run inference.
 ```python
 from sentence_transformers import SentenceTransformer
 import torch.nn.functional as F
@@ -65,7 +60,8 @@ import torch.nn.functional as F
 # The argument "trust_remote_code=True" is required to load the model
 model = SentenceTransformer("pkshatech/RoSEtta-base-ja",trust_remote_code=True)
 
-# Don't forget to add the prefix "query: " for query-side or "passage: " for passage-side texts.
+# Each input text should start with "query: " or "passage: ".
+# For tasks other than retrieval, you can simply use the "query: " prefix.
 sentences = [
     'query: PKSHAはどんな会社ですか？',
     'passage: 研究開発したアルゴリズムを、多くの企業のソフトウエア・オペレーションに導入しています。',
@@ -79,19 +75,56 @@ print(embeddings.shape)
 # Get the similarity scores for the embeddings
 similarities = F.cosine_similarity(embeddings.unsqueeze(0), embeddings.unsqueeze(1), dim=2)
 print(similarities)
-# tensor([[1.0000, 0.5910, 0.4332, 0.5421],
-#         [0.5910, 1.0000, 0.4977, 0.6969],
-#         [0.4332, 0.4977, 1.0000, 0.7475],
-#         [0.5421, 0.6969, 0.7475, 1.0000]])
+# [[1.0000, 0.5910, 0.4332, 0.5421],
+# [0.5910, 1.0000, 0.4977, 0.6969],
+# [0.4332, 0.4977, 1.0000, 0.7475],
+# [0.5421, 0.6969, 0.7475, 1.0000]]
 ```
 
-<!--
 ### Direct Usage (Transformers)
 
-<details><summary>Click to see the direct usage in Transformers</summary>
+You can perform inference using Transformers with the following code:
 
-</details>
--->
+```python
+import torch.nn.functional as F
+from torch import Tensor
+from transformers import AutoTokenizer, AutoModel
+
+def mean_pooling(last_hidden_states: Tensor,attention_mask: Tensor) -> Tensor:
+    emb = last_hidden_states * attention_mask.unsqueeze(-1)
+    emb = emb.sum(dim=1) / attention_mask.sum(dim=1).unsqueeze(-1)
+    return emb
+
+# Download from the 🤗 Hub
+tokenizer = AutoTokenizer.from_pretrained("pkshatech/RoSEtta-base-ja")
+# The argument "trust_remote_code=True" is required to load the model
+model = AutoModel.from_pretrained("pkshatech/RoSEtta-base-ja",trust_remote_code=True)
+
+# Each input text should start with "query: " or "passage: ".
+# For tasks other than retrieval, you can simply use the "query: " prefix.
+sentences = [
+    'query: PKSHAはどんな会社ですか？',
+    'passage: 研究開発したアルゴリズムを、多くの企業のソフトウエア・オペレーションに導入しています。',
+    'query: 日本で一番高い山は？',
+    'passage: 富士山（ふじさん）は、標高3776.12 m、日本最高峰（剣ヶ峰）の独立峰で、その優美な風貌は日本国外でも日本の象徴として広く知られている。',
+]
+
+# Tokenize the input texts
+batch_dict = tokenizer(sentences, max_length=1024, padding=True, truncation=True, return_tensors='pt')
+
+outputs = model(**batch_dict)
+embeddings = mean_pooling(outputs.last_hidden_state, batch_dict['attention_mask'])
+print(embeddings.shape)
+# [4, 768]
+
+# Get the similarity scores for the embeddings
+similarities = F.cosine_similarity(embeddings.unsqueeze(0), embeddings.unsqueeze(1), dim=2)
+print(similarities)
+# [[1.0000, 0.5910, 0.4332, 0.5421],
+# [0.5910, 1.0000, 0.4977, 0.6969],
+# [0.4332, 0.4977, 1.0000, 0.7475],
+# [0.5421, 0.6969, 0.7475, 1.0000]]
+```
 
 <!--
 ### Downstream Usage (Sentence Transformers)