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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid-infer
-Version: 0.3.5
+Version: 0.3.7
 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
@@ -9,7 +9,7 @@ License-File: LICENSE
 Requires-Python: <3.12,>=3.11
 Requires-Dist: accelerate>=0.33
 Requires-Dist: anthropic>=0.40.0
-Requires-Dist: cortexgrid>=0.3.4
+Requires-Dist: cortexgrid>=0.3.12
 Requires-Dist: diffusers>=0.38
 Requires-Dist: fastapi>=0.133
 Requires-Dist: google-genai>=1.0
@@ -132,7 +132,7 @@ async def run():
 asyncio.run(run())
 
 # 4. Teardown.
-cortexgrid.undeploy_model(imp.family, imp.suffix, cortexgrid.IMPORTED)  # free the GPU
+cortexgrid.undeploy_model(deployment.key)  # free the GPU
 cortexgrid.delete_model(imp.family, imp.suffix, cortexgrid.IMPORTED)    # and the weights
 ```
 
@@ -179,20 +179,24 @@ beyond the weights. (A `Hosted` entry asks for none — it holds no weights.)
 ## Compilation
 
 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:
+does unless asked: a replica is eager unless its deployment's config carries
+`compile: "true"`. Each `LocalModel` writes `compile: "false"` onto the model card
+at import, so turning it on is a deployment that asks for it:
 
 ```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)
+cortexgrid.deploy_model(family, suffix, cortexgrid.IMPORTED, wait=True, config={"compile": "true"})
 ```
 
+That is a separate, compiled deployment of the same model, alongside any uncompiled
+one: a model gets one deployment per distinct config, each at its own
+`deployment.url`. `Text2Text`'s `enable_thinking` and `max_total_tokens` are set per
+deployment the same way. Whatever a deployment leaves out comes from the model card,
+so the defaults `config()` writes there at import still apply.
+
 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.
+launch overhead to matter. A deployment 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 —
@@ -233,7 +237,7 @@ Each serve app is tuned rather than defaulted:
   1.2 GB of keys and values re-read per generated token, several times the
   traffic of the model's own weights, to make room for a prompt nobody sent. It
   is pinned instead to the prompt and reply a caller plausibly sends: 4096 tokens,
-  or `max_total_tokens` on the model card, never past the model's own
+  or `max_total_tokens` from the deployment's config, never past the model's own
   `max_position_embeddings`. A reply takes at most half of it, so an
   over-optimistic `max_new_tokens` cannot squeeze the prompt to nothing, and a
   prompt that will not fit in the rest is truncated — keeping the end, since a

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

@@ -104,7 +104,7 @@ async def run():
 asyncio.run(run())
 
 # 4. Teardown.
-cortexgrid.undeploy_model(imp.family, imp.suffix, cortexgrid.IMPORTED)  # free the GPU
+cortexgrid.undeploy_model(deployment.key)  # free the GPU
 cortexgrid.delete_model(imp.family, imp.suffix, cortexgrid.IMPORTED)    # and the weights
 ```
 
@@ -151,20 +151,24 @@ beyond the weights. (A `Hosted` entry asks for none — it holds no weights.)
 ## Compilation
 
 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:
+does unless asked: a replica is eager unless its deployment's config carries
+`compile: "true"`. Each `LocalModel` writes `compile: "false"` onto the model card
+at import, so turning it on is a deployment that asks for it:
 
 ```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)
+cortexgrid.deploy_model(family, suffix, cortexgrid.IMPORTED, wait=True, config={"compile": "true"})
 ```
 
+That is a separate, compiled deployment of the same model, alongside any uncompiled
+one: a model gets one deployment per distinct config, each at its own
+`deployment.url`. `Text2Text`'s `enable_thinking` and `max_total_tokens` are set per
+deployment the same way. Whatever a deployment leaves out comes from the model card,
+so the defaults `config()` writes there at import still apply.
+
 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.
+launch overhead to matter. A deployment 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 —
@@ -205,7 +209,7 @@ Each serve app is tuned rather than defaulted:
   1.2 GB of keys and values re-read per generated token, several times the
   traffic of the model's own weights, to make room for a prompt nobody sent. It
   is pinned instead to the prompt and reply a caller plausibly sends: 4096 tokens,
-  or `max_total_tokens` on the model card, never past the model's own
+  or `max_total_tokens` from the deployment's config, never past the model's own
   `max_position_embeddings`. A reply takes at most half of it, so an
   over-optimistic `max_new_tokens` cannot squeeze the prompt to nothing, and a
   prompt that will not fit in the rest is truncated — keeping the end, since a

{cortexgrid_infer-0.3.5 → cortexgrid_infer-0.3.7}/cortexgrid_infer/core.py

@@ -37,6 +37,27 @@ class CompletionChunk:
     def has_tool_calls(self) -> bool:
         return bool(self.tool_calls)
 
+    @classmethod
+    def from_wire(
+        cls, answered: dict[str, Any], tools_by_name: dict[str, Callable[..., Any]]
+    ) -> CompletionChunk:
+        tool_calls = []
+        for called in answered["tool_calls"]:
+            tool = tools_by_name.get(called["name"])
+            if tool is None:
+                continue
+            bound = partial(tool, **called["arguments"])
+            tool_call = ToolCall(
+                id=called["id"], name=called["name"], arguments=called["arguments"], _func=bound
+            )
+            tool_calls.append(tool_call)
+        return cls(
+            content=answered["content"],
+            thinking=answered["thinking"],
+            tool_calls=tool_calls,
+            finish_reason=answered["finish_reason"],
+        )
+
 
 Message = dict[str, Any]
 ToolSpec = dict[str, Any]

cortexgrid_infer-0.3.7/cortexgrid_infer/protocols/completion.py

@@ -0,0 +1,92 @@
+"""The wire protocol between a cortexgrid-served completion app and its client.
+
+A completion serve app streams plain text from `POST /complete`, with tool calls
+inlined in the text as `<tool_call>{"name": ..., "arguments": {...}}</tool_call>`.
+That is the format an instruct-tuned causal LM emits on its own, so a serve app
+wrapping one forwards its tokens untouched; a serve app wrapping a provider with
+structured tool calls encodes them into the same form with `encode_tool_call`.
+
+Either way the client is the same `ServedCompletingModel`, which streams text
+back as `CompletionChunk`s and reassembles the tool calls.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+import json
+from typing import Any, Sequence
+
+import httpx
+
+from cortexgrid_infer.core import CompletingModel, CompletionChunk, Message, Tool
+from cortexgrid_infer.utils import build_tool_map, normalize_tools
+
+@dataclass
+class ServedCompletingModel(CompletingModel):
+    """Client for a completion app cortexgrid has deployed at `url`.
+
+    Provider-agnostic: every completion serve app speaks the protocol above, so
+    which model is behind it only shows up in `model_id`."""
+
+    url: str
+    model_id: str
+
+    @property
+    def name(self) -> str:
+        return self.model_id
+
+    async def loglikelihoods(
+        self, messages: list[Message], continuations: list[str]
+    ) -> list[float]:
+        body = {"messages": messages, "continuations": continuations}
+        loglikelihoods_url = f"{self.url}/loglikelihoods"
+        async with httpx.AsyncClient(timeout=None) as client:
+            response = await client.post(loglikelihoods_url, json=body)
+        response.raise_for_status()
+        answered = response.json()
+        return answered["loglikelihoods"]
+
+    async def complete(
+        self,
+        messages: list[Message],
+        tools: Sequence[Tool] | None = None,
+        max_new_tokens: int = 16 * 1024,
+        temperature: float = 0.7,
+        top_p: float = 0.9,
+        top_k: int = 50,
+        repetition_penalty: float = 1.0,
+        **kwargs: Any,
+    ) -> AsyncIterator[CompletionChunk]:
+        tool_specs = normalize_tools(tools)
+        tools_by_name = build_tool_map(tools)
+        body: dict[str, Any] = {
+            "messages": messages,
+            "tools": tool_specs,
+            "max_new_tokens": max_new_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+            "top_k": top_k,
+            "repetition_penalty": repetition_penalty,
+            **kwargs,
+        }
+        complete_url = f"{self.url}/complete"
+        async with httpx.AsyncClient(timeout=None) as client:
+            async with client.stream("POST", complete_url, json=body) as response:
+                response.raise_for_status()
+                async for line in response.aiter_lines():
+                    answered = json.loads(line)
+                    chunk = CompletionChunk.from_wire(answered, tools_by_name)
+                    yield chunk
+
+    async def last_hidden_states(
+        self, parts: list[str], continuations_of_each_part: list[list[str]]
+    ) -> AsyncIterator[list[list[float]]]:
+        body = {"parts": parts, "continuations_of_each_part": continuations_of_each_part}
+        last_hidden_states_url = f"{self.url}/last_hidden_states"
+        async with httpx.AsyncClient(timeout=None) as client:
+            async with client.stream("POST", last_hidden_states_url, json=body) as response:
+                response.raise_for_status()
+                async for line in response.aiter_lines():
+                    answered = json.loads(line)
+                    yield answered["last_hidden_states"]

{cortexgrid_infer-0.3.5 → cortexgrid_infer-0.3.7}/cortexgrid_infer/serve_apps/anthropic.py

@@ -19,12 +19,13 @@ from typing import Any
 import cortexgrid
 from anthropic import AsyncAnthropic
 from cortexgrid import serve
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException, status
 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.protocols.completion import ServedCompletingModel
 from cortexgrid_infer.serve_apps.base import HostedModel
+from cortexgrid_infer.serve_apps.streamed_completion import StreamedCompletionWriter
 
 _app = FastAPI()
 
@@ -131,13 +132,14 @@ class AnthropicText2Text(HostedModel):
     def client(cls, url: str, name: str) -> ServedCompletingModel:
         return ServedCompletingModel(url=url, model_id=name)
 
-    def __init__(self, family: str, suffix: str, run_name: str) -> None:
-        config = cortexgrid.model_config(family, suffix, run_name)
+    def __init__(self, deployment: cortexgrid.DeploymentKey) -> None:
+        config = cortexgrid.model_config(deployment)
         missing = {MODEL_PARAM, API_KEY_SECRET_PARAM} - config.keys()
         if missing:
             raise RuntimeError(
-                f"{family}/{suffix}/{run_name} is missing {sorted(missing)} from "
-                "its config; set them on the model card"
+                f"{deployment.family}/{deployment.suffix}/{deployment.run_name} "
+                f"is missing {sorted(missing)} from its config; set them on the "
+                "model card"
             )
         self._model = config[MODEL_PARAM]
         self._client = AsyncAnthropic(
@@ -161,19 +163,30 @@ class AnthropicText2Text(HostedModel):
             create_kwargs["tools"] = tools
 
         async def stream() -> AsyncIterator[bytes]:
+            writer = StreamedCompletionWriter()
             async with self._client.messages.stream(**create_kwargs) as events:
                 async for event in events:
                     if (
                         event.type == "content_block_delta"
                         and event.delta.type == "text_delta"
                     ):
-                        yield event.delta.text.encode("utf-8")
+                        line = writer.read_content(event.delta.text)
+                        yield line
 
                     elif event.type == "content_block_stop":
                         block = events.current_message_snapshot.content[event.index]
                         if block.type == "tool_use":
-                            yield encode_tool_call(
-                                block.name, dict(block.input)
-                            ).encode("utf-8")
-
-        return StreamingResponse(stream(), media_type="text/plain")
+                            arguments = dict(block.input)
+                            line = writer.read_tool_call(block.name, arguments, block.id)
+                            yield line
+            for line in writer.finish():
+                yield line
+
+        return StreamingResponse(stream(), media_type="application/x-ndjson")
+
+    @_app.post("/last_hidden_states")
+    async def last_hidden_states(self, body: dict[str, Any]) -> None:
+        raise HTTPException(
+            status_code=status.HTTP_501_NOT_IMPLEMENTED,
+            detail="Anthropic does not expose its models' hidden states",
+        )

{cortexgrid_infer-0.3.5 → cortexgrid_infer-0.3.7}/cortexgrid_infer/serve_apps/base.py

@@ -31,9 +31,9 @@ _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.
+# `cortexgrid_infer.compiling`). Off unless the deployment's config 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

{cortexgrid_infer-0.3.5 → cortexgrid_infer-0.3.7}/cortexgrid_infer/serve_apps/gemini/base.py

@@ -36,13 +36,13 @@ class GeminiModel(HostedModel):
         holding the key to send it with."""
         return {MODEL_PARAM: model_id, API_KEY_SECRET_PARAM: api_key_secret}
 
-    def __init__(self, family: str, suffix: str, run_name: str) -> None:
-        config = cortexgrid.model_config(family, suffix, run_name)
+    def __init__(self, deployment: cortexgrid.DeploymentKey) -> None:
+        config = cortexgrid.model_config(deployment)
         missing = {MODEL_PARAM, API_KEY_SECRET_PARAM} - config.keys()
         if missing:
             raise RuntimeError(
-                f"{family}/{suffix}/{run_name} is missing {sorted(missing)} from "
-                "its config; set them on the model card"
+                f"{deployment.family}/{deployment.suffix}/{deployment.run_name} is "
+                f"missing {sorted(missing)} from its config; set them on the model card"
             )
         self._config = config
         self._model = config[MODEL_PARAM]

{cortexgrid_infer-0.3.5 → cortexgrid_infer-0.3.7}/cortexgrid_infer/serve_apps/gemini/text2text.py

@@ -12,17 +12,14 @@ from collections.abc import AsyncIterator
 from typing import Any
 
 from cortexgrid import serve
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException, status
 from fastapi.responses import StreamingResponse
 
 from cortexgrid_infer.core import Message, ToolSpec
-from cortexgrid_infer.protocols.completion import (
-    ServedCompletingModel,
-    encode_thinking,
-    encode_tool_call,
-)
+from cortexgrid_infer.protocols.completion import ServedCompletingModel
 from cortexgrid_infer.serve_apps.base import ENABLE_THINKING_PARAM
 from cortexgrid_infer.serve_apps.gemini.base import DEFAULT_API_KEY_SECRET, GeminiModel
+from cortexgrid_infer.serve_apps.streamed_completion import StreamedCompletionWriter
 
 _app = FastAPI()
 
@@ -150,6 +147,7 @@ class GeminiText2Text(GeminiModel):
             generate_config["tools"] = tools
 
         async def stream() -> AsyncIterator[bytes]:
+            writer = StreamedCompletionWriter()
             chunks = await self._client.aio.models.generate_content_stream(
                 model=self._model, contents=contents, config=generate_config
             )
@@ -158,14 +156,27 @@ class GeminiText2Text(GeminiModel):
                     continue
                 for part in chunk.candidates[0].content.parts or []:
                     if part.function_call:
-                        yield encode_tool_call(
-                            part.function_call.name, dict(part.function_call.args or {})
-                        ).encode("utf-8")
+                        arguments = dict(part.function_call.args or {})
+                        line = writer.read_tool_call(
+                            part.function_call.name, arguments, part.function_call.id
+                        )
+                        yield line
 
                     elif part.text and part.thought:
-                        yield encode_thinking(part.text).encode("utf-8")
+                        line = writer.read_thinking(part.text)
+                        yield line
 
                     elif part.text:
-                        yield part.text.encode("utf-8")
-
-        return StreamingResponse(stream(), media_type="text/plain")
+                        line = writer.read_content(part.text)
+                        yield line
+            for line in writer.finish():
+                yield line
+
+        return StreamingResponse(stream(), media_type="application/x-ndjson")
+
+    @_app.post("/last_hidden_states")
+    async def last_hidden_states(self, body: dict[str, Any]) -> None:
+        raise HTTPException(
+            status_code=status.HTTP_501_NOT_IMPLEMENTED,
+            detail="Gemini does not expose its models' hidden states",
+        )

{cortexgrid_infer-0.3.5 → cortexgrid_infer-0.3.7}/cortexgrid_infer/serve_apps/image2mesh.py

@@ -54,10 +54,11 @@ class Image2Mesh(LocalModel):
     def client(cls, url: str, name: str) -> ServedMeshingModel:
         return ServedMeshingModel(url=url, model_id=name)
 
-    def __init__(self, family: str, suffix: str, run_name: str) -> None:
+    def __init__(self, deployment: cortexgrid.DeploymentKey) -> 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)
+        path = cortexgrid.load_model(deployment.family, deployment.suffix, deployment.run_name)
+        self.load(Path(path), self.device)
+        settings = cortexgrid.model_config(deployment)
         requested = settings.get(COMPILE_PARAM, "false") == "true"
         self.compiled = requested and compiling.supported(self.device)
         if self.compiled:

cortexgrid_infer-0.3.7/cortexgrid_infer/serve_apps/streamed_completion.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import itertools
+import json
+import re
+from typing import Any, Sequence
+
+_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),
+    re.compile(r"```tool_call\s*(\{.*?\})\s*```", re.DOTALL),
+]
+
+
+def new_tool_call_id() -> str:
+    number = next(_tool_call_id_counter)
+    return f"call_{number}"
+
+
+def completion_chunk(
+    content: str = "",
+    thinking: str = "",
+    tool_calls: Sequence[dict[str, Any]] = (),
+    finish_reason: str | None = None,
+) -> bytes:
+    chunk = {
+        "content": content,
+        "thinking": thinking,
+        "tool_calls": list(tool_calls),
+        "finish_reason": finish_reason,
+    }
+    line = json.dumps(chunk)
+    return f"{line}\n".encode("utf-8")
+
+
+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 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, 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 one of
+    `openers`, so it must be held back until more text arrives.
+    """
+    for opener in openers:
+        for length in range(1, len(opener)):
+            if text.endswith(opener[:length]):
+                return text[:-length], text[-length:]
+    return text, ""
+
+
+def tool_calls_in(text: str) -> tuple[str, list[dict[str, Any]]]:
+    tool_calls = []
+    clean_text = text
+    for pattern in TOOL_CALL_PATTERNS:
+        matches = pattern.findall(text)
+        for match in matches:
+            try:
+                data = json.loads(match)
+                name = data["name"]
+            except (json.JSONDecodeError, KeyError):
+                continue
+            arguments = data.get("arguments", {})
+            call_id = new_tool_call_id()
+            tool_calls.append({"id": call_id, "name": name, "arguments": arguments})
+        clean_text = pattern.sub("", clean_text)
+    stripped = clean_text.strip()
+    return stripped, tool_calls
+
+
+class StreamedCompletionWriter:
+    def __init__(self) -> None:
+        self._pending = ""
+        self._in_thinking = False
+        self._in_tool_call = False
+        self._tool_call_text = ""
+        self._made_a_tool_call = False
+
+    def read_content(self, text: str) -> bytes:
+        return completion_chunk(content=text)
+
+    def read_thinking(self, text: str) -> bytes:
+        return completion_chunk(thinking=text)
+
+    def read_tool_call(
+        self, name: str, arguments: dict[str, Any], call_id: str | None = None
+    ) -> bytes:
+        self._made_a_tool_call = True
+        tool_call_id = call_id or new_tool_call_id()
+        tool_call = {"id": tool_call_id, "name": name, "arguments": arguments}
+        return completion_chunk(tool_calls=[tool_call])
+
+    def read(self, text: str) -> list[bytes]:
+        if not text:
+            return []
+        if self._in_tool_call:
+            self._tool_call_text += text
+            return []
+        self._pending += text
+        lines: list[bytes] = []
+        openers = [THINKING_OPENER, *TOOL_CALL_OPENERS]
+        while self._pending and not self._in_tool_call:
+            if self._in_thinking:
+                closer_pos = self._pending.find(THINKING_CLOSER)
+                if closer_pos == -1:
+                    thought, self._pending = _split_at_potential_prefix(
+                        self._pending, [THINKING_CLOSER]
+                    )
+                    if thought:
+                        lines.append(completion_chunk(thinking=thought))
+                    break
+                if closer_pos:
+                    thought = self._pending[:closer_pos]
+                    lines.append(completion_chunk(thinking=thought))
+                self._pending = self._pending[closer_pos + len(THINKING_CLOSER) :]
+                self._in_thinking = False
+                continue
+
+            opener_pos = _find_opener(self._pending, openers)
+            if opener_pos is None:
+                safe, self._pending = _split_at_potential_prefix(self._pending, openers)
+                if safe:
+                    lines.append(completion_chunk(content=safe))
+                break
+
+            before = self._pending[:opener_pos]
+            if before:
+                lines.append(completion_chunk(content=before))
+            if self._pending.startswith(THINKING_OPENER, opener_pos):
+                self._pending = self._pending[opener_pos + len(THINKING_OPENER) :]
+                self._in_thinking = True
+            else:
+                self._tool_call_text = self._pending[opener_pos:]
+                self._in_tool_call = True
+                self._pending = ""
+        return lines
+
+    def finish(self) -> list[bytes]:
+        lines: list[bytes] = []
+        if self._pending and self._in_thinking:
+            lines.append(completion_chunk(thinking=self._pending))
+        elif self._pending:
+            lines.append(completion_chunk(content=self._pending))
+        clean_text, tool_calls = tool_calls_in(self._tool_call_text)
+        if clean_text:
+            lines.append(completion_chunk(content=clean_text))
+        if tool_calls:
+            lines.append(completion_chunk(tool_calls=tool_calls))
+        elif not self._made_a_tool_call:
+            lines.append(completion_chunk(finish_reason="stop"))
+        return lines

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

@@ -84,8 +84,8 @@ class Text2Image(LocalModel):
     def client(cls, url: str, name: str) -> ServedGeneratingModel:
         return ServedGeneratingModel(url=url, model_id=name)
 
-    def __init__(self, family: str, suffix: str, run_name: str) -> None:
-        path = cortexgrid.load_model(family, suffix, run_name)
+    def __init__(self, deployment: cortexgrid.DeploymentKey) -> None:
+        path = cortexgrid.load_model(deployment.family, deployment.suffix, deployment.run_name)
         self._device = detect_device()
         pipe = _pipeline_class(path).from_pretrained(path, torch_dtype=torch.bfloat16)
         pipe = pipe.to(self._device)
@@ -95,7 +95,7 @@ class Text2Image(LocalModel):
             vae.enable_tiling()
         self._pipe = pipe
         # A denoising schedule is a loop of identically shaped forwards: capture once, replay.
-        settings = cortexgrid.model_config(family, suffix, run_name)
+        settings = cortexgrid.model_config(deployment)
         requested = settings.get(COMPILE_PARAM, "false") == "true"
         self._compiled = (
             compiling.compile_pipeline(pipe, self._device, compiling.Mode.GRAPHED)

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

@@ -10,8 +10,10 @@ its tokens are forwarded untouched.
 from __future__ import annotations
 
 import asyncio
+import copy
+import json
 import logging
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Callable
 from threading import Thread
 from typing import Any
 
@@ -22,7 +24,9 @@ from fastapi.responses import StreamingResponse
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
+    Cache,
     CompileConfig,
+    PreTrainedModel,
     TextIteratorStreamer,
 )
 import torch
@@ -35,6 +39,7 @@ from cortexgrid_infer.serve_apps.base import (
     ENABLE_THINKING_PARAM,
     LocalModel,
 )
+from cortexgrid_infer.serve_apps.streamed_completion import StreamedCompletionWriter
 
 
 log = logging.getLogger(__name__)
@@ -51,6 +56,9 @@ _app = FastAPI()
 MAX_TOTAL_TOKENS_PARAM = "max_total_tokens"
 DEFAULT_MAX_TOTAL_TOKENS = 4096
 
+DTYPE_PARAM = "dtype"
+DEFAULT_DTYPE = "float16"
+
 _RESERVED_BODY_KEYS = {
     "messages",
     "tools",
@@ -59,9 +67,55 @@ _RESERVED_BODY_KEYS = {
     "top_p",
     "top_k",
     "repetition_penalty",
+    "choices",
 }
 
 
+def allowed_next_tokens(
+    choices: list[list[int]], generated: list[int], end_of_text: int
+) -> list[int]:
+    matching = [choice for choice in choices if choice[: len(generated)] == generated]
+    continuing = {choice[len(generated)] for choice in matching if len(choice) > len(generated)}
+    finished = any(choice == generated for choice in matching)
+    allowed = sorted(continuing | {end_of_text}) if finished else sorted(continuing)
+    return allowed or [end_of_text]
+
+
+def continuation_loglikelihood(
+    model: PreTrainedModel, prompt_ids: list[int], continuation_ids: list[int]
+) -> float:
+    read = torch.tensor([prompt_ids + continuation_ids], device=model.device)
+    with torch.inference_mode():
+        outputs = model(input_ids=read, logits_to_keep=len(continuation_ids) + 1)
+    predicting_the_continuation = outputs.logits[0, :-1].float()
+    log_probabilities = torch.log_softmax(predicting_the_continuation, dim=-1)
+    positions = torch.arange(len(continuation_ids), device=log_probabilities.device)
+    continuation = torch.tensor(continuation_ids, device=log_probabilities.device)
+    return float(log_probabilities[positions, continuation].sum())
+
+
+def read_further(model: PreTrainedModel, read: Cache | None, ids: list[int]) -> Cache:
+    unread = torch.tensor([ids], device=model.device)
+    with torch.inference_mode():
+        outputs = model(input_ids=unread, past_key_values=read, use_cache=True, logits_to_keep=1)
+    return outputs.past_key_values
+
+
+def last_hidden_state_after(model: PreTrainedModel, read: Cache, ids: list[int]) -> torch.Tensor:
+    read_for_the_continuation = copy.deepcopy(read)
+    continuation = torch.tensor([ids], device=model.device)
+    with torch.inference_mode():
+        outputs = model(
+            input_ids=continuation,
+            past_key_values=read_for_the_continuation,
+            use_cache=True,
+            output_hidden_states=True,
+            logits_to_keep=1,
+        )
+    last_layer = outputs.hidden_states[-1]
+    return last_layer[0, -1].float()
+
+
 @serve.ingress(_app)
 class Text2Text(LocalModel):
     @classmethod
@@ -72,17 +126,19 @@ class Text2Text(LocalModel):
     def client(cls, url: str, name: str) -> ServedCompletingModel:
         return ServedCompletingModel(url=url, model_id=name)
 
-    def __init__(self, family: str, suffix: str, run_name: str) -> None:
-        path = cortexgrid.load_model(family, suffix, run_name)
+    def __init__(self, deployment: cortexgrid.DeploymentKey) -> None:
+        path = cortexgrid.load_model(
+            deployment.family, deployment.suffix, deployment.run_name
+        )
+        settings = cortexgrid.model_config(deployment)
         self._device = detect_device()
         self._tokenizer = AutoTokenizer.from_pretrained(path)
         if self._tokenizer.pad_token is None:
             self._tokenizer.pad_token = self._tokenizer.eos_token
         self._model = AutoModelForCausalLM.from_pretrained(
-            str(path), torch_dtype=torch.float16
+            str(path), torch_dtype=getattr(torch, settings.get(DTYPE_PARAM, DEFAULT_DTYPE))
         )
         self._model.to(self._device)  # type: ignore
-        settings = cortexgrid.model_config(family, suffix, run_name)
         context = getattr(
             self._model.config, "max_position_embeddings", DEFAULT_MAX_TOTAL_TOKENS
         )
@@ -125,6 +181,82 @@ class Text2Text(LocalModel):
             inputs[key] = value[..., -limit:]
         return inputs
 
+    def _prompt_ids(self, messages: list[dict[str, Any]]) -> list[int]:
+        prompt = self._tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=self._enable_thinking,
+        )
+        return self._tokenizer(prompt)["input_ids"]
+
+    def _continuation_ids(self, text: str) -> list[int]:
+        return self._tokenizer(text, add_special_tokens=False)["input_ids"]
+
+    def _only_the_choices(
+        self, choices: list[str], prompt_length: int
+    ) -> Callable[[int, torch.Tensor], list[int]]:
+        choice_ids = [self._continuation_ids(choice) for choice in choices]
+        end_of_text = self._tokenizer.eos_token_id
+
+        def allowed(_batch: int, sequence: torch.Tensor) -> list[int]:
+            generated = sequence[prompt_length:].tolist()
+            return allowed_next_tokens(choice_ids, generated, end_of_text)
+
+        return allowed
+
+    def _loglikelihoods(
+        self, messages: list[dict[str, Any]], continuations: list[str]
+    ) -> list[float]:
+        prompt_ids = self._prompt_ids(messages)
+        continuation_ids = [self._continuation_ids(continuation) for continuation in continuations]
+        return [
+            continuation_loglikelihood(self._model, prompt_ids, ids) for ids in continuation_ids
+        ]
+
+    @_app.post("/loglikelihoods")
+    async def loglikelihoods(self, body: dict[str, Any]) -> dict[str, list[float]]:
+        loop = asyncio.get_running_loop()
+        loglikelihoods = await loop.run_in_executor(
+            None, self._loglikelihoods, body["messages"], body["continuations"]
+        )
+        return {"loglikelihoods": loglikelihoods}
+
+    def _read_part(
+        self, read: Cache | None, part: str, continuations: list[str]
+    ) -> tuple[Cache, list[list[float]]]:
+        tokenized = self._tokenizer(part, add_special_tokens=read is None)
+        part_ids = tokenized["input_ids"]
+        read = read_further(self._model, read, part_ids)
+        continuation_ids = [self._continuation_ids(continuation) for continuation in continuations]
+        last_hidden_states = [
+            last_hidden_state_after(self._model, read, ids) for ids in continuation_ids
+        ]
+        answered = [hidden_state.tolist() for hidden_state in last_hidden_states]
+        return read, answered
+
+    async def _last_hidden_states_of_each_part(
+        self, parts: list[str], continuations_of_each_part: list[list[str]]
+    ) -> AsyncIterator[bytes]:
+        loop = asyncio.get_running_loop()
+        read: Cache | None = None
+        for part, continuations in zip(parts, continuations_of_each_part):
+            read, last_hidden_states = await loop.run_in_executor(
+                None, self._read_part, read, part, continuations
+            )
+            answered = json.dumps({"last_hidden_states": last_hidden_states})
+            line = f"{answered}\n".encode("utf-8")
+            yield line
+
+    @_app.post("/last_hidden_states")
+    async def last_hidden_states(self, body: dict[str, Any]) -> StreamingResponse:
+        last_hidden_states_of_each_part = self._last_hidden_states_of_each_part(
+            body["parts"], body["continuations_of_each_part"]
+        )
+        return StreamingResponse(
+            last_hidden_states_of_each_part, media_type="application/x-ndjson"
+        )
+
     @_app.post("/complete")
     async def complete(self, body: dict[str, Any]) -> StreamingResponse:
         messages = body["messages"]
@@ -136,6 +268,7 @@ class Text2Text(LocalModel):
         top_p = body.get("top_p", 0.9)
         top_k = body.get("top_k", 50)
         repetition_penalty = body.get("repetition_penalty", 1.0)
+        choices = body.get("choices")
         extra = {k: v for k, v in body.items() if k not in _RESERVED_BODY_KEYS}
 
         prompt = self._tokenizer.apply_chat_template(
@@ -149,6 +282,9 @@ class Text2Text(LocalModel):
             self._tokenizer(prompt, return_tensors="pt"),
             self._max_total_tokens - max_new_tokens,
         ).to(self._model.device)
+        if choices:
+            prompt_length = inputs["input_ids"].shape[-1]
+            extra["prefix_allowed_tokens_fn"] = self._only_the_choices(choices, prompt_length)
 
         queue: asyncio.Queue[str | None] = asyncio.Queue()
         error: list[BaseException] = []
@@ -190,13 +326,16 @@ class Text2Text(LocalModel):
         loop.run_in_executor(None, generate)
 
         async def stream() -> AsyncIterator[bytes]:
+            writer = StreamedCompletionWriter()
             while True:
                 text = await queue.get()
                 if text is None:
                     break
-                if text:
-                    yield text.encode("utf-8")
+                for line in writer.read(text):
+                    yield line
             if error:
                 raise error[0]
+            for line in writer.finish():
+                yield line
 
-        return StreamingResponse(stream(), media_type="text/plain")
+        return StreamingResponse(stream(), media_type="application/x-ndjson")

{cortexgrid_infer-0.3.5 → cortexgrid_infer-0.3.7}/cortexgrid_infer/serve_apps/text_rewriter.py

@@ -50,15 +50,15 @@ class TextRewriter(LocalModel):
     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)
+    def __init__(self, deployment: cortexgrid.DeploymentKey) -> None:
+        path = cortexgrid.load_model(deployment.family, deployment.suffix, deployment.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)
+        settings = cortexgrid.model_config(deployment)
         # 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`.

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

@@ -1,6 +1,6 @@
 [project]
 name = "cortexgrid-infer"
-version = "0.3.5"
+version = "0.3.7"
 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"
@@ -15,7 +15,9 @@ dependencies = [
     # 0.3.4: jobs return results (remote -> JobFuture), register_model for a
     # model that stages no weights, model_config, and deploy_model skipping a
     # redundant PUT - which is why this library has no reuse guard of its own.
-    "cortexgrid>=0.3.4",
+    # 0.3.12: DeploymentKey - a serve-app is constructed with its deployment's
+    # key and reads that deployment's config with model_config(key).
+    "cortexgrid>=0.3.12",
     "torch>=2.0.0",
     # 4.47: CompileConfig, which the completion serve app hands to `generate`
     # alongside a static cache so transformers compiles the decode step itself.

cortexgrid_infer-0.3.5/cortexgrid_infer/protocols/completion.py

@@ -1,222 +0,0 @@
-"""The wire protocol between a cortexgrid-served completion app and its client.
-
-A completion serve app streams plain text from `POST /complete`, with tool calls
-inlined in the text as `<tool_call>{"name": ..., "arguments": {...}}</tool_call>`.
-That is the format an instruct-tuned causal LM emits on its own, so a serve app
-wrapping one forwards its tokens untouched; a serve app wrapping a provider with
-structured tool calls encodes them into the same form with `encode_tool_call`.
-
-Either way the client is the same `ServedCompletingModel`, which streams text
-back as `CompletionChunk`s and reassembles the tool calls.
-"""
-
-from __future__ import annotations
-
-from collections.abc import AsyncIterator
-from dataclasses import dataclass
-from functools import partial
-import itertools
-import json
-import re
-from typing import Any, Callable, Sequence
-
-import httpx
-
-from cortexgrid_infer.core import (
-    CompletingModel,
-    CompletionChunk,
-    Message,
-    Tool,
-    ToolCall,
-)
-from cortexgrid_infer.utils import build_tool_map, normalize_tools
-
-_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),
-    re.compile(r"```tool_call\s*(\{.*?\})\s*```", re.DOTALL),
-]
-
-
-def encode_tool_call(name: str, arguments: dict[str, Any]) -> str:
-    """Render a tool call in the form the client parses back out of the stream.
-
-    For serve apps whose upstream reports tool calls as structured data rather
-    than as text the model generated - they are re-encoded here so one client
-    handles every completion app."""
-    return f'<tool_call>{json.dumps({"name": name, "arguments": arguments})}</tool_call>'
-
-
-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 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, 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 one of
-    `openers`, so it must be held back until more text arrives.
-    """
-    for opener in openers:
-        for length in range(1, len(opener)):
-            if text.endswith(opener[:length]):
-                return text[:-length], text[-length:]
-    return text, ""
-
-
-def parse_tool_calls(
-    text: str, tool_map: dict[str, Callable[..., Any]]
-) -> tuple[str, list[ToolCall]]:
-    """Parse tool calls from generated text and bind them to their functions."""
-    tool_calls: list[ToolCall] = []
-    clean_text = text
-
-    for pattern in TOOL_CALL_PATTERNS:
-        matches = pattern.findall(text)
-        for match in matches:
-            try:
-                data = json.loads(match)
-                name = data["name"]
-                arguments = data.get("arguments", {})
-                if name in tool_map:
-                    tool_calls.append(
-                        ToolCall(
-                            id=f"call_{next(_tool_call_id_counter)}",
-                            name=name,
-                            arguments=arguments,
-                            _func=partial(tool_map[name], **arguments),
-                        )
-                    )
-            except (json.JSONDecodeError, KeyError):
-                continue
-        clean_text = pattern.sub("", clean_text)
-
-    return clean_text.strip(), tool_calls
-
-
-@dataclass
-class ServedCompletingModel(CompletingModel):
-    """Client for a completion app cortexgrid has deployed at `url`.
-
-    Provider-agnostic: every completion serve app speaks the protocol above, so
-    which model is behind it only shows up in `model_id`."""
-
-    url: str
-    model_id: str
-
-    @property
-    def name(self) -> str:
-        return self.model_id
-
-    async def complete(
-        self,
-        messages: list[Message],
-        tools: Sequence[Tool] | None = None,
-        max_new_tokens: int = 16 * 1024,
-        temperature: float = 0.7,
-        top_p: float = 0.9,
-        top_k: int = 50,
-        repetition_penalty: float = 1.0,
-        **kwargs: Any,
-    ) -> AsyncIterator[CompletionChunk]:
-        tool_specs = normalize_tools(tools)
-        tool_map = build_tool_map(tools)
-
-        body: dict[str, Any] = {
-            "messages": messages,
-            "tools": tool_specs,
-            "max_new_tokens": max_new_tokens,
-            "temperature": temperature,
-            "top_p": top_p,
-            "top_k": top_k,
-            "repetition_penalty": repetition_penalty,
-            **kwargs,
-        }
-
-        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(
-                "POST", f"{self.url}/complete", json=body
-            ) as response:
-                response.raise_for_status()
-                async for text in response.aiter_text():
-                    if not text:
-                        continue
-
-                    if in_tool_call:
-                        tool_call_text += text
-                        continue
-
-                    pending += text
-
-                    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)
-                        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:
-            clean_text, tool_calls = parse_tool_calls(tool_call_text, tool_map)
-            if clean_text:
-                yield CompletionChunk(content=clean_text)
-            if tool_calls:
-                yield CompletionChunk(tool_calls=tool_calls)
-            else:
-                yield CompletionChunk(finish_reason="stop")
-        else:
-            yield CompletionChunk(finish_reason="stop")