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

This commit is contained in:
Tejas Dharamsi
2026-05-13 22:01:07 -07:00
committed by GitHub
parent a6a6c3119b
commit c016246b0f
@@ -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)