fix: memmap mode=r, tool parser json.loads, greedy decode

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)

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):
+        if not (isinstance(temperature, (int, float)) and temperature >= 0):
-            raise ValueError("temperature must be a positive number")
+            raise ValueError("temperature must be a non-negative number")
         if not (
             isinstance(frequency_penalty, (int, float))
             and -2.0 <= frequency_penalty <= 2.0

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),

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)]