From 88ec786e39cc2d6ded988e1f09227b5f477e25bc Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sun, 19 Jul 2026 16:38:28 +0800 Subject: [PATCH] fix: memmap mode=r, tool parser json.loads, greedy decode --- astrai/inference/api/tool_parser.py | 31 +++++++++++++++++++++++------ astrai/inference/engine.py | 4 ++-- astrai/inference/sample.py | 19 ++++++++++++++++++ astrai/serialization/dataset.py | 2 +- 4 files changed, 47 insertions(+), 9 deletions(-) diff --git a/astrai/inference/api/tool_parser.py b/astrai/inference/api/tool_parser.py index f72d2b0..3ed72b8 100644 --- a/astrai/inference/api/tool_parser.py +++ b/astrai/inference/api/tool_parser.py @@ -7,6 +7,7 @@ Subclasses may optionally consume ``token_ids`` for token-level parsing (e.g. Harmony / VLM-style parsers). """ +import json import re import uuid from abc import ABC, abstractmethod @@ -117,6 +118,29 @@ def _parse_tool_call_json(json_str: str, complete: bool): Returns ``(name, args, valid)``. """ + if complete: + try: + obj = json.loads(json_str) + except json.JSONDecodeError: + return None, "", False + name = obj.get("name") + if not isinstance(name, str) or not name: + return None, "", False + args = obj.get("arguments") + if isinstance(args, dict): + if not args: + args = "" + else: + args = json.dumps(args, ensure_ascii=False) + args = args[1:-1].rstrip() + elif isinstance(args, list): + args = json.dumps(args, ensure_ascii=False) if args else "" + elif isinstance(args, str): + pass + else: + args = str(args) if args is not None else "" + return name, args, True + name_match = re.search(r'"name"\s*:\s*"([^"]*)"', json_str) if not name_match: return None, "", False @@ -127,8 +151,6 @@ def _parse_tool_call_json(json_str: str, complete: bool): return name, "", True raw = args_match.group(1).rstrip() - if complete and raw.endswith("}"): - raw = raw[:-1].rstrip() if raw.startswith("{"): inner = raw[1:].rstrip() if inner.endswith("}"): @@ -156,9 +178,6 @@ def _find_tool_calls(text: str, start_pos: int = 0): break json_str = text[brace:end] - if not _TOOL_CALL_HEAD_RE.search(json_str): - pos = end - continue name, args, valid = _parse_tool_call_json(json_str, complete=True) if not valid or name is None: @@ -186,7 +205,7 @@ def _find_partial_tool_call(text: str, start_pos: int = 0): return None json_str = text[brace:] - if not _TOOL_CALL_HEAD_RE.search(json_str): + if '"name"' not in json_str: return None name, args, valid = _parse_tool_call_json(json_str, complete=False) diff --git a/astrai/inference/engine.py b/astrai/inference/engine.py index 19ea9f8..9181bf4 100644 --- a/astrai/inference/engine.py +++ b/astrai/inference/engine.py @@ -82,8 +82,8 @@ class GenerationRequest: raise ValueError("top_k must be a non-negative integer") if not (0.0 <= top_p <= 1.0): raise ValueError("top_p must be a float between 0.0 and 1.0") - if not (isinstance(temperature, (int, float)) and temperature > 0): - raise ValueError("temperature must be a positive number") + if not (isinstance(temperature, (int, float)) and temperature >= 0): + raise ValueError("temperature must be a non-negative number") if not ( isinstance(frequency_penalty, (int, float)) and -2.0 <= frequency_penalty <= 2.0 diff --git a/astrai/inference/sample.py b/astrai/inference/sample.py index 5ac0104..40b0256 100644 --- a/astrai/inference/sample.py +++ b/astrai/inference/sample.py @@ -263,6 +263,12 @@ class SamplingPipeline(BaseSamplingStrategy): logits = strategy.apply(logits, filter_value, input_ids, input_mask) return logits + @staticmethod + def _is_greedy(temperature: Union[float, Tensor]) -> bool: + if isinstance(temperature, Tensor): + return temperature.numel() == 1 and temperature.item() == 0 + return temperature == 0 + @torch.inference_mode() def sample( self, @@ -273,6 +279,9 @@ class SamplingPipeline(BaseSamplingStrategy): ) -> Tensor: """Apply strategies then sample (softmax + multinomial). + Short-circuits to ``argmax`` when temperature is exactly 0 + (deterministic / greedy decode). + Args: logits: Raw logits ``[batch, vocab_size]``. input_ids: Previously generated token IDs ``[batch, seq_len]``. @@ -281,6 +290,11 @@ class SamplingPipeline(BaseSamplingStrategy): Returns: Sampled token IDs ``[batch]``. """ + for s in self.strategies: + if isinstance(s, TemperatureStrategy) and self._is_greedy(s.temperature): + return logits.argmax(dim=-1) + break + return torch.multinomial( torch.softmax( self.apply(logits, filter_value, input_ids, input_mask), dim=-1 @@ -304,6 +318,9 @@ def sample( Shortcut for ``SamplingPipeline(...).sample(logits)``. + When **temperature** is exactly 0 (scalar or single-element tensor) + the function short-circuits to ``argmax`` for deterministic decode. + Args: logits: Raw logits ``[batch, vocab_size]``. frequency_penalty: Penalty per occurrence for repeated tokens @@ -314,6 +331,8 @@ def sample( Returns: Sampled token IDs ``[batch]``. """ + if SamplingPipeline._is_greedy(temperature): + return logits.argmax(dim=-1) return SamplingPipeline( [ TemperatureStrategy(temperature), diff --git a/astrai/serialization/dataset.py b/astrai/serialization/dataset.py index 08037fc..84629c0 100644 --- a/astrai/serialization/dataset.py +++ b/astrai/serialization/dataset.py @@ -100,7 +100,7 @@ def load_bin(file_path: str) -> Dict[str, List[Tensor]]: arr = np.memmap( os.path.join(file_path, f"{key}.bin"), dtype=info["dtype"], - mode="r+", + mode="r", shape=tuple(info["shape"]), ) segments[key] = [torch.from_numpy(arr)]