Update README.md

README.md

@@ -157,6 +157,121 @@ On LMEB-Dialogue, a compact embedding model paired with our Nano reranker, which
 ![lmeb](./assets/lmeb.jpg)
 ![lmeb_emb](./assets/lmeb_emb.jpg)
 
+## Usage
+### 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-Nano-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.985496461391449, 0.00017952796770259738]
+rankings: [{'corpus_id': 0, 'score': 0.985496461391449}, {'corpus_id': 1, 'score': 0.00017952796770259738}]
+'''
+
+```
+
+
 #### 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.