Integrate with Sentence Transformers v5.4

1_LogitScore/config.json

@@ -0,0 +1,4 @@
+{
+    "true_token_id": 9454,
+    "false_token_id": null
+}

README.md

@@ -40,18 +40,47 @@ This model is released under a non-commercial license. If you'd like a commercia
 
 ## How to Use
 
+### Using Sentence Transformers
+
+Install Sentence Transformers:
+
+```bash
+pip install sentence_transformers
+```
+
+Then load the model and score query/document pairs. `model.predict` returns the raw "Yes" logit per pair; rankings can be used directly. To map the logits to a 0-1 score range, apply a temperature-scaled sigmoid: `sigmoid(score / 5)`.
+
 ```python
 from sentence_transformers import CrossEncoder
 
-model = CrossEncoder("zeroentropy/zerank-2", trust_remote_code=True)
+model = CrossEncoder("zeroentropy/zerank-2")
 
 query_documents = [
     ("What is 2+2?", "4"),
     ("What is 2+2?", "The answer is definitely 1 million"),
 ]
 
-scores = model.predict(query_documents)
+scores = model.predict(query_documents, convert_to_tensor=True)
 print(scores)
+# tensor([ 5.4062, -4.5000], device='cuda:0', dtype=torch.bfloat16)
+
+# Optional: convert to 0-1 probabilities
+probabilities = (scores / 5).sigmoid()
+print(probabilities)
+# tensor([0.7461, 0.2891], device='cuda:0', dtype=torch.bfloat16)
+```
+
+You can also use `model.rank` to score and sort a list of documents for a single query:
+
+```python
+rankings = model.rank(
+    "What is 2+2?",
+    ["4", "The answer is definitely 1 million"],
+)
+for r in rankings:
+    print(r)
+# {'corpus_id': 0, 'score': np.float32(5.40625)}
+# {'corpus_id': 1, 'score': np.float32(-4.5)}
 ```
 
 The model can also be inferenced using ZeroEntropy's [/models/rerank](https://docs.zeroentropy.dev/api-reference/models/rerank) endpoint, and on [AWS Marketplace](https://aws.amazon.com/marketplace/pp/prodview-o7avk66msiukc).

chat_template.jinja

@@ -1,3 +1,12 @@
+{%- set zerank_query = messages | selectattr("role", "eq", "query") | map(attribute="content") | list -%}
+{%- set zerank_document = messages | selectattr("role", "eq", "document") | map(attribute="content") | list -%}
+{%- if zerank_query and zerank_document -%}
+{{- '<|im_start|>system\n' + (zerank_query | first) + '<|im_end|>\n' -}}
+{{- '<|im_start|>user\n' + (zerank_document | first) + '<|im_end|>\n' -}}
+{%- if add_generation_prompt -%}
+{{- '<|im_start|>assistant\n' -}}
+{%- endif -%}
+{%- else -%}
 {%- if tools %}
     {{- '<|im_start|>system\n' }}
     {%- if messages[0].role == 'system' %}
@@ -86,4 +95,5 @@
     {%- if enable_thinking is defined and enable_thinking is false %}
         {{- '<think>\n\n</think>\n\n' }}
     {%- endif %}
-{%- endif %}
+{%- endif %}
+{%- endif -%}

config_sentence_transformers.json

@@ -0,0 +1,9 @@
+{
+  "__version__": {
+    "sentence_transformers": "5.4.0"
+  },
+  "activation_fn": "torch.nn.modules.linear.Identity",
+  "default_prompt_name": null,
+  "model_type": "CrossEncoder",
+  "prompts": {}
+}

modules.json

@@ -0,0 +1,14 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.base.modules.transformer.Transformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_LogitScore",
+    "type": "sentence_transformers.cross_encoder.modules.logit_score.LogitScore"
+  }
+]

sentence_bert_config.json

@@ -0,0 +1,20 @@
+{
+    "transformer_task": "text-generation",
+    "modality_config": {
+        "text": {
+            "method": "forward",
+            "method_output_name": "logits"
+        },
+        "message": {
+            "method": "forward",
+            "method_output_name": "logits",
+            "format": "flat"
+        }
+    },
+    "module_output_name": "causal_logits",
+    "processing_kwargs": {
+        "chat_template": {
+            "add_generation_prompt": true
+        }
+    }
+}