Update README.md

README.md

@@ -158,6 +158,75 @@ On LMEB-Dialogue, a compact embedding model paired with our Nano reranker, which
 ![lmeb_emb](./assets/lmeb_emb.jpg)
 
 ## Usage
+### Using Sentence Transformers
+
+KaLM-Reranker-V1-Nano-R2 can be loaded as a modular Sentence Transformers
+`CrossEncoder`. This integration requires `sentence-transformers>=5.6,<6`,
+`transformers>=5.3,<6`, and the PyTorch backend. The repository contains custom
+modeling code, so load only trusted revisions and pass `trust_remote_code=True`.
+
+```bash
+pip install "sentence-transformers>=5.6,<6" "transformers>=5.3,<6"
+```
+
+```python
+import torch
+from sentence_transformers import CrossEncoder
+
+model = CrossEncoder(
+    "KaLM-Embedding/KaLM-Reranker-V1-Nano-R2",
+    trust_remote_code=True,
+    device="cuda",
+    model_kwargs={"dtype": torch.bfloat16, "chunk_size": 4},
+)
+
+query = "What is the capital of China?"
+documents = [
+    "The capital of China is Beijing.",
+    "Gravity attracts bodies toward one another.",
+]
+pairs = [(query, document) for document in documents]
+
+# The default output is P(yes).
+scores = model.predict(pairs)
+rankings = model.rank(query, documents, return_documents=True)
+
+# CrossEncoder prompts are interpreted as KaLM task instructions.
+instruction = "Given a web search query, retrieve passages that answer the query."
+custom_scores = model.predict(pairs, prompt=instruction)
+custom_rankings = model.rank(query, documents, prompt=instruction)
+
+# Use Identity to return yes_logit - no_logit instead of P(yes).
+margins = model.predict(pairs, activation_fn=torch.nn.Identity())
+
+print(f"scores: {scores}")
+print(f"rankings: {rankings}")
+print(f"custom_scores: {custom_scores}")
+print(f"custom_rankings: {custom_rankings}")
+print(f"margins: {margins}")
+
+'''
+scores: [9.8549646e-01 1.7952797e-04]
+rankings: [{'corpus_id': 0, 'score': 0.98549646, 'text': 'The capital of China is Beijing.'}, {'corpus_id': 1, 'score': 0.00017952797, 'text': 'Gravity attracts bodies toward one another.'}]
+custom_scores: [9.7771388e-01 7.4846226e-05]
+custom_rankings: [{'corpus_id': 0, 'score': 0.9777139}, {'corpus_id': 1, 'score': 7.4846226e-05}]
+margins: [ 4.21875 -8.625  ]
+'''
+
+```
+
+Inputs must be ordered as `(query, document)`. By default, queries are
+truncated to 512 tokens and documents to 1024 tokens. `chunk_size=4` performs a
+mask-aware mean over each consecutive group of four encoder token states before
+passing the compressed encoder output to the decoder. Set `chunk_size=None` to
+disable compression, or change `model[0].chunk_size` after loading.
+
+For CPU inference, use `device="cpu"` and
+`model_kwargs={"dtype": torch.float32, "chunk_size": 4}`. Only the PyTorch
+inference backend is currently supported; training, ONNX, and OpenVINO are not
+included in this release.
+
+
 ### Using transformers
 ```python
 import argparse