cortexgrid-infer 0.3.3__tar.gz → 0.3.4__tar.gz

{cortexgrid_infer-0.3.3 → cortexgrid_infer-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid-infer
-Version: 0.3.3
+Version: 0.3.4
 Summary: Serve apps per task (text-to-text, text-to-image, image-to-mesh) and importers per source (HuggingFace) for running models on cortexgrid
 Project-URL: Homepage, https://github.com/robodatalab/cortexgrid-infer
 Project-URL: Repository, https://github.com/robodatalab/cortexgrid-infer
@@ -245,7 +245,7 @@ it is the floor; only quantizing the weights or running fewer forwards moves it.
 
 ## Inference
 
-`complete` streams `CompletionChunk`s (`.content`, `.tool_calls`, `.finish_reason`):
+`complete` streams `CompletionChunk`s (`.content`, `.thinking`, `.tool_calls`, `.finish_reason`):
 
 ```python
 imp = mg.HuggingFaceImporter("Qwen/Qwen2.5-7B-Instruct", mg.Text2Text)

{cortexgrid_infer-0.3.3 → cortexgrid_infer-0.3.4}/README.md

@@ -217,7 +217,7 @@ it is the floor; only quantizing the weights or running fewer forwards moves it.
 
 ## Inference
 
-`complete` streams `CompletionChunk`s (`.content`, `.tool_calls`, `.finish_reason`):
+`complete` streams `CompletionChunk`s (`.content`, `.thinking`, `.tool_calls`, `.finish_reason`):
 
 ```python
 imp = mg.HuggingFaceImporter("Qwen/Qwen2.5-7B-Instruct", mg.Text2Text)

{cortexgrid_infer-0.3.3 → cortexgrid_infer-0.3.4}/cortexgrid_infer/core.py

@@ -25,6 +25,7 @@ class ToolCall:
 @dataclass
 class CompletionChunk:
     content: str = ""
+    thinking: str = ""
     tool_calls: list[ToolCall] = field(default_factory=list)
     finish_reason: str | None = None
 

{cortexgrid_infer-0.3.3 → cortexgrid_infer-0.3.4}/cortexgrid_infer/importers/base.py

@@ -95,5 +95,8 @@ class Importer(ModelEntry):
         have, so a hand-set value is never overwritten."""
         return self.serve_app.requirements(self.weights())
 
+    def config(self) -> dict[str, str]:
+        return self.serve_app.config()
+
     def client(self, url: str) -> DeployedModel:
         return self.serve_app.client(url, self.model_id)

{cortexgrid_infer-0.3.3 → cortexgrid_infer-0.3.4}/cortexgrid_infer/protocols/completion.py

@@ -35,6 +35,9 @@ _tool_call_id_counter = itertools.count()
 
 TOOL_CALL_OPENERS = ["<tool_call>", "<|tool_call|>", "```tool_call"]
 
+THINKING_OPENER = "<think>"
+THINKING_CLOSER = "</think>"
+
 TOOL_CALL_PATTERNS = [
     re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL),
     re.compile(r"<\|tool_call\|>\s*(\{.*?\})\s*<\|/tool_call\|>", re.DOTALL),
@@ -51,23 +54,27 @@ def encode_tool_call(name: str, arguments: dict[str, Any]) -> str:
     return f'<tool_call>{json.dumps({"name": name, "arguments": arguments})}</tool_call>'
 
 
-def _find_opener(text: str) -> int | None:
-    """Find the earliest tool call opener position in text, or None."""
+def encode_thinking(text: str) -> str:
+    return f"{THINKING_OPENER}{text}{THINKING_CLOSER}"
+
+
+def _find_opener(text: str, openers: Sequence[str]) -> int | None:
+    """Find the earliest position of any of `openers` in text, or None."""
     earliest = None
-    for opener in TOOL_CALL_OPENERS:
+    for opener in openers:
         idx = text.find(opener)
         if idx != -1 and (earliest is None or idx < earliest):
             earliest = idx
     return earliest
 
 
-def _split_at_potential_prefix(text: str) -> tuple[str, str]:
+def _split_at_potential_prefix(text: str, openers: Sequence[str]) -> tuple[str, str]:
     """Split text into (safe_to_stream, potential_opener_prefix).
 
-    The second part is a suffix that could be the beginning of a tool call
-    opening tag, so it must be held back until more text arrives.
+    The second part is a suffix that could be the beginning of one of
+    `openers`, so it must be held back until more text arrives.
     """
-    for opener in TOOL_CALL_OPENERS:
+    for opener in openers:
         for length in range(1, len(opener)):
             if text.endswith(opener[:length]):
                 return text[:-length], text[-length:]
@@ -144,8 +151,10 @@ class ServedCompletingModel(CompletingModel):
         }
 
         pending = ""
+        in_thinking = False
         in_tool_call = False
         tool_call_text = ""
+        openers = [THINKING_OPENER, *TOOL_CALL_OPENERS]
 
         async with httpx.AsyncClient(timeout=None) as client:
             async with client.stream(
@@ -162,22 +171,43 @@ class ServedCompletingModel(CompletingModel):
 
                     pending += text
 
-                    opener_pos = _find_opener(pending)
-                    if opener_pos is not None:
+                    while pending and not in_tool_call:
+                        if in_thinking:
+                            closer_pos = pending.find(THINKING_CLOSER)
+                            if closer_pos == -1:
+                                thought, pending = _split_at_potential_prefix(
+                                    pending, [THINKING_CLOSER]
+                                )
+                                if thought:
+                                    yield CompletionChunk(thinking=thought)
+                                break
+                            if closer_pos:
+                                yield CompletionChunk(thinking=pending[:closer_pos])
+                            pending = pending[closer_pos + len(THINKING_CLOSER) :]
+                            in_thinking = False
+                            continue
+
+                        opener_pos = _find_opener(pending, openers)
+                        if opener_pos is None:
+                            safe, pending = _split_at_potential_prefix(pending, openers)
+                            if safe:
+                                yield CompletionChunk(content=safe)
+                            break
+
                         before = pending[:opener_pos]
                         if before:
                             yield CompletionChunk(content=before)
-                        tool_call_text = pending[opener_pos:]
-                        in_tool_call = True
-                        pending = ""
-                        continue
-
-                    safe, held = _split_at_potential_prefix(pending)
-                    if safe:
-                        yield CompletionChunk(content=safe)
-                    pending = held
-
-        if pending:
+                        if pending.startswith(THINKING_OPENER, opener_pos):
+                            pending = pending[opener_pos + len(THINKING_OPENER) :]
+                            in_thinking = True
+                        else:
+                            tool_call_text = pending[opener_pos:]
+                            in_tool_call = True
+                            pending = ""
+
+        if pending and in_thinking:
+            yield CompletionChunk(thinking=pending)
+        elif pending:
             yield CompletionChunk(content=pending)
 
         if tool_call_text:

{cortexgrid_infer-0.3.3 → cortexgrid_infer-0.3.4}/cortexgrid_infer/serve_apps/base.py

@@ -28,6 +28,8 @@ from cortexgrid_infer.core import DeployedModel
 
 _GIB = float(1 << 30)
 
+ENABLE_THINKING_PARAM = "enable_thinking"
+
 # Weights are only part of what a replica holds. The multiplier covers CUDA
 # context, activations and fragmentation; the floor covers the KV cache of a
 # long context, which does not scale with model size the way activations do.
@@ -94,6 +96,10 @@ class LocalModel:
             vram_gb=max(round(size * _VRAM_HEADROOM + _VRAM_FLOOR_GB, 1), cls.min_vram_gb),
         )
 
+    @classmethod
+    def config(cls) -> dict[str, str]:
+        return {}
+
     @classmethod
     def client(cls, url: str, name: str) -> DeployedModel:
         """A client for this app deployed at `url`, serving the model `name`."""

{cortexgrid_infer-0.3.3 → cortexgrid_infer-0.3.4}/cortexgrid_infer/serve_apps/gemini/base.py

@@ -44,6 +44,7 @@ class GeminiModel(HostedModel):
                 f"{family}/{suffix}/{run_name} is missing {sorted(missing)} from "
                 "its config; set them on the model card"
             )
+        self._config = config
         self._model = config[MODEL_PARAM]
         self._client = genai.Client(
             api_key=cortexgrid.get_secret(config[API_KEY_SECRET_PARAM])

{cortexgrid_infer-0.3.3 → cortexgrid_infer-0.3.4}/cortexgrid_infer/serve_apps/gemini/text2text.py

@@ -16,8 +16,13 @@ from fastapi import FastAPI
 from fastapi.responses import StreamingResponse
 
 from cortexgrid_infer.core import Message, ToolSpec
-from cortexgrid_infer.protocols.completion import ServedCompletingModel, encode_tool_call
-from cortexgrid_infer.serve_apps.gemini.base import GeminiModel
+from cortexgrid_infer.protocols.completion import (
+    ServedCompletingModel,
+    encode_thinking,
+    encode_tool_call,
+)
+from cortexgrid_infer.serve_apps.base import ENABLE_THINKING_PARAM
+from cortexgrid_infer.serve_apps.gemini.base import DEFAULT_API_KEY_SECRET, GeminiModel
 
 _app = FastAPI()
 
@@ -109,6 +114,18 @@ def to_gemini_tools(
 
 @serve.ingress(_app)
 class GeminiText2Text(GeminiModel):
+    @classmethod
+    def config(
+        cls,
+        model_id: str,
+        api_key_secret: str = DEFAULT_API_KEY_SECRET,
+        enable_thinking: str = "true",
+    ) -> dict[str, str]:
+        return {
+            **super().config(model_id, api_key_secret),
+            ENABLE_THINKING_PARAM: enable_thinking,
+        }
+
     @classmethod
     def client(cls, url: str, name: str) -> ServedCompletingModel:
         return ServedCompletingModel(url=url, model_id=name)
@@ -121,6 +138,11 @@ class GeminiText2Text(GeminiModel):
         generate_config: dict[str, Any] = {
             "max_output_tokens": body.get("max_new_tokens", 16 * 1024),
             "temperature": body.get("temperature", 0.7),
+            "thinking_config": (
+                {"include_thoughts": True}
+                if self._config.get(ENABLE_THINKING_PARAM, "true") == "true"
+                else {"thinking_budget": 0}
+            ),
         }
         if system_prompt:
             generate_config["system_instruction"] = system_prompt
@@ -140,7 +162,10 @@ class GeminiText2Text(GeminiModel):
                             part.function_call.name, dict(part.function_call.args or {})
                         ).encode("utf-8")
 
-                    elif part.text and not part.thought:
+                    elif part.text and part.thought:
+                        yield encode_thinking(part.text).encode("utf-8")
+
+                    elif part.text:
                         yield part.text.encode("utf-8")
 
         return StreamingResponse(stream(), media_type="text/plain")

{cortexgrid_infer-0.3.3 → cortexgrid_infer-0.3.4}/cortexgrid_infer/serve_apps/text2text.py

@@ -30,7 +30,7 @@ import torch
 from cortexgrid_infer import compiling
 from cortexgrid_infer.device import detect_device
 from cortexgrid_infer.protocols.completion import ServedCompletingModel
-from cortexgrid_infer.serve_apps.base import LocalModel
+from cortexgrid_infer.serve_apps.base import ENABLE_THINKING_PARAM, LocalModel
 
 
 log = logging.getLogger(__name__)
@@ -60,6 +60,10 @@ _RESERVED_BODY_KEYS = {
 
 @serve.ingress(_app)
 class Text2Text(LocalModel):
+    @classmethod
+    def config(cls) -> dict[str, str]:
+        return {ENABLE_THINKING_PARAM: "true"}
+
     @classmethod
     def client(cls, url: str, name: str) -> ServedCompletingModel:
         return ServedCompletingModel(url=url, model_id=name)
@@ -81,6 +85,7 @@ class Text2Text(LocalModel):
         self._max_total_tokens = min(
             int(settings.get(MAX_TOTAL_TOKENS_PARAM, DEFAULT_MAX_TOTAL_TOKENS)), context
         )
+        self._enable_thinking = settings.get(ENABLE_THINKING_PARAM, "true") == "true"
         # A static cache of fixed length is what makes decode compilable: it
         # pins the shape every forward sees, so the loop is captured once
         # instead of recompiled per token. transformers then compiles `generate`
@@ -133,6 +138,7 @@ class Text2Text(LocalModel):
             tools=tools,
             tokenize=False,
             add_generation_prompt=True,
+            enable_thinking=self._enable_thinking,
         )
         inputs = self._truncate(
             self._tokenizer(prompt, return_tensors="pt"),

{cortexgrid_infer-0.3.3 → cortexgrid_infer-0.3.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cortexgrid-infer"
-version = "0.3.3"
+version = "0.3.4"
 description = "Serve apps per task (text-to-text, text-to-image, image-to-mesh) and importers per source (HuggingFace) for running models on cortexgrid"
 readme = "README.md"
 license = "Apache-2.0"