ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: unify greedy check with _is_greedy helper

Browse Source 
- Replace batch-scattered temperature==0 checks with (temperature == 0).all()

- Reuse _is_greedy in standalone sample() function
This commit is contained in:
ViperEkura
2026-08-07 14:14:15 +08:00
parent 6f49738991
commit ef1bb6f401
1 changed files with 2 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/inference/sample.py
+2 -6
View File
@@ -266,7 +266,7 @@ class SamplingPipeline(BaseSamplingStrategy):
@staticmethod @staticmethod
def _is_greedy(temperature: Union[float, Tensor]) -> bool: def _is_greedy(temperature: Union[float, Tensor]) -> bool:
if isinstance(temperature, Tensor): if isinstance(temperature, Tensor):
return temperature.numel() == 1 and temperature.item() == 0 return bool((temperature == 0).all())
return temperature == 0 return temperature == 0
@torch.inference_mode() @torch.inference_mode()
@@ -364,11 +364,7 @@ def sample(
``chosen_logprobs`` has shape ``[batch]``. ``chosen_logprobs`` has shape ``[batch]``.
""" """
greedy = ( greedy = (
( bool((temperature == 0).all())
isinstance(temperature, Tensor)
and temperature.numel() == 1
and temperature.item() == 0
)
if isinstance(temperature, Tensor) if isinstance(temperature, Tensor)
else temperature == 0 else temperature == 0
) )