diff --git a/instill/helpers/__init__.py b/instill/helpers/__init__.py index e69de29..496028f 100644 --- a/instill/helpers/__init__.py +++ b/instill/helpers/__init__.py @@ -0,0 +1,5 @@ +from instill.helpers.protobufs.parse import ( + Metadata, + construct_infer_response, + construct_metadata_response, +) diff --git a/instill/helpers/protobufs/__init__.py b/instill/helpers/protobufs/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/instill/helpers/protobufs/__init__.py @@ -0,0 +1 @@ + diff --git a/instill/helpers/protobufs/parse.py b/instill/helpers/protobufs/parse.py new file mode 100644 index 0000000..8838d0e --- /dev/null +++ b/instill/helpers/protobufs/parse.py @@ -0,0 +1,77 @@ +from dataclasses import dataclass +from typing import List + +from instill.helpers.protobufs.ray_pb2 import ( + ModelMetadataRequest, + ModelMetadataResponse, + RayServiceCallRequest, + RayServiceCallResponse, + InferTensor, +) + + +@dataclass +class Metadata: + name: str + datatype: str + shape: list + + +def construct_metadata_response( + req: ModelMetadataRequest, + inputs: List[Metadata], + outputs: List[Metadata], +) -> ModelMetadataResponse: + resp = ModelMetadataResponse( + name=req.name, + versions=req.version, + framework="python", + inputs=[], + outputs=[], + ) + + for i in inputs: + resp.inputs.append( + ModelMetadataResponse.TensorMetadata( + name=i.name, + datatype=i.datatype, + shape=i.shape, + ) + ) + + for o in outputs: + resp.outputs.append( + ModelMetadataResponse.TensorMetadata( + name=o.name, + datatype=o.datatype, + shape=o.shape, + ) + ) + + return resp + + +def construct_infer_response( + req: RayServiceCallRequest, + outputs: List[Metadata], + raw_outputs, +) -> RayServiceCallResponse: + resp = RayServiceCallResponse( + model_name=req.model_name, + model_version=req.model_version, + outputs=[], + raw_output_contents=[], + ) + + for o in outputs: + resp.outputs.append( + InferTensor( + name=o.name, + datatype=o.datatype, + shape=o.shape, + ) + ) + + resp.raw_output_contents.append(raw_outputs) + + return resp