refactor: 替换魔法字符串为_STOP sentinel,修复generator清理逻辑

This commit is contained in:
2026-05-06 20:37:16 +08:00
parent b89f8436ea
commit ffff05b2c6
3 changed files with 9 additions and 15 deletions
+3 -4
View File
@@ -12,6 +12,8 @@ from torch import Tensor
from astrai.model.automodel import AutoModel
from astrai.tokenize import AutoTokenizer
_STOP = object()
class _RadixNode:
"""Internal node for the radix tree prefix cache.
@@ -290,9 +292,6 @@ def apply_sampling_strategies(
) -> Tensor:
"""Applies temperature scaling, top-k filtering, and top-p (nucleus) filtering.
Operates on a clone of the input logits to avoid in-place modification
of the inference tensor.
Args:
logits: Raw logits tensor of shape (batch, vocab_size).
temperature: Temperature scaling factor (1.0 = no scaling).
@@ -704,7 +703,7 @@ class InferenceScheduler:
for t in tasks:
if t.is_finished(self.tokenizer.stop_ids):
if t.stream_callback:
t.stream_callback("[DONE]")
t.stream_callback(_STOP)
def _run_generation_loop(self) -> None:
"""Main generation loop run in a daemon thread.