minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[Rerank] Early-exit logprob scan and hoist math import (#25046)

Browse Source
This commit is contained in:
Tejas Dharamsi
2026-05-13 22:01:07 -07:00
committed by GitHub
parent a6a6c3119b
commit c016246b0f
1 changed files with 7 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/entrypoints/openai/serving_rerank.py
+7 -2
View File
@@ -1,5 +1,6 @@
import heapq import heapq
import logging import logging
import math
from typing import Any, Dict, List, Optional, Union from typing import Any, Dict, List, Optional, Union
from fastapi import Request from fastapi import Request
@@ -533,8 +534,6 @@ class OpenAIServingRerank(OpenAIServingBase):
def _extract_score_from_logprobs(self, ret: Dict[str, Any]) -> float: def _extract_score_from_logprobs(self, ret: Dict[str, Any]) -> float:
"""Extract reranking score from generation response with logprobs.""" """Extract reranking score from generation response with logprobs."""
import math
# Get logprobs from the response # Get logprobs from the response
meta_info = ret.get("meta_info", {}) meta_info = ret.get("meta_info", {})
output_top_logprobs = meta_info.get("output_top_logprobs", []) output_top_logprobs = meta_info.get("output_top_logprobs", [])
@@ -546,13 +545,19 @@ class OpenAIServingRerank(OpenAIServingBase):
# Format: list of tuples (logprob, token_id, token_text) # Format: list of tuples (logprob, token_id, token_text)
p_yes = 0.0 p_yes = 0.0
p_no = 0.0 p_no = 0.0
found_yes = False
found_no = False
for item in top_logprobs: for item in top_logprobs:
logprob, token_id = item[0], item[1] logprob, token_id = item[0], item[1]
if token_id == self._yes_token_id: if token_id == self._yes_token_id:
p_yes = math.exp(logprob) p_yes = math.exp(logprob)
found_yes = True
elif token_id == self._no_token_id: elif token_id == self._no_token_id:
p_no = math.exp(logprob) p_no = math.exp(logprob)
found_no = True
if found_yes and found_no:
break
return _qwen3_rerank_score(p_yes, p_no) return _qwen3_rerank_score(p_yes, p_no)