[Tunix] Use self.config.return_logprobs to determine if logprobs should be returned.

wang2yn84 · The tunix Authors · commit a08371bb86d2 · 2026-03-25T14:49:50.000-07:00
PiperOrigin-RevId: 889446968
diff --git a/tunix/generate/vllm_sampler.py b/tunix/generate/vllm_sampler.py
@@ -505,5 +505,5 @@ def __call__(
         logits=None,
         tokens=out_tokens[0],
         padded_prompt_tokens=all_input_ids,
-        logprobs=out_logprobs[0] if len(out_logprobs[0]) else None,
+        logprobs=out_logprobs[0] if self.config.return_logprobs else None,
     )

Original file line number	Diff line number	Diff line change
`@@ -505,5 +505,5 @@ def __call__(`
`505`	`505`	`logits=None,`
`506`	`506`	`tokens=out_tokens[0],`
`507`	`507`	`padded_prompt_tokens=all_input_ids,`
`508`		`- logprobs=out_logprobs[0] if len(out_logprobs[0]) else None,`
	`508`	`+ logprobs=out_logprobs[0] if self.config.return_logprobs else None,`
`509`	`509`	`)`