Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 6200a71

Browse files
mandelmandel
authored
feat: all to pass directly backend or processor objects to granite-io blocks (#973)
Signed-off-by: Louis Mandel <[email protected]>
1 parent f07f980 commit 6200a71

File tree

5 files changed

+31
-32
lines changed

5 files changed

+31
-32
lines changed

‎pdl-live-react/src/pdl_ast.d.ts‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2354,6 +2354,11 @@ export type PdlContext17 =
23542354
export type PdlId17 = string | null
23552355
export type PdlIsLeaf17 = true
23562356
export type Kind17 = "model"
2357+
/**
2358+
* Model name used by the backend.
2359+
*
2360+
*/
2361+
export type Model = LocalizedExpression | string
23572362
/**
23582363
* Messages to send to the model.
23592364
*
@@ -2411,11 +2416,6 @@ export type Backend =
24112416
| {
24122417
[k: string]: unknown
24132418
}
2414-
/**
2415-
* IO Processor name.
2416-
*
2417-
*/
2418-
export type Processor = LocalizedExpression | string | null
24192419
/**
24202420
* Parameters sent to the model.
24212421
*
@@ -3210,7 +3210,7 @@ export interface GraniteioModelBlock {
32103210
pdl__timing?: PdlTiming | null
32113211
pdl__is_leaf?: PdlIsLeaf17
32123212
kind?: Kind17
3213-
model: unknown
3213+
model: Model
32143214
input?: Input
32153215
modelResponse?: Modelresponse
32163216
/**
@@ -3221,7 +3221,7 @@ export interface GraniteioModelBlock {
32213221
pdl__model_input?: PdlModelInput
32223222
platform?: Platform
32233223
backend: Backend
3224-
processor?: Processor
3224+
processor?: unknown
32253225
parameters?: Parameters
32263226
}
32273227
/**

‎pdl-live-react/src/pdl_ast_utils.ts‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { match, P } from "ts-pattern"
22

3-
import { Backend, PdlBlock, Processor } from "./pdl_ast"
3+
import { Backend, PdlBlock } from "./pdl_ast"
44
import { ExpressionT, isArgs } from "./helpers"
55

66
export function map_block_children(
@@ -74,10 +74,7 @@ export function map_block_children(
7474
: undefined
7575
// @ts-expect-error: f_expr does not preserve the type of the expression
7676
const backend: Backend = f_expr(block.backend)
77-
// @ts-expect-error: f_expr does not preserve the type of the expression
78-
const processor: Processor = block.processor
79-
? f_expr(block.processor)
80-
: undefined
77+
const processor = block.processor ? f_expr(block.processor) : undefined
8178
return {
8279
...block,
8380
model,

‎src/pdl/pdl-schema.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4461,7 +4461,6 @@
44614461
{
44624462
"$ref": "#/$defs/LocalizedExpression_TypeVar_"
44634463
},
4464-
{},
44654464
{
44664465
"type": "string"
44674466
}
@@ -4626,6 +4625,7 @@
46264625
{
46274626
"type": "string"
46284627
},
4628+
{},
46294629
{
46304630
"type": "null"
46314631
}

‎src/pdl/pdl_ast.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -548,13 +548,13 @@ class GraniteioModelBlock(ModelBlock):
548548
platform: Literal[ModelPlatform.GRANITEIO] = ModelPlatform.GRANITEIO
549549
"""Optional field to ensure that the block is using granite-io.
550550
"""
551-
model: ExpressionType[object]
551+
model: ExpressionType[str]
552552
"""Model name used by the backend.
553553
"""
554-
backend: ExpressionType[str | dict[str, Any]]
554+
backend: ExpressionType[str | dict[str, Any | object]]
555555
"""Backend name and configuration.
556556
"""
557-
processor: Optional[ExpressionType[str]] = None
557+
processor: Optional[ExpressionType[str | object]] = None
558558
"""IO Processor name.
559559
"""
560560
parameters: Optional[ExpressionType[dict[str, Any]]] = None

‎src/pdl/pdl_granite_io.py‎

Lines changed: 18 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -19,44 +19,46 @@
1919
class GraniteioModel:
2020
@staticmethod
2121
def processor_of_block(block: GraniteioModelBlock):
22-
model = value_of_expr(block.model)
22+
from granite_io import make_backend, make_io_processor
23+
from granite_io.backend.base import Backend
24+
from granite_io.io import InputOutputProcessor
25+
26+
processor = value_of_expr(block.processor)
27+
if isinstance(processor, InputOutputProcessor):
28+
return processor
29+
model: str = value_of_expr(block.model)
2330
backend = value_of_expr(block.backend)
2431
assert isinstance(model, str), f"The model should be a string: {model}"
25-
assert isinstance(
26-
backend, (dict, str)
27-
), f"The backend should be a string or a dictionary: {backend}"
2832
match backend:
2933
case {"transformers": device}:
30-
assert isinstance(backend, dict)
31-
from granite_io import make_backend
32-
3334
backend = make_backend(
3435
"transformers",
3536
{
3637
"model_name": model,
3738
"device": device,
3839
},
3940
)
40-
case backend_name if isinstance(backend_name, str):
41-
from granite_io import make_backend
42-
41+
case str():
4342
backend = make_backend(
44-
backend_name,
43+
backend,
4544
{
4645
"model_name": model,
4746
},
4847
)
48+
case Backend():
49+
pass
4950
case _:
5051
assert False, f"Unexpected backend: {backend}"
51-
if block.processor is None:
52+
if processor is None:
5253
processor_name = model
5354
else:
54-
processor_name = value_of_expr(block.processor)
55+
assert isinstance(
56+
processor, str
57+
), f"The processor should be a string: {processor}"
58+
processor_name = value_of_expr(processor)
5559
assert isinstance(
5660
processor_name, str
5761
), f"The processor should be a string: {processor_name}"
58-
from granite_io import make_io_processor
59-
6062
io_processor = make_io_processor(processor_name, backend=backend)
6163
return io_processor
6264

@@ -87,7 +89,7 @@ async def async_generate_text(
8789
inputs
8890
)
8991
try: # TODO: update when new version of granite-io is released
90-
message = result.next_message.model_dump()
92+
message = result.next_message.model_dump() # pyright: ignore
9193
except AttributeError:
9294
message = result.results[0].next_message.model_dump()
9395
raw_result = result.model_dump()

0 commit comments

Comments
 (0)