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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (33) hide show
  1. {cortexgrid_infer-0.3.2 → cortexgrid_infer-0.3.3}/PKG-INFO +6 -5
  2. {cortexgrid_infer-0.3.2 → cortexgrid_infer-0.3.3}/README.md +5 -4
  3. {cortexgrid_infer-0.3.2 → cortexgrid_infer-0.3.3}/cortexgrid_infer/__init__.py +17 -6
  4. {cortexgrid_infer-0.3.2 → cortexgrid_infer-0.3.3}/cortexgrid_infer/core.py +19 -0
  5. {cortexgrid_infer-0.3.2 → cortexgrid_infer-0.3.3}/cortexgrid_infer/importers/base.py +1 -1
  6. {cortexgrid_infer-0.3.2 → cortexgrid_infer-0.3.3}/cortexgrid_infer/importers/huggingface.py +1 -1
  7. cortexgrid_infer-0.3.3/cortexgrid_infer/protocols/__init__.py +14 -0
  8. {cortexgrid_infer-0.3.2/cortexgrid_infer → cortexgrid_infer-0.3.3/cortexgrid_infer/protocols}/imaging.py +1 -1
  9. {cortexgrid_infer-0.3.2/cortexgrid_infer → cortexgrid_infer-0.3.3/cortexgrid_infer/protocols}/meshing.py +3 -3
  10. cortexgrid_infer-0.3.3/cortexgrid_infer/protocols/rewriting.py +44 -0
  11. {cortexgrid_infer-0.3.2 → cortexgrid_infer-0.3.3}/cortexgrid_infer/registry.py +1 -1
  12. cortexgrid_infer-0.3.3/cortexgrid_infer/serve_apps/__init__.py +23 -0
  13. {cortexgrid_infer-0.3.2/cortexgrid_infer/models → cortexgrid_infer-0.3.3/cortexgrid_infer/serve_apps}/anthropic.py +2 -2
  14. cortexgrid_infer-0.3.3/cortexgrid_infer/serve_apps/gemini/__init__.py +11 -0
  15. {cortexgrid_infer-0.3.2/cortexgrid_infer/models → cortexgrid_infer-0.3.3/cortexgrid_infer/serve_apps}/gemini/base.py +1 -1
  16. {cortexgrid_infer-0.3.2/cortexgrid_infer/models → cortexgrid_infer-0.3.3/cortexgrid_infer/serve_apps}/gemini/text2image.py +3 -3
  17. {cortexgrid_infer-0.3.2/cortexgrid_infer/models → cortexgrid_infer-0.3.3/cortexgrid_infer/serve_apps}/gemini/text2text.py +2 -2
  18. {cortexgrid_infer-0.3.2/cortexgrid_infer/models → cortexgrid_infer-0.3.3/cortexgrid_infer/serve_apps}/image2mesh.py +5 -5
  19. {cortexgrid_infer-0.3.2/cortexgrid_infer/models → cortexgrid_infer-0.3.3/cortexgrid_infer/serve_apps}/text2image.py +3 -3
  20. {cortexgrid_infer-0.3.2/cortexgrid_infer/models → cortexgrid_infer-0.3.3/cortexgrid_infer/serve_apps}/text2text.py +2 -2
  21. cortexgrid_infer-0.3.3/cortexgrid_infer/serve_apps/text_rewriter.py +54 -0
  22. {cortexgrid_infer-0.3.2 → cortexgrid_infer-0.3.3}/pyproject.toml +3 -3
  23. cortexgrid_infer-0.3.2/cortexgrid_infer/models/__init__.py +0 -21
  24. cortexgrid_infer-0.3.2/cortexgrid_infer/models/gemini/__init__.py +0 -11
  25. {cortexgrid_infer-0.3.2 → cortexgrid_infer-0.3.3}/.gitignore +0 -0
  26. {cortexgrid_infer-0.3.2 → cortexgrid_infer-0.3.3}/LICENSE +0 -0
  27. {cortexgrid_infer-0.3.2 → cortexgrid_infer-0.3.3}/cortexgrid_infer/compiling.py +0 -0
  28. {cortexgrid_infer-0.3.2 → cortexgrid_infer-0.3.3}/cortexgrid_infer/device.py +0 -0
  29. {cortexgrid_infer-0.3.2 → cortexgrid_infer-0.3.3}/cortexgrid_infer/importers/__init__.py +0 -0
  30. {cortexgrid_infer-0.3.2/cortexgrid_infer → cortexgrid_infer-0.3.3/cortexgrid_infer/protocols}/completion.py +0 -0
  31. {cortexgrid_infer-0.3.2 → cortexgrid_infer-0.3.3}/cortexgrid_infer/py.typed +0 -0
  32. {cortexgrid_infer-0.3.2/cortexgrid_infer/models → cortexgrid_infer-0.3.3/cortexgrid_infer/serve_apps}/base.py +0 -0
  33. {cortexgrid_infer-0.3.2 → cortexgrid_infer-0.3.3}/cortexgrid_infer/utils.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: cortexgrid-infer
3
- Version: 0.3.2
3
+ Version: 0.3.3
4
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
5
5
  Project-URL: Homepage, https://github.com/robodatalab/cortexgrid-infer
6
6
  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
49
49
 
50
50
  A model is two things, and the library keeps them apart:
51
51
 
52
- - **What runs it** — a serve app in `cortexgrid_infer.models`, named for the task it
53
- runs, not for where the model came from. It loads the weights, answers the task's
54
- routes, and names the client that speaks them. One serve app runs every model of
55
- its task, whichever source staged it.
52
+ - **What runs it** — a serve app in `cortexgrid_infer.serve_apps`, named for the task
53
+ it runs, not for where the model came from. It loads the weights, answers the task's
54
+ routes (defined in `cortexgrid_infer.protocols`), and names the client that speaks
55
+ them. One serve app runs every model of its task, whichever source staged it.
56
56
  - **Where its weights come from** — an importer in `cortexgrid_infer.importers`, one
57
57
  per source. It names the model in the registry, downloads the weights, and sizes
58
58
  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:
63
63
  | `Text2Text` | text to text | any causal LM in the transformers layout, via `AutoModelForCausalLM` | `ServedCompletingModel` (`complete`) |
64
64
  | `Text2Image` | text (and image) to image | any diffusers pipeline | `ServedGeneratingModel` (`generate`) |
65
65
  | `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) |
66
+ | `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) |
66
67
  | `AnthropicText2Text` | text to text | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
67
68
  | `GeminiText2Text` | text to text | the Gemini API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
68
69
  | `GeminiText2Image` | text to image | the Gemini API, forwarded from the cluster | `ServedGeneratingModel` (`generate`) |
@@ -21,10 +21,10 @@ Upgrading from 0.2.x? See [MIGRATION.md](https://github.com/robodatalab/cortexgr
21
21
 
22
22
  A model is two things, and the library keeps them apart:
23
23
 
24
- - **What runs it** — a serve app in `cortexgrid_infer.models`, named for the task it
25
- runs, not for where the model came from. It loads the weights, answers the task's
26
- routes, and names the client that speaks them. One serve app runs every model of
27
- its task, whichever source staged it.
24
+ - **What runs it** — a serve app in `cortexgrid_infer.serve_apps`, named for the task
25
+ it runs, not for where the model came from. It loads the weights, answers the task's
26
+ routes (defined in `cortexgrid_infer.protocols`), and names the client that speaks
27
+ them. One serve app runs every model of its task, whichever source staged it.
28
28
  - **Where its weights come from** — an importer in `cortexgrid_infer.importers`, one
29
29
  per source. It names the model in the registry, downloads the weights, and sizes
30
30
  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:
35
35
  | `Text2Text` | text to text | any causal LM in the transformers layout, via `AutoModelForCausalLM` | `ServedCompletingModel` (`complete`) |
36
36
  | `Text2Image` | text (and image) to image | any diffusers pipeline | `ServedGeneratingModel` (`generate`) |
37
37
  | `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) |
38
+ | `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) |
38
39
  | `AnthropicText2Text` | text to text | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
39
40
  | `GeminiText2Text` | text to text | the Gemini API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
40
41
  | `GeminiText2Image` | text to image | the Gemini API, forwarded from the cluster | `ServedGeneratingModel` (`generate`) |
@@ -7,8 +7,9 @@ nothing else, split along the two things a model is:
7
7
 
8
8
  - where its weights come from - an importer per source (`importers`), which
9
9
  names, fetches and sizes them;
10
- - what runs it - a serve app per task (`models`), which loads the weights,
11
- answers the task's routes, and names the client that speaks them.
10
+ - what runs it - a serve app per task (`serve_apps`), which loads the weights,
11
+ answers the task's routes (`protocols`), and names the client that speaks
12
+ them.
12
13
 
13
14
  imp = cortexgrid_infer.HuggingFaceImporter("Qwen/Qwen2.5-0.5B-Instruct", cortexgrid_infer.Text2Text)
14
15
  with imp:
@@ -37,18 +38,23 @@ from cortexgrid_infer.core import (
37
38
  GeneratingModel,
38
39
  MeshingModel,
39
40
  Message,
41
+ RewritingModel,
40
42
  Tool,
41
43
  ToolCall,
42
44
  ToolSpec,
43
45
  complete,
44
46
  generate,
45
47
  mesh,
48
+ rewrite,
49
+ )
50
+ from cortexgrid_infer.protocols import (
51
+ ServedCompletingModel,
52
+ ServedGeneratingModel,
53
+ ServedMeshingModel,
54
+ ServedRewritingModel,
46
55
  )
47
- from cortexgrid_infer.completion import ServedCompletingModel
48
- from cortexgrid_infer.imaging import ServedGeneratingModel
49
- from cortexgrid_infer.meshing import ServedMeshingModel
50
56
  from cortexgrid_infer.device import detect_device
51
- from cortexgrid_infer.models import (
57
+ from cortexgrid_infer.serve_apps import (
52
58
  AnthropicText2Text,
53
59
  GeminiText2Image,
54
60
  GeminiText2Text,
@@ -57,6 +63,7 @@ from cortexgrid_infer.models import (
57
63
  LocalModel,
58
64
  Text2Image,
59
65
  Text2Text,
66
+ TextRewriter,
60
67
  Weights,
61
68
  )
62
69
  from cortexgrid_infer.registry import Hosted, ModelEntry, split_model_id
@@ -70,6 +77,7 @@ __all__ = [
70
77
  "GeneratedMesh",
71
78
  "GeneratingModel",
72
79
  "MeshingModel",
80
+ "RewritingModel",
73
81
  "Message",
74
82
  "Tool",
75
83
  "ToolCall",
@@ -77,17 +85,20 @@ __all__ = [
77
85
  "complete",
78
86
  "generate",
79
87
  "mesh",
88
+ "rewrite",
80
89
  "ModelDeployFailed",
81
90
  "detect_device",
82
91
  "ServedCompletingModel",
83
92
  "ServedGeneratingModel",
84
93
  "ServedMeshingModel",
94
+ "ServedRewritingModel",
85
95
  "LocalModel",
86
96
  "HostedModel",
87
97
  "Weights",
88
98
  "Text2Text",
89
99
  "Text2Image",
90
100
  "Image2Mesh",
101
+ "TextRewriter",
91
102
  "AnthropicText2Text",
92
103
  "GeminiText2Text",
93
104
  "GeminiText2Image",
@@ -109,6 +109,15 @@ class MeshingModel(DeployedModel):
109
109
  async def mesh(self, image: bytes, **kwargs: Any) -> GeneratedMesh: ...
110
110
 
111
111
 
112
+ class RewritingModel(DeployedModel):
113
+ """A deployed model that rewrites one text into another - corrects,
114
+ paraphrases, summarises or translates it. Like generation, a single
115
+ request/response."""
116
+
117
+ @abc.abstractmethod
118
+ async def rewrite(self, text: str, max_new_tokens: int = 512, **kwargs: Any) -> str: ...
119
+
120
+
112
121
  async def complete(
113
122
  deployed_model: CompletingModel,
114
123
  messages: list[Message],
@@ -142,3 +151,13 @@ async def generate(
142
151
  async def mesh(deployed_model: MeshingModel, image: bytes, **kwargs: Any) -> GeneratedMesh:
143
152
  """Single image-to-mesh request to the deployed model."""
144
153
  return await deployed_model.mesh(image, **kwargs)
154
+
155
+
156
+ async def rewrite(
157
+ deployed_model: RewritingModel,
158
+ text: str,
159
+ max_new_tokens: int = 512,
160
+ **kwargs: Any,
161
+ ) -> str:
162
+ """Single rewriting request to the deployed model."""
163
+ return await deployed_model.rewrite(text, max_new_tokens=max_new_tokens, **kwargs)
@@ -30,8 +30,8 @@ from typing import Self
30
30
  import cortexgrid
31
31
 
32
32
  from cortexgrid_infer.core import DeployedModel
33
- from cortexgrid_infer.models.base import LocalModel, Weights
34
33
  from cortexgrid_infer.registry import ModelEntry, split_model_id
34
+ from cortexgrid_infer.serve_apps.base import LocalModel, Weights
35
35
 
36
36
 
37
37
  class Importer(ModelEntry):
@@ -15,7 +15,7 @@ from huggingface_hub import HfApi, get_safetensors_metadata, snapshot_download
15
15
  from huggingface_hub.errors import NotASafetensorsRepoError
16
16
 
17
17
  from cortexgrid_infer.importers.base import Importer
18
- from cortexgrid_infer.models.base import LocalModel, Weights
18
+ from cortexgrid_infer.serve_apps.base import LocalModel, Weights
19
19
 
20
20
  # Files that are weights rather than configs, tokenizers or docs. Used for the
21
21
  # fallback estimate, where a repo has no top-level safetensors index.
@@ -0,0 +1,14 @@
1
+ """The wire protocols between a serve app and its client, one per task, each with
2
+ the client that speaks it."""
3
+
4
+ from cortexgrid_infer.protocols.completion import ServedCompletingModel
5
+ from cortexgrid_infer.protocols.imaging import ServedGeneratingModel
6
+ from cortexgrid_infer.protocols.meshing import ServedMeshingModel
7
+ from cortexgrid_infer.protocols.rewriting import ServedRewritingModel
8
+
9
+ __all__ = [
10
+ "ServedCompletingModel",
11
+ "ServedGeneratingModel",
12
+ "ServedMeshingModel",
13
+ "ServedRewritingModel",
14
+ ]
@@ -7,7 +7,7 @@ as null, take the app's defaults. The reply carries the picture as a base64 PNG
7
7
  under `image`, its `width` and `height`, the options the app resolved, and
8
8
  `duration_s`.
9
9
 
10
- `cortexgrid_infer.models.text2image.Text2Image` speaks it for any diffusers
10
+ `cortexgrid_infer.serve_apps.text2image.Text2Image` speaks it for any diffusers
11
11
  pipeline, so one client, `ServedGeneratingModel`, serves every model of the
12
12
  task whoever made it.
13
13
  """
@@ -7,9 +7,9 @@ the mesh as base64 little-endian arrays - `vertices` float32 x, y, z, `faces`
7
7
  int32 indices, `colours` uint8 r, g, b per vertex - in the frame
8
8
  `GeneratedMesh` documents, and `params`, what the model resolved.
9
9
 
10
- Every such app speaks it through `cortexgrid_infer.models.image2mesh.Image2Mesh`,
11
- so one client, `ServedMeshingModel`, serves every model of the task whoever
12
- made it.
10
+ Every such app speaks it through
11
+ `cortexgrid_infer.serve_apps.image2mesh.Image2Mesh`, so one client,
12
+ `ServedMeshingModel`, serves every model of the task whoever made it.
13
13
  """
14
14
 
15
15
  from __future__ import annotations
@@ -0,0 +1,44 @@
1
+ """The wire protocol between a cortexgrid-served rewriting app and its client.
2
+
3
+ A rewriting serve app answers `POST /rewrite`: the request carries the text to
4
+ rewrite under `text`, plus whatever generation options the model takes
5
+ (`max_new_tokens`, `num_beams`, say) as further fields. The reply carries the
6
+ rewritten text under `text`. A request without `max_new_tokens` takes the app's
7
+ default: left to transformers, it would be 20 tokens, cutting all but the
8
+ shortest text.
9
+
10
+ Rewriting is one text in, one text out - correcting grammar, paraphrasing,
11
+ summarising, translating - with no conversation, tools or streaming, which is
12
+ what an encoder-decoder model does. One client, `ServedRewritingModel`, serves
13
+ every model of the task whoever made it.
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ from dataclasses import dataclass
19
+ from typing import Any
20
+
21
+ import httpx
22
+
23
+ from cortexgrid_infer.core import RewritingModel
24
+
25
+
26
+ @dataclass
27
+ class ServedRewritingModel(RewritingModel):
28
+ """The client of any rewriting serve app deployed at `url`."""
29
+
30
+ url: str
31
+ model_id: str
32
+
33
+ @property
34
+ def name(self) -> str:
35
+ return self.model_id
36
+
37
+ async def rewrite(
38
+ self, text: str, max_new_tokens: int = 512, **generation_options: Any
39
+ ) -> str:
40
+ body = {"text": text, "max_new_tokens": max_new_tokens, **generation_options}
41
+ async with httpx.AsyncClient(timeout=None) as client:
42
+ response = await client.post(f"{self.url}/rewrite", json=body)
43
+ response.raise_for_status()
44
+ return response.json()["text"]
@@ -28,7 +28,7 @@ import abc
28
28
  import cortexgrid
29
29
 
30
30
  from cortexgrid_infer.core import DeployedModel
31
- from cortexgrid_infer.models.base import HostedModel
31
+ from cortexgrid_infer.serve_apps.base import HostedModel
32
32
 
33
33
 
34
34
  def split_model_id(model_id: str) -> tuple[str, str]:
@@ -0,0 +1,23 @@
1
+ """The serve apps cortexgrid runs, one per task, each working with a model of
2
+ that task whichever source it came from."""
3
+
4
+ from cortexgrid_infer.serve_apps.anthropic import AnthropicText2Text
5
+ from cortexgrid_infer.serve_apps.base import HostedModel, LocalModel, Weights
6
+ from cortexgrid_infer.serve_apps.gemini import GeminiText2Image, GeminiText2Text
7
+ from cortexgrid_infer.serve_apps.image2mesh import Image2Mesh
8
+ from cortexgrid_infer.serve_apps.text2image import Text2Image
9
+ from cortexgrid_infer.serve_apps.text2text import Text2Text
10
+ from cortexgrid_infer.serve_apps.text_rewriter import TextRewriter
11
+
12
+ __all__ = [
13
+ "LocalModel",
14
+ "HostedModel",
15
+ "Weights",
16
+ "Text2Text",
17
+ "Text2Image",
18
+ "Image2Mesh",
19
+ "TextRewriter",
20
+ "AnthropicText2Text",
21
+ "GeminiText2Text",
22
+ "GeminiText2Image",
23
+ ]
@@ -22,9 +22,9 @@ from cortexgrid import serve
22
22
  from fastapi import FastAPI
23
23
  from fastapi.responses import StreamingResponse
24
24
 
25
- from cortexgrid_infer.completion import ServedCompletingModel, encode_tool_call
26
25
  from cortexgrid_infer.core import Message, ToolSpec
27
- from cortexgrid_infer.models.base import HostedModel
26
+ from cortexgrid_infer.protocols.completion import ServedCompletingModel, encode_tool_call
27
+ from cortexgrid_infer.serve_apps.base import HostedModel
28
28
 
29
29
  _app = FastAPI()
30
30
 
@@ -0,0 +1,11 @@
1
+ """The serve apps that forward to the Gemini API, one per task."""
2
+
3
+ from cortexgrid_infer.serve_apps.gemini.base import GeminiModel
4
+ from cortexgrid_infer.serve_apps.gemini.text2image import GeminiText2Image
5
+ from cortexgrid_infer.serve_apps.gemini.text2text import GeminiText2Text
6
+
7
+ __all__ = [
8
+ "GeminiModel",
9
+ "GeminiText2Text",
10
+ "GeminiText2Image",
11
+ ]
@@ -12,7 +12,7 @@ from __future__ import annotations
12
12
  import cortexgrid
13
13
  from google import genai
14
14
 
15
- from cortexgrid_infer.models.base import HostedModel
15
+ from cortexgrid_infer.serve_apps.base import HostedModel
16
16
 
17
17
  # Keys the registry entry must carry for the deployment to reach the API.
18
18
  MODEL_PARAM = "model"
@@ -1,7 +1,7 @@
1
1
  """The text-to-image serve app that forwards to the Gemini API.
2
2
 
3
3
  It answers the same `POST /generate` route as `Text2Image` (see
4
- `cortexgrid_infer.imaging`), so the same client talks to it. It takes a prompt
4
+ `cortexgrid_infer.protocols.imaging`), so the same client talks to it. It takes a prompt
5
5
  only. Of the sampling options, `size` and `seed` have a Gemini counterpart;
6
6
  `steps` and `guidance` are diffusion settings, and are reported back as unused.
7
7
  """
@@ -18,8 +18,8 @@ from fastapi import FastAPI, HTTPException
18
18
  from google.genai import types
19
19
  from PIL import Image
20
20
 
21
- from cortexgrid_infer.imaging import ServedGeneratingModel
22
- from cortexgrid_infer.models.gemini.base import GeminiModel
21
+ from cortexgrid_infer.protocols.imaging import ServedGeneratingModel
22
+ from cortexgrid_infer.serve_apps.gemini.base import GeminiModel
23
23
 
24
24
  _app = FastAPI()
25
25
 
@@ -15,9 +15,9 @@ from cortexgrid import serve
15
15
  from fastapi import FastAPI
16
16
  from fastapi.responses import StreamingResponse
17
17
 
18
- from cortexgrid_infer.completion import ServedCompletingModel, encode_tool_call
19
18
  from cortexgrid_infer.core import Message, ToolSpec
20
- from cortexgrid_infer.models.gemini.base import GeminiModel
19
+ from cortexgrid_infer.protocols.completion import ServedCompletingModel, encode_tool_call
20
+ from cortexgrid_infer.serve_apps.gemini.base import GeminiModel
21
21
 
22
22
  _app = FastAPI()
23
23
 
@@ -6,8 +6,8 @@ There is no one loader for image-to-mesh models the way `AutoModelForCausalLM`
6
6
  or a diffusers `model_index.json` serves a whole family of repos: each model
7
7
  ships its own code. So a model's serve app subclasses `Image2Mesh` and supplies
8
8
  `load` and `make_mesh`; the subclass inherits the route, whose protocol is
9
- `cortexgrid_infer.meshing`'s, and the client that speaks it. cortexgrid bundles
10
- the subclass's own file, so the model's code travels with it.
9
+ `cortexgrid_infer.protocols.meshing`'s, and the client that speaks it.
10
+ cortexgrid bundles the subclass's own file, so the model's code travels with it.
11
11
 
12
12
  class MyMesh(Image2Mesh):
13
13
  min_vram_gb = 6.0
@@ -31,11 +31,11 @@ from fastapi import FastAPI
31
31
  from PIL import Image, ImageOps
32
32
  import torch
33
33
 
34
- from cortexgrid_infer import meshing
35
34
  from cortexgrid_infer.core import GeneratedMesh
36
35
  from cortexgrid_infer.device import detect_device
37
- from cortexgrid_infer.meshing import ServedMeshingModel
38
- from cortexgrid_infer.models.base import LocalModel
36
+ from cortexgrid_infer.protocols import meshing
37
+ from cortexgrid_infer.protocols.meshing import ServedMeshingModel
38
+ from cortexgrid_infer.serve_apps.base import LocalModel
39
39
 
40
40
 
41
41
  _app = FastAPI()
@@ -1,5 +1,5 @@
1
1
  """The text-to-image serve app: a diffusers pipeline loaded from the cortexgrid
2
- registry, answering `POST /generate` (see `cortexgrid_infer.imaging`).
2
+ registry, answering `POST /generate` (see `cortexgrid_infer.protocols.imaging`).
3
3
 
4
4
  The pipeline class is not hardcoded: it is read from the model's
5
5
  ``model_index.json`` (``_class_name``) and resolved against ``diffusers``, so
@@ -28,8 +28,8 @@ import torch
28
28
 
29
29
  from cortexgrid_infer import compiling
30
30
  from cortexgrid_infer.device import detect_device
31
- from cortexgrid_infer.imaging import ServedGeneratingModel
32
- from cortexgrid_infer.models.base import LocalModel
31
+ from cortexgrid_infer.protocols.imaging import ServedGeneratingModel
32
+ from cortexgrid_infer.serve_apps.base import LocalModel
33
33
 
34
34
 
35
35
  _app = FastAPI()
@@ -28,9 +28,9 @@ from transformers import (
28
28
  import torch
29
29
 
30
30
  from cortexgrid_infer import compiling
31
- from cortexgrid_infer.completion import ServedCompletingModel
32
31
  from cortexgrid_infer.device import detect_device
33
- from cortexgrid_infer.models.base import LocalModel
32
+ from cortexgrid_infer.protocols.completion import ServedCompletingModel
33
+ from cortexgrid_infer.serve_apps.base import LocalModel
34
34
 
35
35
 
36
36
  log = logging.getLogger(__name__)
@@ -0,0 +1,54 @@
1
+ """The text-rewriting serve app: an encoder-decoder LM loaded from the cortexgrid
2
+ registry, answering `POST /rewrite` (see `cortexgrid_infer.protocols.rewriting`).
3
+
4
+ Loads with `AutoModelForSeq2SeqLM`, so it runs any encoder-decoder model (T5,
5
+ BART, Marian, ...) whose weights are in the transformers layout, whichever
6
+ importer staged them. Any task prefix the model expects (`gec: ` for grammar
7
+ correction, say) is the caller's to put in the text.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ from typing import Any
13
+
14
+ import cortexgrid
15
+ from cortexgrid import serve
16
+ from fastapi import FastAPI
17
+ from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
18
+ import torch
19
+
20
+ from cortexgrid_infer.device import detect_device
21
+ from cortexgrid_infer.protocols.rewriting import ServedRewritingModel
22
+ from cortexgrid_infer.serve_apps.base import LocalModel
23
+
24
+
25
+ _app = FastAPI()
26
+
27
+ # Output budget for a request that names none: T5's training context, and
28
+ # generation stops at the end-of-sequence token well before it for most text.
29
+ DEFAULT_MAX_NEW_TOKENS = 512
30
+
31
+
32
+ @serve.ingress(_app)
33
+ class TextRewriter(LocalModel):
34
+ @classmethod
35
+ def client(cls, url: str, name: str) -> ServedRewritingModel:
36
+ return ServedRewritingModel(url=url, model_id=name)
37
+
38
+ def __init__(self, family: str, suffix: str, run_name: str) -> None:
39
+ path = cortexgrid.load_model(family, suffix, run_name)
40
+ self._device = detect_device()
41
+ self._tokenizer = AutoTokenizer.from_pretrained(path)
42
+ self._model = AutoModelForSeq2SeqLM.from_pretrained(
43
+ str(path), torch_dtype=torch.bfloat16
44
+ )
45
+ self._model.to(self._device)
46
+
47
+ @_app.post("/rewrite")
48
+ async def rewrite(self, body: dict[str, Any]) -> dict[str, str]:
49
+ generation_options = {"max_new_tokens": DEFAULT_MAX_NEW_TOKENS} | {
50
+ key: value for key, value in body.items() if key != "text"
51
+ }
52
+ inputs = self._tokenizer(body["text"], return_tensors="pt").to(self._device)
53
+ output = self._model.generate(**inputs, **generation_options)
54
+ return {"text": self._tokenizer.decode(output[0], skip_special_tokens=True)}
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "cortexgrid-infer"
3
- version = "0.3.2"
3
+ version = "0.3.3"
4
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"
5
5
  readme = "README.md"
6
6
  license = "Apache-2.0"
@@ -22,7 +22,7 @@ dependencies = [
22
22
  "transformers>=4.47.0",
23
23
  "huggingface_hub>=0.20.0",
24
24
  "anthropic>=0.40.0",
25
- # models.gemini. The unified Google Gen AI SDK, not the deprecated
25
+ # serve_apps.gemini. The unified Google Gen AI SDK, not the deprecated
26
26
  # google-generativeai.
27
27
  "google-genai>=1.0",
28
28
  "python-dotenv>=1.0.0",
@@ -30,7 +30,7 @@ dependencies = [
30
30
  # The serve-apps run on the cluster's FastAPI, which Ray 2.58's Serve
31
31
  # requires to be >=0.133.
32
32
  "fastapi>=0.133",
33
- # Text-to-image serving (models.text2image). FLUX.2 pipelines land in
33
+ # Text-to-image serving (serve_apps.text2image). FLUX.2 pipelines land in
34
34
  # diffusers 0.38; the Qwen3 text encoder needs sentencepiece + protobuf.
35
35
  "diffusers>=0.38",
36
36
  "pillow>=10.0",
@@ -1,21 +0,0 @@
1
- """The serve apps cortexgrid runs, one per task, each working with a model of
2
- that task whichever source it came from."""
3
-
4
- from cortexgrid_infer.models.anthropic import AnthropicText2Text
5
- from cortexgrid_infer.models.base import HostedModel, LocalModel, Weights
6
- from cortexgrid_infer.models.gemini import GeminiText2Image, GeminiText2Text
7
- from cortexgrid_infer.models.image2mesh import Image2Mesh
8
- from cortexgrid_infer.models.text2image import Text2Image
9
- from cortexgrid_infer.models.text2text import Text2Text
10
-
11
- __all__ = [
12
- "LocalModel",
13
- "HostedModel",
14
- "Weights",
15
- "Text2Text",
16
- "Text2Image",
17
- "Image2Mesh",
18
- "AnthropicText2Text",
19
- "GeminiText2Text",
20
- "GeminiText2Image",
21
- ]
@@ -1,11 +0,0 @@
1
- """The serve apps that forward to the Gemini API, one per task."""
2
-
3
- from cortexgrid_infer.models.gemini.base import GeminiModel
4
- from cortexgrid_infer.models.gemini.text2image import GeminiText2Image
5
- from cortexgrid_infer.models.gemini.text2text import GeminiText2Text
6
-
7
- __all__ = [
8
- "GeminiModel",
9
- "GeminiText2Text",
10
- "GeminiText2Image",
11
- ]