Update README.md

README.md

@@ -91,20 +91,94 @@ On LMEB, reranking models demonstrate a clear advantage, with even the 0.27B Nan
 # Usage
 ```python
 import argparse
-
-from kalm_reranker import KaLMReranker
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser()
+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-Small"
+        default="KaLM-Embedding/KaLM-Reranker-V1-Small",
+        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("--device", default=None)
-    args = parser.parse_args()
+    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
+
 
-    reranker = KaLMReranker(args.model, device=args.device)
+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.",
@@ -115,11 +189,15 @@ def main() -> None:
     pairs = [(query, document) for document in documents]
     print("scores:", reranker.predict(pairs, instruction=instruction))
     print("rankings:", reranker.rank(query, documents, instruction=instruction))
+    
+    # scores: [0.9999822378158569, 3.187565198459197e-06]
+    # rankings: [{'corpus_id': 0, 'score': 0.9999822378158569}, {'corpus_id': 1, 'score': 3.187565198459197e-06}]
 
 
 if __name__ == "__main__":
     main()
 
+
 ```
 
 # Citation