fix: fix parsing token_logprobs in sglang_wrapper

Author: ChenZiHong-Gavin

graphgen/models/llm/api/http_client.py

@@ -163,7 +163,7 @@ async def generate_topk_per_token(
         **extra: Any,
     ) -> List[Token]:
         body = self._build_body(text, history or [])
-        body["max_tokens"] = 5
+        body["max_tokens"] = 1
         if self.topk_per_token > 0:
             body["logprobs"] = True
             body["top_logprobs"] = self.topk_per_token

graphgen/models/llm/api/openai_client.py

@@ -105,8 +105,8 @@ async def generate_topk_per_token(
             kwargs["logprobs"] = True
             kwargs["top_logprobs"] = self.topk_per_token
 
-        # Limit max_tokens to 5 to avoid long completions
-        kwargs["max_tokens"] = 5
+        # Limit max_tokens to 1 to avoid long completions
+        kwargs["max_tokens"] = 1
 
         completion = await self.client.chat.completions.create(  # pylint: disable=E1125
             model=self.model_name, **kwargs

graphgen/models/llm/local/sglang_wrapper.py

@@ -5,7 +5,6 @@
 from graphgen.bases.datatypes import Token
 
 
-# TODO: implement SGLangWrapper methods
 class SGLangWrapper(BaseLLMWrapper):
     """
     Async inference backend based on SGLang offline engine.
@@ -59,43 +58,39 @@ def _build_sampling_params(
             params["logprobs"] = topk
         return params
 
-    def _prep_prompt(self, text: str, history: Optional[List[str]] = None) -> str:
+    def _prep_prompt(self, text: str, history: Optional[List[dict]] = None) -> str:
         """Convert raw text (+ optional history) into a single prompt string."""
         parts = []
         if self.system_prompt:
             parts.append(self.system_prompt)
         if history:
             assert len(history) % 2 == 0, "History must have even length (u/a turns)."
-            parts.extend(history)
+            parts.extend([item["content"] for item in history])
         parts.append(text)
         return "\n".join(parts)
 
     def _tokens_from_output(self, output: Dict[str, Any]) -> List[Token]:
-        """
-        Convert SGLang logprobs output into List[Token].
-        SGLang returns:
-            output['logprobs'][t] -> {
-                "token": <str>,
-                "logprob": <float>,
-                "top_k_tokens": [...],
-                "top_k_logprobs": [...],
-            }
-        """
         tokens: List[Token] = []
-        if "logprobs" not in output or not output["logprobs"]:
-            return tokens
 
-        for entry in output["logprobs"]:
-            token_str = entry["token"]
-            logprob = entry["logprob"]
-            prob = math.exp(logprob)
+        meta = output.get("meta_info", {})
+        logprobs = meta.get("output_token_logprobs", [])
+        topks = meta.get("output_top_logprobs", [])
+
+        tokenizer = self.engine.tokenizer_manager.tokenizer
+
+        for idx, (lp, tid, _) in enumerate(logprobs):
+            prob = math.exp(lp)
+            tok_str = tokenizer.decode([tid])
 
             top_candidates = []
-            if self.topk > 0 and "top_k_tokens" in entry:
-                for tok, lp in zip(entry["top_k_tokens"], entry["top_k_logprobs"]):
-                    top_candidates.append(Token(tok, math.exp(lp)))
+            if self.topk > 0 and idx < len(topks):
+                for t_lp, t_tid, _ in topks[idx][: self.topk]:
+                    top_candidates.append(
+                        Token(text=tokenizer.decode([t_tid]), prob=math.exp(t_lp))
+                    )
+
+            tokens.append(Token(text=tok_str, prob=prob, top_candidates=top_candidates))
 
-            tokens.append(Token(token_str, prob, top_candidates=top_candidates))
         return tokens
 
     async def generate_answer(
@@ -112,7 +107,7 @@ async def generate_answer(
             topk=0,  # no logprobs needed for simple generation
         )
 
-        outputs = self.engine.generate([prompt], sampling_params)
+        outputs = await self.engine.async_generate([prompt], sampling_params)
         return self.filter_think_tags(outputs[0]["text"])
 
     async def generate_topk_per_token(
@@ -125,45 +120,23 @@ async def generate_topk_per_token(
         sampling_params = self._build_sampling_params(
             temperature=self.temperature,
             top_p=self.top_p,
-            max_tokens=5,  # keep short for token-level analysis
+            max_tokens=1,  # keep short for token-level analysis
             topk=self.topk,
-            logprobs=True,
         )
 
-        outputs = self.engine.generate([prompt], sampling_params)
+        outputs = await self.engine.async_generate(
+            [prompt], sampling_params, return_logprob=True, top_logprobs_num=5
+        )
+        print(outputs)
         return self._tokens_from_output(outputs[0])
 
     async def generate_inputs_prob(
         self, text: str, history: Optional[List[str]] = None, **extra: Any
     ) -> List[Token]:
-        """
-        Return per-token probabilities for the *input* sequence.
-        SGLang offline engine does not expose this directly; we emulate by
-        generating 0 new tokens with logprobs enabled (returns prompt logprobs).
-        """
-        prompt = self._prep_prompt(text, history)
-        sampling_params = self._build_sampling_params(
-            temperature=0.0,  # deterministic
-            top_p=1.0,
-            max_tokens=0,  # generate nothing
-            topk=self.topk,
-            logprobs=True,
+        raise NotImplementedError(
+            "SGLangWrapper does not support per-token logprobs yet."
         )
 
-        outputs = self.engine.generate([prompt], sampling_params)
-        # SGLang returns prompt logprobs under key 'prompt_logprobs' when max_new_tokens=0
-        prompt_logprobs = outputs[0].get("prompt_logprobs", [])
-        tokens: List[Token] = []
-        for entry in prompt_logprobs:
-            tokens.append(
-                Token(
-                    text=entry["token"],
-                    prob=math.exp(entry["logprob"]),
-                    top_candidates=[],  # SGLang does not give top-k for prompt tokens
-                )
-            )
-        return tokens
-
     def shutdown(self) -> None:
         """Gracefully shutdown the SGLang engine."""
         if hasattr(self, "engine"):