Update README.md

README.md

@@ -157,6 +157,188 @@ On LMEB-Dialogue, a compact embedding model paired with our Nano reranker, which
 ![lmeb](./assets/lmeb.jpg)
 ![lmeb_emb](./assets/lmeb_emb.jpg)
 
+### Using Sentence Transformers
+
+KaLM-Reranker-V1-Large-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-Large-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.8677725e-01 9.0574264e-05]
+rankings: [{'corpus_id': 0, 'score': 0.98677725, 'text': 'The capital of China is Beijing.'}, {'corpus_id': 1, 'score': 9.057426e-05, 'text': 'Gravity attracts bodies toward one another.'}]
+custom_scores: [9.7241479e-01 3.2129523e-05]
+custom_rankings: [{'corpus_id': 0, 'score': 0.9724148}, {'corpus_id': 1, 'score': 3.2129523e-05}]
+margins: [ 4.3125  -9.30925]
+'''
+
+```
+
+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
+from typing import Optional
+
+
+def optional_positive_int(value: str) -> Optional[int]:
+    if value.lower() == "none":
+        return None
+    try:
+        parsed = int(value)
+    except ValueError as error:
+        raise argparse.ArgumentTypeError(
+            "must be a positive integer or 'none'"
+        ) from error
+    if parsed <= 0:
+        raise argparse.ArgumentTypeError("must be a positive integer or 'none'")
+    return parsed
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--model",
+        default="KaLM-Embedding/KaLM-Reranker-V1-Large-R2",
+        help="Hugging Face model ID or local checkpoint path.",
+    )
+    parser.add_argument(
+        "--device",
+        default=None,
+        help="Inference device, such as 'cuda', 'cuda:0', or 'cpu'.",
+    )
+    parser.add_argument(
+        "--dtype",
+        default=None,
+        choices=("bfloat16", "bf16", "float16", "fp16", "float32", "fp32"),
+        help="Model parameter dtype. By default, use BF16 on CUDA and FP32 on CPU.",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=32,
+        help="Number of query-document pairs scored per inference batch.",
+    )
+    parser.add_argument(
+        "--query-max-length",
+        type=int,
+        default=512,
+        help=(
+            "Maximum tokens in the raw query before it is inserted into the "
+            "decoder prompt; prompt tokens are not included in this limit."
+        ),
+    )
+    parser.add_argument(
+        "--reranker-max-length",
+        type=int,
+        default=1024,
+        help=(
+            "Maximum encoder tokens for '<Document>: {passage}'. This is not a "
+            "combined query-document context limit."
+        ),
+    )
+    parser.add_argument(
+        "--chunk-size",
+        type=optional_positive_int,
+        default=4,
+        metavar="N|none",
+        help=(
+            "Number of encoder token hidden states per mean-pooled chunk; use "
+            "'none' to disable encoder chunk pooling."
+        ),
+    )
+    return parser
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+
+    from kalm_reranker import KaLMReranker
+
+    reranker = KaLMReranker(
+        args.model,
+        device=args.device,
+        dtype=args.dtype,
+        batch_size=args.batch_size,
+        query_max_length=args.query_max_length,
+        max_length=args.reranker_max_length,
+        chunk_size=args.chunk_size,
+    )
+    query = "What is the capital of China?"
+    documents = [
+        "The capital of China is Beijing.",
+        "Gravity attracts bodies toward one another.",
+    ]
+    instruction = "Given a query, retrieve documents that answer the query."
+
+    pairs = [(query, document) for document in documents]
+    print("scores:", reranker.predict(pairs, instruction=instruction))
+    print("rankings:", reranker.rank(query, documents, instruction=instruction))
+    
+
+if __name__ == "__main__":
+    main()
+
+'''
+scores: [0.9867772459983826, 9.0574256319087e-05]
+rankings: [{'corpus_id': 0, 'score': 0.9867772459983826}, {'corpus_id': 1, 'score': 9.0574256319087e-05}]
+'''
+
+```
+
 #### Ablation on multi-stage training
 
 Across all three model sizes and all seven compression ratios, performance on BEIR and MIRACL improves consistently from Stage 1 to Stage 3, demonstrating the effectiveness of our multi-stage training pipeline. Concretely, Stage 1 establishes a robust foundation for document reranking, distillation in Stage 2 substantially improves performance, and Stage 3 yields further modest gains. More importantly, robustness to compression generally improves across the three training stages. For example, from Stage 1 to Stage 3, the performance retention of KaLM-Reranker-V1-Nano at r = 128 relative to r = 2 increases from 92.88% to 93.80% on BEIR and from 90.93% to 92.15% on MIRACL.