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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid-infer
-Version: 0.3.2
+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
@@ -49,10 +49,10 @@ Upgrading from 0.2.x? See [MIGRATION.md](https://github.com/robodatalab/cortexgr
 
 A model is two things, and the library keeps them apart:
 
-- **What runs it** — a serve app in `cortexgrid_infer.models`, named for the task it
-  runs, not for where the model came from. It loads the weights, answers the task's
-  routes, and names the client that speaks them. One serve app runs every model of
-  its task, whichever source staged it.
+- **What runs it** — a serve app in `cortexgrid_infer.serve_apps`, named for the task
+  it runs, not for where the model came from. It loads the weights, answers the task's
+  routes (defined in `cortexgrid_infer.protocols`), and names the client that speaks
+  them. One serve app runs every model of its task, whichever source staged it.
 - **Where its weights come from** — an importer in `cortexgrid_infer.importers`, one
   per source. It names the model in the registry, downloads the weights, and sizes
   them from the source's metadata. It is handed the serve app, so one importer takes
@@ -63,6 +63,7 @@ A model is two things, and the library keeps them apart:
 | `Text2Text` | text to text | any causal LM in the transformers layout, via `AutoModelForCausalLM` | `ServedCompletingModel` (`complete`) |
 | `Text2Image` | text (and image) to image | any diffusers pipeline | `ServedGeneratingModel` (`generate`) |
 | `Image2Mesh` — a base: subclass it with the model's `load` and `make_mesh` | image to mesh | the model the subclass loads, behind the task's `POST /mesh` | `ServedMeshingModel` (`mesh`: one picture in, a `GeneratedMesh` out) |
+| `TextRewriter` | text to text, rewritten | any encoder-decoder LM (T5, BART, Marian, ...) in the transformers layout, via `AutoModelForSeq2SeqLM` | `ServedRewritingModel` (`rewrite`: one text in, one text out) |
 | `AnthropicText2Text` | text to text | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
 | `GeminiText2Text` | text to text | the Gemini API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
 | `GeminiText2Image` | text to image | the Gemini API, forwarded from the cluster | `ServedGeneratingModel` (`generate`) |
@@ -244,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.2 → cortexgrid_infer-0.3.4}/README.md

@@ -21,10 +21,10 @@ Upgrading from 0.2.x? See [MIGRATION.md](https://github.com/robodatalab/cortexgr
 
 A model is two things, and the library keeps them apart:
 
-- **What runs it** — a serve app in `cortexgrid_infer.models`, named for the task it
-  runs, not for where the model came from. It loads the weights, answers the task's
-  routes, and names the client that speaks them. One serve app runs every model of
-  its task, whichever source staged it.
+- **What runs it** — a serve app in `cortexgrid_infer.serve_apps`, named for the task
+  it runs, not for where the model came from. It loads the weights, answers the task's
+  routes (defined in `cortexgrid_infer.protocols`), and names the client that speaks
+  them. One serve app runs every model of its task, whichever source staged it.
 - **Where its weights come from** — an importer in `cortexgrid_infer.importers`, one
   per source. It names the model in the registry, downloads the weights, and sizes
   them from the source's metadata. It is handed the serve app, so one importer takes
@@ -35,6 +35,7 @@ A model is two things, and the library keeps them apart:
 | `Text2Text` | text to text | any causal LM in the transformers layout, via `AutoModelForCausalLM` | `ServedCompletingModel` (`complete`) |
 | `Text2Image` | text (and image) to image | any diffusers pipeline | `ServedGeneratingModel` (`generate`) |
 | `Image2Mesh` — a base: subclass it with the model's `load` and `make_mesh` | image to mesh | the model the subclass loads, behind the task's `POST /mesh` | `ServedMeshingModel` (`mesh`: one picture in, a `GeneratedMesh` out) |
+| `TextRewriter` | text to text, rewritten | any encoder-decoder LM (T5, BART, Marian, ...) in the transformers layout, via `AutoModelForSeq2SeqLM` | `ServedRewritingModel` (`rewrite`: one text in, one text out) |
 | `AnthropicText2Text` | text to text | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
 | `GeminiText2Text` | text to text | the Gemini API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
 | `GeminiText2Image` | text to image | the Gemini API, forwarded from the cluster | `ServedGeneratingModel` (`generate`) |
@@ -216,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.2 → cortexgrid_infer-0.3.4}/cortexgrid_infer/__init__.py

@@ -7,8 +7,9 @@ nothing else, split along the two things a model is:
 
 - where its weights come from - an importer per source (`importers`), which
   names, fetches and sizes them;
-- what runs it - a serve app per task (`models`), which loads the weights,
-  answers the task's routes, and names the client that speaks them.
+- what runs it - a serve app per task (`serve_apps`), which loads the weights,
+  answers the task's routes (`protocols`), and names the client that speaks
+  them.
 
     imp = cortexgrid_infer.HuggingFaceImporter("Qwen/Qwen2.5-0.5B-Instruct", cortexgrid_infer.Text2Text)
     with imp:
@@ -37,18 +38,23 @@ from cortexgrid_infer.core import (
     GeneratingModel,
     MeshingModel,
     Message,
+    RewritingModel,
     Tool,
     ToolCall,
     ToolSpec,
     complete,
     generate,
     mesh,
+    rewrite,
+)
+from cortexgrid_infer.protocols import (
+    ServedCompletingModel,
+    ServedGeneratingModel,
+    ServedMeshingModel,
+    ServedRewritingModel,
 )
-from cortexgrid_infer.completion import ServedCompletingModel
-from cortexgrid_infer.imaging import ServedGeneratingModel
-from cortexgrid_infer.meshing import ServedMeshingModel
 from cortexgrid_infer.device import detect_device
-from cortexgrid_infer.models import (
+from cortexgrid_infer.serve_apps import (
     AnthropicText2Text,
     GeminiText2Image,
     GeminiText2Text,
@@ -57,6 +63,7 @@ from cortexgrid_infer.models import (
     LocalModel,
     Text2Image,
     Text2Text,
+    TextRewriter,
     Weights,
 )
 from cortexgrid_infer.registry import Hosted, ModelEntry, split_model_id
@@ -70,6 +77,7 @@ __all__ = [
     "GeneratedMesh",
     "GeneratingModel",
     "MeshingModel",
+    "RewritingModel",
     "Message",
     "Tool",
     "ToolCall",
@@ -77,17 +85,20 @@ __all__ = [
     "complete",
     "generate",
     "mesh",
+    "rewrite",
     "ModelDeployFailed",
     "detect_device",
     "ServedCompletingModel",
     "ServedGeneratingModel",
     "ServedMeshingModel",
+    "ServedRewritingModel",
     "LocalModel",
     "HostedModel",
     "Weights",
     "Text2Text",
     "Text2Image",
     "Image2Mesh",
+    "TextRewriter",
     "AnthropicText2Text",
     "GeminiText2Text",
     "GeminiText2Image",

{cortexgrid_infer-0.3.2 → 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
 
@@ -109,6 +110,15 @@ class MeshingModel(DeployedModel):
     async def mesh(self, image: bytes, **kwargs: Any) -> GeneratedMesh: ...
 
 
+class RewritingModel(DeployedModel):
+    """A deployed model that rewrites one text into another - corrects,
+    paraphrases, summarises or translates it. Like generation, a single
+    request/response."""
+
+    @abc.abstractmethod
+    async def rewrite(self, text: str, max_new_tokens: int = 512, **kwargs: Any) -> str: ...
+
+
 async def complete(
     deployed_model: CompletingModel,
     messages: list[Message],
@@ -142,3 +152,13 @@ async def generate(
 async def mesh(deployed_model: MeshingModel, image: bytes, **kwargs: Any) -> GeneratedMesh:
     """Single image-to-mesh request to the deployed model."""
     return await deployed_model.mesh(image, **kwargs)
+
+
+async def rewrite(
+    deployed_model: RewritingModel,
+    text: str,
+    max_new_tokens: int = 512,
+    **kwargs: Any,
+) -> str:
+    """Single rewriting request to the deployed model."""
+    return await deployed_model.rewrite(text, max_new_tokens=max_new_tokens, **kwargs)

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

@@ -30,8 +30,8 @@ from typing import Self
 import cortexgrid
 
 from cortexgrid_infer.core import DeployedModel
-from cortexgrid_infer.models.base import LocalModel, Weights
 from cortexgrid_infer.registry import ModelEntry, split_model_id
+from cortexgrid_infer.serve_apps.base import LocalModel, Weights
 
 
 class Importer(ModelEntry):
@@ -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.2 → cortexgrid_infer-0.3.4}/cortexgrid_infer/importers/huggingface.py

@@ -15,7 +15,7 @@ from huggingface_hub import HfApi, get_safetensors_metadata, snapshot_download
 from huggingface_hub.errors import NotASafetensorsRepoError
 
 from cortexgrid_infer.importers.base import Importer
-from cortexgrid_infer.models.base import LocalModel, Weights
+from cortexgrid_infer.serve_apps.base import LocalModel, Weights
 
 # Files that are weights rather than configs, tokenizers or docs. Used for the
 # fallback estimate, where a repo has no top-level safetensors index.

cortexgrid_infer-0.3.4/cortexgrid_infer/protocols/__init__.py

@@ -0,0 +1,14 @@
+"""The wire protocols between a serve app and its client, one per task, each with
+the client that speaks it."""
+
+from cortexgrid_infer.protocols.completion import ServedCompletingModel
+from cortexgrid_infer.protocols.imaging import ServedGeneratingModel
+from cortexgrid_infer.protocols.meshing import ServedMeshingModel
+from cortexgrid_infer.protocols.rewriting import ServedRewritingModel
+
+__all__ = [
+    "ServedCompletingModel",
+    "ServedGeneratingModel",
+    "ServedMeshingModel",
+    "ServedRewritingModel",
+]

{cortexgrid_infer-0.3.2/cortexgrid_infer → 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.2/cortexgrid_infer → cortexgrid_infer-0.3.4/cortexgrid_infer/protocols}/imaging.py

@@ -7,7 +7,7 @@ as null, take the app's defaults. The reply carries the picture as a base64 PNG
 under `image`, its `width` and `height`, the options the app resolved, and
 `duration_s`.
 
-`cortexgrid_infer.models.text2image.Text2Image` speaks it for any diffusers
+`cortexgrid_infer.serve_apps.text2image.Text2Image` speaks it for any diffusers
 pipeline, so one client, `ServedGeneratingModel`, serves every model of the
 task whoever made it.
 """

{cortexgrid_infer-0.3.2/cortexgrid_infer → cortexgrid_infer-0.3.4/cortexgrid_infer/protocols}/meshing.py

@@ -7,9 +7,9 @@ the mesh as base64 little-endian arrays - `vertices` float32 x, y, z, `faces`
 int32 indices, `colours` uint8 r, g, b per vertex - in the frame
 `GeneratedMesh` documents, and `params`, what the model resolved.
 
-Every such app speaks it through `cortexgrid_infer.models.image2mesh.Image2Mesh`,
-so one client, `ServedMeshingModel`, serves every model of the task whoever
-made it.
+Every such app speaks it through
+`cortexgrid_infer.serve_apps.image2mesh.Image2Mesh`, so one client,
+`ServedMeshingModel`, serves every model of the task whoever made it.
 """
 
 from __future__ import annotations

cortexgrid_infer-0.3.4/cortexgrid_infer/protocols/rewriting.py

@@ -0,0 +1,44 @@
+"""The wire protocol between a cortexgrid-served rewriting app and its client.
+
+A rewriting serve app answers `POST /rewrite`: the request carries the text to
+rewrite under `text`, plus whatever generation options the model takes
+(`max_new_tokens`, `num_beams`, say) as further fields. The reply carries the
+rewritten text under `text`. A request without `max_new_tokens` takes the app's
+default: left to transformers, it would be 20 tokens, cutting all but the
+shortest text.
+
+Rewriting is one text in, one text out - correcting grammar, paraphrasing,
+summarising, translating - with no conversation, tools or streaming, which is
+what an encoder-decoder model does. One client, `ServedRewritingModel`, serves
+every model of the task whoever made it.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import httpx
+
+from cortexgrid_infer.core import RewritingModel
+
+
+@dataclass
+class ServedRewritingModel(RewritingModel):
+    """The client of any rewriting serve app deployed at `url`."""
+
+    url: str
+    model_id: str
+
+    @property
+    def name(self) -> str:
+        return self.model_id
+
+    async def rewrite(
+        self, text: str, max_new_tokens: int = 512, **generation_options: Any
+    ) -> str:
+        body = {"text": text, "max_new_tokens": max_new_tokens, **generation_options}
+        async with httpx.AsyncClient(timeout=None) as client:
+            response = await client.post(f"{self.url}/rewrite", json=body)
+            response.raise_for_status()
+            return response.json()["text"]

{cortexgrid_infer-0.3.2 → cortexgrid_infer-0.3.4}/cortexgrid_infer/registry.py

@@ -28,7 +28,7 @@ import abc
 import cortexgrid
 
 from cortexgrid_infer.core import DeployedModel
-from cortexgrid_infer.models.base import HostedModel
+from cortexgrid_infer.serve_apps.base import HostedModel
 
 
 def split_model_id(model_id: str) -> tuple[str, str]:

cortexgrid_infer-0.3.4/cortexgrid_infer/serve_apps/__init__.py

@@ -0,0 +1,23 @@
+"""The serve apps cortexgrid runs, one per task, each working with a model of
+that task whichever source it came from."""
+
+from cortexgrid_infer.serve_apps.anthropic import AnthropicText2Text
+from cortexgrid_infer.serve_apps.base import HostedModel, LocalModel, Weights
+from cortexgrid_infer.serve_apps.gemini import GeminiText2Image, GeminiText2Text
+from cortexgrid_infer.serve_apps.image2mesh import Image2Mesh
+from cortexgrid_infer.serve_apps.text2image import Text2Image
+from cortexgrid_infer.serve_apps.text2text import Text2Text
+from cortexgrid_infer.serve_apps.text_rewriter import TextRewriter
+
+__all__ = [
+    "LocalModel",
+    "HostedModel",
+    "Weights",
+    "Text2Text",
+    "Text2Image",
+    "Image2Mesh",
+    "TextRewriter",
+    "AnthropicText2Text",
+    "GeminiText2Text",
+    "GeminiText2Image",
+]

{cortexgrid_infer-0.3.2/cortexgrid_infer/models → cortexgrid_infer-0.3.4/cortexgrid_infer/serve_apps}/anthropic.py

@@ -22,9 +22,9 @@ from cortexgrid import serve
 from fastapi import FastAPI
 from fastapi.responses import StreamingResponse
 
-from cortexgrid_infer.completion import ServedCompletingModel, encode_tool_call
 from cortexgrid_infer.core import Message, ToolSpec
-from cortexgrid_infer.models.base import HostedModel
+from cortexgrid_infer.protocols.completion import ServedCompletingModel, encode_tool_call
+from cortexgrid_infer.serve_apps.base import HostedModel
 
 _app = FastAPI()
 

{cortexgrid_infer-0.3.2/cortexgrid_infer/models → 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.4/cortexgrid_infer/serve_apps/gemini/__init__.py

@@ -0,0 +1,11 @@
+"""The serve apps that forward to the Gemini API, one per task."""
+
+from cortexgrid_infer.serve_apps.gemini.base import GeminiModel
+from cortexgrid_infer.serve_apps.gemini.text2image import GeminiText2Image
+from cortexgrid_infer.serve_apps.gemini.text2text import GeminiText2Text
+
+__all__ = [
+    "GeminiModel",
+    "GeminiText2Text",
+    "GeminiText2Image",
+]

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

@@ -12,7 +12,7 @@ from __future__ import annotations
 import cortexgrid
 from google import genai
 
-from cortexgrid_infer.models.base import HostedModel
+from cortexgrid_infer.serve_apps.base import HostedModel
 
 # Keys the registry entry must carry for the deployment to reach the API.
 MODEL_PARAM = "model"
@@ -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.2/cortexgrid_infer/models → cortexgrid_infer-0.3.4/cortexgrid_infer/serve_apps}/gemini/text2image.py

@@ -1,7 +1,7 @@
 """The text-to-image serve app that forwards to the Gemini API.
 
 It answers the same `POST /generate` route as `Text2Image` (see
-`cortexgrid_infer.imaging`), so the same client talks to it. It takes a prompt
+`cortexgrid_infer.protocols.imaging`), so the same client talks to it. It takes a prompt
 only. Of the sampling options, `size` and `seed` have a Gemini counterpart;
 `steps` and `guidance` are diffusion settings, and are reported back as unused.
 """
@@ -18,8 +18,8 @@ from fastapi import FastAPI, HTTPException
 from google.genai import types
 from PIL import Image
 
-from cortexgrid_infer.imaging import ServedGeneratingModel
-from cortexgrid_infer.models.gemini.base import GeminiModel
+from cortexgrid_infer.protocols.imaging import ServedGeneratingModel
+from cortexgrid_infer.serve_apps.gemini.base import GeminiModel
 
 _app = FastAPI()
 

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

@@ -15,9 +15,14 @@ from cortexgrid import serve
 from fastapi import FastAPI
 from fastapi.responses import StreamingResponse
 
-from cortexgrid_infer.completion import ServedCompletingModel, encode_tool_call
 from cortexgrid_infer.core import Message, ToolSpec
-from cortexgrid_infer.models.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.2/cortexgrid_infer/models → cortexgrid_infer-0.3.4/cortexgrid_infer/serve_apps}/image2mesh.py

@@ -6,8 +6,8 @@ There is no one loader for image-to-mesh models the way `AutoModelForCausalLM`
 or a diffusers `model_index.json` serves a whole family of repos: each model
 ships its own code. So a model's serve app subclasses `Image2Mesh` and supplies
 `load` and `make_mesh`; the subclass inherits the route, whose protocol is
-`cortexgrid_infer.meshing`'s, and the client that speaks it. cortexgrid bundles
-the subclass's own file, so the model's code travels with it.
+`cortexgrid_infer.protocols.meshing`'s, and the client that speaks it.
+cortexgrid bundles the subclass's own file, so the model's code travels with it.
 
     class MyMesh(Image2Mesh):
         min_vram_gb = 6.0
@@ -31,11 +31,11 @@ from fastapi import FastAPI
 from PIL import Image, ImageOps
 import torch
 
-from cortexgrid_infer import meshing
 from cortexgrid_infer.core import GeneratedMesh
 from cortexgrid_infer.device import detect_device
-from cortexgrid_infer.meshing import ServedMeshingModel
-from cortexgrid_infer.models.base import LocalModel
+from cortexgrid_infer.protocols import meshing
+from cortexgrid_infer.protocols.meshing import ServedMeshingModel
+from cortexgrid_infer.serve_apps.base import LocalModel
 
 
 _app = FastAPI()

{cortexgrid_infer-0.3.2/cortexgrid_infer/models → cortexgrid_infer-0.3.4/cortexgrid_infer/serve_apps}/text2image.py

@@ -1,5 +1,5 @@
 """The text-to-image serve app: a diffusers pipeline loaded from the cortexgrid
-registry, answering `POST /generate` (see `cortexgrid_infer.imaging`).
+registry, answering `POST /generate` (see `cortexgrid_infer.protocols.imaging`).
 
 The pipeline class is not hardcoded: it is read from the model's
 ``model_index.json`` (``_class_name``) and resolved against ``diffusers``, so
@@ -28,8 +28,8 @@ import torch
 
 from cortexgrid_infer import compiling
 from cortexgrid_infer.device import detect_device
-from cortexgrid_infer.imaging import ServedGeneratingModel
-from cortexgrid_infer.models.base import LocalModel
+from cortexgrid_infer.protocols.imaging import ServedGeneratingModel
+from cortexgrid_infer.serve_apps.base import LocalModel
 
 
 _app = FastAPI()

{cortexgrid_infer-0.3.2/cortexgrid_infer/models → cortexgrid_infer-0.3.4/cortexgrid_infer/serve_apps}/text2text.py

@@ -28,9 +28,9 @@ from transformers import (
 import torch
 
 from cortexgrid_infer import compiling
-from cortexgrid_infer.completion import ServedCompletingModel
 from cortexgrid_infer.device import detect_device
-from cortexgrid_infer.models.base import LocalModel
+from cortexgrid_infer.protocols.completion import ServedCompletingModel
+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.4/cortexgrid_infer/serve_apps/text_rewriter.py

@@ -0,0 +1,54 @@
+"""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)}

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

@@ -1,6 +1,6 @@
 [project]
 name = "cortexgrid-infer"
-version = "0.3.2"
+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"
@@ -22,7 +22,7 @@ dependencies = [
     "transformers>=4.47.0",
     "huggingface_hub>=0.20.0",
     "anthropic>=0.40.0",
-    # models.gemini. The unified Google Gen AI SDK, not the deprecated
+    # serve_apps.gemini. The unified Google Gen AI SDK, not the deprecated
     # google-generativeai.
     "google-genai>=1.0",
     "python-dotenv>=1.0.0",
@@ -30,7 +30,7 @@ dependencies = [
     # The serve-apps run on the cluster's FastAPI, which Ray 2.58's Serve
     # requires to be >=0.133.
     "fastapi>=0.133",
-    # Text-to-image serving (models.text2image). FLUX.2 pipelines land in
+    # Text-to-image serving (serve_apps.text2image). FLUX.2 pipelines land in
     # diffusers 0.38; the Qwen3 text encoder needs sentencepiece + protobuf.
     "diffusers>=0.38",
     "pillow>=10.0",

cortexgrid_infer-0.3.2/cortexgrid_infer/models/__init__.py

@@ -1,21 +0,0 @@
-"""The serve apps cortexgrid runs, one per task, each working with a model of
-that task whichever source it came from."""
-
-from cortexgrid_infer.models.anthropic import AnthropicText2Text
-from cortexgrid_infer.models.base import HostedModel, LocalModel, Weights
-from cortexgrid_infer.models.gemini import GeminiText2Image, GeminiText2Text
-from cortexgrid_infer.models.image2mesh import Image2Mesh
-from cortexgrid_infer.models.text2image import Text2Image
-from cortexgrid_infer.models.text2text import Text2Text
-
-__all__ = [
-    "LocalModel",
-    "HostedModel",
-    "Weights",
-    "Text2Text",
-    "Text2Image",
-    "Image2Mesh",
-    "AnthropicText2Text",
-    "GeminiText2Text",
-    "GeminiText2Image",
-]

cortexgrid_infer-0.3.2/cortexgrid_infer/models/gemini/__init__.py

@@ -1,11 +0,0 @@
-"""The serve apps that forward to the Gemini API, one per task."""
-
-from cortexgrid_infer.models.gemini.base import GeminiModel
-from cortexgrid_infer.models.gemini.text2image import GeminiText2Image
-from cortexgrid_infer.models.gemini.text2text import GeminiText2Text
-
-__all__ = [
-    "GeminiModel",
-    "GeminiText2Text",
-    "GeminiText2Image",
-]