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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid-infer
-Version: 0.3.3
+Version: 0.3.5
 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
@@ -178,7 +178,23 @@ beyond the weights. (A `Hosted` entry asks for none — it holds no weights.)
 
 ## Compilation
 
-`Text2Text` and `Text2Image` run `torch.compile` over the model they load, because
+Every serve app that runs weights can run them through `torch.compile`, and none
+does unless asked: a replica is eager unless its model card carries
+`compile: "true"`. Each `LocalModel` writes `compile: "false"` onto the card at
+import, so turning it on is an edit of the card and a redeploy:
+
+```python
+config = cortexgrid.model_config(family, suffix, cortexgrid.IMPORTED)
+cortexgrid.set_model_config(family, suffix, cortexgrid.IMPORTED, {**config, "compile": "true"})
+cortexgrid.deploy_model(family, suffix, cortexgrid.IMPORTED, wait=True)
+```
+
+It is off by default because it is not free: every input shape pays a warm-up on
+the request that first brings it, and each CUDA graph holds memory of its own. It
+pays back where traffic keeps to a few shapes and the model is small enough for
+launch overhead to matter. A card asking for it off CUDA is served eager.
+
+Compiled, a serve app runs `torch.compile` over the model it loads, because
 eager PyTorch bills per *operation* and a transformer forward is mostly glue —
 views, broadcasts, permutes, elementwise adds — with only a small fraction of the
 dispatches being the matrix multiplies that do the arithmetic. While the tensors
@@ -197,7 +213,7 @@ fits the work it does:
 | `Mode.GRAPHED` (`reduce-overhead`) | fuses **and** captures a CUDA graph, so a forward replays one graph instead of issuing its launches | a loop of identically shaped forwards, where the capture is reused every iteration. A graph cannot go dynamic, so each distinct shape captures separately |
 | `Mode.FUSED` (`default`) | fuses only | work whose shapes move under it, where there is nothing stable to capture and dynamo is free to compile one dynamic kernel set for all of them |
 
-Both serve apps are tuned rather than defaulted:
+Each serve app is tuned rather than defaulted:
 
 - **`Text2Image`** compiles the denoiser (`transformer` or `unet`)
   and every `text_encoder*` as `Mode.GRAPHED`. A denoising schedule is the ideal case —
@@ -224,12 +240,25 @@ Both serve apps are tuned rather than defaulted:
   chat template puts the turn to answer last — with a warning naming both
   lengths. Serving it whole would resize the cache and recompile the decode loop,
   which costs more than the generation itself.
-
-The image app asks to be configured for nothing, for the same reason neither asks
-which dtype to load in; the completion app takes `max_total_tokens` alone, because
-how much cache a caller needs is the one thing the model cannot tell it. Both
-happen on CUDA and nowhere else — inductor is weakest off it, and a host
-outrunning its accelerator is a GPU problem to begin with.
+- **`TextRewriter`** decodes as `Text2Text` does — a static cache, `Mode.GRAPHED`,
+  transformers compiling the step — with the cache sized to the 512 new tokens a
+  reply may take, and a larger `max_new_tokens` capped to it. An encoder-decoder
+  has a second cache, though: `generate` sizes cross-attention to the encoded
+  text, so each distinct text length would capture the decode loop afresh. The
+  text is therefore padded up to the next of 32, 64, 128, 256 or 512 tokens —
+  one capture per bucket, the padding masked out — and a longer one is cut to
+  512, keeping its start, with a warning naming both lengths.
+- **`Image2Mesh`** cannot know what to compile: each subclass brings its own
+  model. It calls the subclass's `compile(device)` after `load`, which does
+  nothing unless overridden; a subclass compiles its own modules there with
+  `compiling.compile_module`, in the mode that fits them.
+
+Beyond the switch, the image app asks to be configured for nothing, for the same
+reason none asks which dtype to load in; the completion app takes
+`max_total_tokens` alone, because how much cache a caller needs is the one thing
+the model cannot tell it. All of it happens on CUDA and nowhere else — inductor
+is weakest off it, and a host outrunning its accelerator is a GPU problem to
+begin with.
 
 Two things worth knowing:
 
@@ -245,7 +274,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)
@@ -338,10 +367,10 @@ diffusion settings with no Gemini counterpart, and come back as `None`.
 | `Importer` | Base of the importers: identity, scratch directory, `source`, and `requirements()` / `client(url)` taken from the serve app. A new source implements `download(local_dir)` and `weights()`. |
 | `Hosted(model_id, serve_app, **settings)` | Entry for a model hosted elsewhere; no weights. `settings` are the serve app's `config` arguments. |
 | `ModelEntry` | Base of both: `model_id`, `family`, `suffix`, `serve_app`, `requirements()`, `config()`, `client(url)`. |
-| `Text2Text`, `Text2Image`, `Image2Mesh` | Serve apps that run weights (`LocalModel`s). |
+| `Text2Text`, `Text2Image`, `TextRewriter`, `Image2Mesh` | Serve apps that run weights (`LocalModel`s). |
 | `AnthropicText2Text` | Serve app that forwards to the Anthropic API (a `HostedModel`). |
 | `GeminiText2Text`, `GeminiText2Image` | Serve apps that forward to the Gemini API (`HostedModel`s sharing the `GeminiModel` base). |
-| `LocalModel` | Base of the serve apps that run weights: `ignore_patterns()`, `bytes_per_param`, `min_vram_gb`, `requirements(weights)`, `client(url, name)`. |
+| `LocalModel` | Base of the serve apps that run weights: `ignore_patterns()`, `bytes_per_param`, `min_vram_gb`, `requirements(weights)`, `config()` (the model card's defaults, `compile: "false"` among them), `client(url, name)`. |
 | `HostedModel` | Base of the serve apps that forward: `config(model_id, **settings)`, `requirements()`, `client(url, name)`. |
 | `Weights(params=None, file_bytes=0)` | What an importer reads about the weights, handed to the serve app to size a replica. |
 | `complete(model, messages, tools=None, max_new_tokens=2048, temperature=0.7, **kw)` | Async stream of `CompletionChunk` for a `CompletingModel`. |
@@ -370,7 +399,8 @@ Extending it:
   (there is an `encode_tool_call` for upstreams that report them structurally) —
   then `ServedCompletingModel` is its client, unchanged.
 - **A model no task's serve app can load** gets its own serve app — an `Image2Mesh`
-  subclass, say — and, if its source is unusual too, its own importer.
+  subclass, say, with `load` and `make_mesh`, and `compile` if compiling it pays —
+  and, if its source is unusual too, its own importer.
 
 ## Notes
 

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

@@ -150,7 +150,23 @@ beyond the weights. (A `Hosted` entry asks for none — it holds no weights.)
 
 ## Compilation
 
-`Text2Text` and `Text2Image` run `torch.compile` over the model they load, because
+Every serve app that runs weights can run them through `torch.compile`, and none
+does unless asked: a replica is eager unless its model card carries
+`compile: "true"`. Each `LocalModel` writes `compile: "false"` onto the card at
+import, so turning it on is an edit of the card and a redeploy:
+
+```python
+config = cortexgrid.model_config(family, suffix, cortexgrid.IMPORTED)
+cortexgrid.set_model_config(family, suffix, cortexgrid.IMPORTED, {**config, "compile": "true"})
+cortexgrid.deploy_model(family, suffix, cortexgrid.IMPORTED, wait=True)
+```
+
+It is off by default because it is not free: every input shape pays a warm-up on
+the request that first brings it, and each CUDA graph holds memory of its own. It
+pays back where traffic keeps to a few shapes and the model is small enough for
+launch overhead to matter. A card asking for it off CUDA is served eager.
+
+Compiled, a serve app runs `torch.compile` over the model it loads, because
 eager PyTorch bills per *operation* and a transformer forward is mostly glue —
 views, broadcasts, permutes, elementwise adds — with only a small fraction of the
 dispatches being the matrix multiplies that do the arithmetic. While the tensors
@@ -169,7 +185,7 @@ fits the work it does:
 | `Mode.GRAPHED` (`reduce-overhead`) | fuses **and** captures a CUDA graph, so a forward replays one graph instead of issuing its launches | a loop of identically shaped forwards, where the capture is reused every iteration. A graph cannot go dynamic, so each distinct shape captures separately |
 | `Mode.FUSED` (`default`) | fuses only | work whose shapes move under it, where there is nothing stable to capture and dynamo is free to compile one dynamic kernel set for all of them |
 
-Both serve apps are tuned rather than defaulted:
+Each serve app is tuned rather than defaulted:
 
 - **`Text2Image`** compiles the denoiser (`transformer` or `unet`)
   and every `text_encoder*` as `Mode.GRAPHED`. A denoising schedule is the ideal case —
@@ -196,12 +212,25 @@ Both serve apps are tuned rather than defaulted:
   chat template puts the turn to answer last — with a warning naming both
   lengths. Serving it whole would resize the cache and recompile the decode loop,
   which costs more than the generation itself.
-
-The image app asks to be configured for nothing, for the same reason neither asks
-which dtype to load in; the completion app takes `max_total_tokens` alone, because
-how much cache a caller needs is the one thing the model cannot tell it. Both
-happen on CUDA and nowhere else — inductor is weakest off it, and a host
-outrunning its accelerator is a GPU problem to begin with.
+- **`TextRewriter`** decodes as `Text2Text` does — a static cache, `Mode.GRAPHED`,
+  transformers compiling the step — with the cache sized to the 512 new tokens a
+  reply may take, and a larger `max_new_tokens` capped to it. An encoder-decoder
+  has a second cache, though: `generate` sizes cross-attention to the encoded
+  text, so each distinct text length would capture the decode loop afresh. The
+  text is therefore padded up to the next of 32, 64, 128, 256 or 512 tokens —
+  one capture per bucket, the padding masked out — and a longer one is cut to
+  512, keeping its start, with a warning naming both lengths.
+- **`Image2Mesh`** cannot know what to compile: each subclass brings its own
+  model. It calls the subclass's `compile(device)` after `load`, which does
+  nothing unless overridden; a subclass compiles its own modules there with
+  `compiling.compile_module`, in the mode that fits them.
+
+Beyond the switch, the image app asks to be configured for nothing, for the same
+reason none asks which dtype to load in; the completion app takes
+`max_total_tokens` alone, because how much cache a caller needs is the one thing
+the model cannot tell it. All of it happens on CUDA and nowhere else — inductor
+is weakest off it, and a host outrunning its accelerator is a GPU problem to
+begin with.
 
 Two things worth knowing:
 
@@ -217,7 +246,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)
@@ -310,10 +339,10 @@ diffusion settings with no Gemini counterpart, and come back as `None`.
 | `Importer` | Base of the importers: identity, scratch directory, `source`, and `requirements()` / `client(url)` taken from the serve app. A new source implements `download(local_dir)` and `weights()`. |
 | `Hosted(model_id, serve_app, **settings)` | Entry for a model hosted elsewhere; no weights. `settings` are the serve app's `config` arguments. |
 | `ModelEntry` | Base of both: `model_id`, `family`, `suffix`, `serve_app`, `requirements()`, `config()`, `client(url)`. |
-| `Text2Text`, `Text2Image`, `Image2Mesh` | Serve apps that run weights (`LocalModel`s). |
+| `Text2Text`, `Text2Image`, `TextRewriter`, `Image2Mesh` | Serve apps that run weights (`LocalModel`s). |
 | `AnthropicText2Text` | Serve app that forwards to the Anthropic API (a `HostedModel`). |
 | `GeminiText2Text`, `GeminiText2Image` | Serve apps that forward to the Gemini API (`HostedModel`s sharing the `GeminiModel` base). |
-| `LocalModel` | Base of the serve apps that run weights: `ignore_patterns()`, `bytes_per_param`, `min_vram_gb`, `requirements(weights)`, `client(url, name)`. |
+| `LocalModel` | Base of the serve apps that run weights: `ignore_patterns()`, `bytes_per_param`, `min_vram_gb`, `requirements(weights)`, `config()` (the model card's defaults, `compile: "false"` among them), `client(url, name)`. |
 | `HostedModel` | Base of the serve apps that forward: `config(model_id, **settings)`, `requirements()`, `client(url, name)`. |
 | `Weights(params=None, file_bytes=0)` | What an importer reads about the weights, handed to the serve app to size a replica. |
 | `complete(model, messages, tools=None, max_new_tokens=2048, temperature=0.7, **kw)` | Async stream of `CompletionChunk` for a `CompletingModel`. |
@@ -342,7 +371,8 @@ Extending it:
   (there is an `encode_tool_call` for upstreams that report them structurally) —
   then `ServedCompletingModel` is its client, unchanged.
 - **A model no task's serve app can load** gets its own serve app — an `Image2Mesh`
-  subclass, say — and, if its source is unusual too, its own importer.
+  subclass, say, with `load` and `make_mesh`, and `compile` if compiling it pays —
+  and, if its source is unusual too, its own importer.
 
 ## Notes
 

{cortexgrid_infer-0.3.3 → cortexgrid_infer-0.3.5}/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.5}/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.5}/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.5}/cortexgrid_infer/serve_apps/base.py

@@ -28,6 +28,14 @@ from cortexgrid_infer.core import DeployedModel
 
 _GIB = float(1 << 30)
 
+ENABLE_THINKING_PARAM = "enable_thinking"
+
+# Whether a replica runs its model through `torch.compile` (see
+# `cortexgrid_infer.compiling`). Off unless the model card says "true": compiled
+# serving pays a warm-up per input shape, which is worth it only where traffic
+# repeats those shapes.
+COMPILE_PARAM = "compile"
+
 # 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 +102,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 {COMPILE_PARAM: "false"}
+
     @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.5}/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.5}/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.5}/cortexgrid_infer/serve_apps/image2mesh.py

@@ -13,6 +13,7 @@ cortexgrid bundles the subclass's own file, so the model's code travels with it.
         min_vram_gb = 6.0
         def load(self, path, device): ...
         def make_mesh(self, image, **options) -> GeneratedMesh: ...
+        def compile(self, device): ...  # optional; see `Image2Mesh.compile`
 
     imp = HuggingFaceImporter("org/my-mesh-model", MyMesh)
 """
@@ -31,11 +32,12 @@ from fastapi import FastAPI
 from PIL import Image, ImageOps
 import torch
 
+from cortexgrid_infer import compiling
 from cortexgrid_infer.core import GeneratedMesh
 from cortexgrid_infer.device import detect_device
 from cortexgrid_infer.protocols import meshing
 from cortexgrid_infer.protocols.meshing import ServedMeshingModel
-from cortexgrid_infer.serve_apps.base import LocalModel
+from cortexgrid_infer.serve_apps.base import COMPILE_PARAM, LocalModel
 
 
 _app = FastAPI()
@@ -55,11 +57,26 @@ class Image2Mesh(LocalModel):
     def __init__(self, family: str, suffix: str, run_name: str) -> None:
         self.device = detect_device()
         self.load(Path(cortexgrid.load_model(family, suffix, run_name)), self.device)
+        settings = cortexgrid.model_config(family, suffix, run_name)
+        requested = settings.get(COMPILE_PARAM, "false") == "true"
+        self.compiled = requested and compiling.supported(self.device)
+        if self.compiled:
+            self.compile(self.device)
 
     def load(self, path: Path, device: torch.device) -> None:
         """Build the model from its files at `path`, on `device`."""
         raise NotImplementedError
 
+    def compile(self, device: torch.device) -> None:
+        """Compile what `load` built, on `device`.
+
+        Called after `load`, and only when the model card asks for it on a
+        device worth compiling on. Does nothing unless overridden: which of a
+        model's modules pay back compiling, and in which `compiling.Mode`, is
+        the model's to know - a denoising loop of fixed shapes suits
+        `Mode.GRAPHED`, a decoder queried over a grid whose size moves suits
+        `Mode.FUSED`. `compiling.compile_module` does the rest."""
+
     def make_mesh(self, image: Image.Image, **options: Any) -> GeneratedMesh:
         """The mesh of the object in `image`, in the frame `GeneratedMesh`
         documents, with the options the request carried; `params` says what the

{cortexgrid_infer-0.3.3 → cortexgrid_infer-0.3.5}/cortexgrid_infer/serve_apps/text2image.py

@@ -29,7 +29,7 @@ import torch
 from cortexgrid_infer import compiling
 from cortexgrid_infer.device import detect_device
 from cortexgrid_infer.protocols.imaging import ServedGeneratingModel
-from cortexgrid_infer.serve_apps.base import LocalModel
+from cortexgrid_infer.serve_apps.base import COMPILE_PARAM, LocalModel
 
 
 _app = FastAPI()
@@ -95,8 +95,12 @@ class Text2Image(LocalModel):
             vae.enable_tiling()
         self._pipe = pipe
         # A denoising schedule is a loop of identically shaped forwards: capture once, replay.
-        self._compiled = compiling.compile_pipeline(
-            pipe, self._device, compiling.Mode.GRAPHED
+        settings = cortexgrid.model_config(family, suffix, run_name)
+        requested = settings.get(COMPILE_PARAM, "false") == "true"
+        self._compiled = (
+            compiling.compile_pipeline(pipe, self._device, compiling.Mode.GRAPHED)
+            if requested
+            else []
         )
 
     @_app.post("/generate")

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

@@ -30,7 +30,11 @@ 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 (
+    COMPILE_PARAM,
+    ENABLE_THINKING_PARAM,
+    LocalModel,
+)
 
 
 log = logging.getLogger(__name__)
@@ -60,6 +64,10 @@ _RESERVED_BODY_KEYS = {
 
 @serve.ingress(_app)
 class Text2Text(LocalModel):
+    @classmethod
+    def config(cls) -> dict[str, str]:
+        return {**super().config(), ENABLE_THINKING_PARAM: "true"}
+
     @classmethod
     def client(cls, url: str, name: str) -> ServedCompletingModel:
         return ServedCompletingModel(url=url, model_id=name)
@@ -81,11 +89,13 @@ 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`
         # itself, which is why the module is not compiled here as well.
-        self._compiled = compiling.supported(self._device)
+        requested = settings.get(COMPILE_PARAM, "false") == "true"
+        self._compiled = requested and compiling.supported(self._device)
         if self._compiled:
             config = self._model.generation_config
             config.cache_implementation = "static"
@@ -133,6 +143,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.5/cortexgrid_infer/serve_apps/text_rewriter.py

@@ -0,0 +1,107 @@
+"""The text-rewriting serve app: an encoder-decoder LM loaded from the cortexgrid
+registry, answering `POST /rewrite` (see `cortexgrid_infer.protocols.rewriting`).
+
+Loads with `AutoModelForSeq2SeqLM`, so it runs any encoder-decoder model (T5,
+BART, Marian, ...) whose weights are in the transformers layout, whichever
+importer staged them. Any task prefix the model expects (`gec: ` for grammar
+correction, say) is the caller's to put in the text.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import cortexgrid
+from cortexgrid import serve
+from fastapi import FastAPI
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, CompileConfig
+import torch
+
+from cortexgrid_infer import compiling
+from cortexgrid_infer.device import detect_device
+from cortexgrid_infer.protocols.rewriting import ServedRewritingModel
+from cortexgrid_infer.serve_apps.base import COMPILE_PARAM, LocalModel
+
+
+log = logging.getLogger(__name__)
+
+_app = FastAPI()
+
+# Output budget for a request that names none: T5's training context, and
+# generation stops at the end-of-sequence token well before it for most text.
+DEFAULT_MAX_NEW_TOKENS = 512
+
+# Lengths a compiled replica pads a text up to. `generate` sizes the
+# cross-attention cache to the encoded text, so every distinct length would
+# capture the decode loop afresh; rounding up to a power of two bounds that at
+# one capture per bucket, at the cost of a little masked-out padding.
+_INPUT_BUCKETS = (32, 64, 128, 256, 512)
+
+
+def _bucket(length: int) -> int:
+    """The smallest of `_INPUT_BUCKETS` that holds `length` tokens, else the largest."""
+    return next((b for b in _INPUT_BUCKETS if b >= length), _INPUT_BUCKETS[-1])
+
+
+@serve.ingress(_app)
+class TextRewriter(LocalModel):
+    @classmethod
+    def client(cls, url: str, name: str) -> ServedRewritingModel:
+        return ServedRewritingModel(url=url, model_id=name)
+
+    def __init__(self, family: str, suffix: str, run_name: str) -> None:
+        path = cortexgrid.load_model(family, suffix, run_name)
+        self._device = detect_device()
+        self._tokenizer = AutoTokenizer.from_pretrained(path)
+        self._model = AutoModelForSeq2SeqLM.from_pretrained(
+            str(path), torch_dtype=torch.bfloat16
+        )
+        self._model.to(self._device)
+        settings = cortexgrid.model_config(family, suffix, run_name)
+        # As in Text2Text: a static cache pins the decoder's self-attention to
+        # one shape, and transformers compiles the decode step itself. The
+        # cross-attention cache is pinned by padding the text; see `_tokenize`.
+        requested = settings.get(COMPILE_PARAM, "false") == "true"
+        self._compiled = requested and compiling.supported(self._device)
+        if self._compiled:
+            config = self._model.generation_config
+            config.cache_implementation = "static"
+            config.max_cache_len = DEFAULT_MAX_NEW_TOKENS
+            config.compile_config = CompileConfig(mode=compiling.Mode.GRAPHED)
+
+    def _tokenize(self, text: str) -> Any:
+        """The model's inputs for `text`, on the device.
+
+        Compiled, the text is padded up to its bucket, and one longer than the
+        last bucket is cut to it, keeping its start."""
+        if not self._compiled:
+            return self._tokenizer(text, return_tensors="pt").to(self._device)
+        length = len(self._tokenizer(text)["input_ids"])
+        bucket = _bucket(length)
+        if length > bucket:
+            log.warning(
+                "truncating text from %d to %d tokens, the most this replica takes",
+                length, bucket,
+            )
+        return self._tokenizer(
+            text,
+            padding="max_length",
+            truncation=True,
+            max_length=bucket,
+            return_tensors="pt",
+        ).to(self._device)
+
+    @_app.post("/rewrite")
+    async def rewrite(self, body: dict[str, Any]) -> dict[str, str]:
+        generation_options = {"max_new_tokens": DEFAULT_MAX_NEW_TOKENS} | {
+            key: value for key, value in body.items() if key != "text"
+        }
+        if self._compiled:
+            # A reply longer than the cache would reallocate it and recompile.
+            generation_options["max_new_tokens"] = min(
+                generation_options["max_new_tokens"], DEFAULT_MAX_NEW_TOKENS
+            )
+        inputs = self._tokenize(body["text"])
+        output = self._model.generate(**inputs, **generation_options)
+        return {"text": self._tokenizer.decode(output[0], skip_special_tokens=True)}

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

@@ -1,6 +1,6 @@
 [project]
 name = "cortexgrid-infer"
-version = "0.3.3"
+version = "0.3.5"
 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"

cortexgrid_infer-0.3.3/cortexgrid_infer/serve_apps/text_rewriter.py

@@ -1,54 +0,0 @@
-"""The text-rewriting serve app: an encoder-decoder LM loaded from the cortexgrid
-registry, answering `POST /rewrite` (see `cortexgrid_infer.protocols.rewriting`).
-
-Loads with `AutoModelForSeq2SeqLM`, so it runs any encoder-decoder model (T5,
-BART, Marian, ...) whose weights are in the transformers layout, whichever
-importer staged them. Any task prefix the model expects (`gec: ` for grammar
-correction, say) is the caller's to put in the text.
-"""
-
-from __future__ import annotations
-
-from typing import Any
-
-import cortexgrid
-from cortexgrid import serve
-from fastapi import FastAPI
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-import torch
-
-from cortexgrid_infer.device import detect_device
-from cortexgrid_infer.protocols.rewriting import ServedRewritingModel
-from cortexgrid_infer.serve_apps.base import LocalModel
-
-
-_app = FastAPI()
-
-# Output budget for a request that names none: T5's training context, and
-# generation stops at the end-of-sequence token well before it for most text.
-DEFAULT_MAX_NEW_TOKENS = 512
-
-
-@serve.ingress(_app)
-class TextRewriter(LocalModel):
-    @classmethod
-    def client(cls, url: str, name: str) -> ServedRewritingModel:
-        return ServedRewritingModel(url=url, model_id=name)
-
-    def __init__(self, family: str, suffix: str, run_name: str) -> None:
-        path = cortexgrid.load_model(family, suffix, run_name)
-        self._device = detect_device()
-        self._tokenizer = AutoTokenizer.from_pretrained(path)
-        self._model = AutoModelForSeq2SeqLM.from_pretrained(
-            str(path), torch_dtype=torch.bfloat16
-        )
-        self._model.to(self._device)
-
-    @_app.post("/rewrite")
-    async def rewrite(self, body: dict[str, Any]) -> dict[str, str]:
-        generation_options = {"max_new_tokens": DEFAULT_MAX_NEW_TOKENS} | {
-            key: value for key, value in body.items() if key != "text"
-        }
-        inputs = self._tokenizer(body["text"], return_tensors="pt").to(self._device)
-        output = self._model.generate(**inputs, **generation_options)
-        return {"text": self._tokenizer.decode(output[0], skip_special_tokens=True)}