add tests

Author: pepesi

.github/workflows/tests-my.yml

@@ -2,11 +2,19 @@
 from .image import PILImageCodec
 from .json import HuggingfaceSingleJSONCodec
 from .jsonlist import HuggingfaceListJSONCodec
+from .numpylist import NumpyListCodec
+from .conversation import HuggingfaceConversationCodec
+from .raw import RawCodec
+from .utils import EqualUtil
 
 __all__ = [
-    MultiInputRequestCodec,
-    HuggingfaceRequestCodec,
-    PILImageCodec,
-    HuggingfaceSingleJSONCodec,
-    HuggingfaceListJSONCodec,
+    "MultiInputRequestCodec",
+    "HuggingfaceRequestCodec",
+    "PILImageCodec",
+    "HuggingfaceSingleJSONCodec",
+    "HuggingfaceListJSONCodec",
+    "HuggingfaceConversationCodec",
+    "NumpyListCodec",
+    "RawCodec",
+    "EqualUtil",
 ]

runtimes/huggingface/mlserver_huggingface/codecs/__init__.py

@@ -1,4 +1,4 @@
-from typing import Optional, Type, Any, Dict, AnyStr, List, Union
+from typing import Optional, Type, Any, Dict, List, Union
 from mlserver.codecs.utils import (
     has_decoded,
     _save_decoded,
@@ -40,38 +40,39 @@ class MultiInputRequestCodec(RequestCodec):
     Huggingface codecs is prefered, then mlserver's
     """
 
-    ContentType: str = StringCodec.ContentType
-    DefaultCodec: Type[InputCodecTy] = StringCodec
+    DefaultCodec: Type["InputCodecTy"] = StringCodec
     InputCodecsWithPriority: List[Type[InputCodecTy]] = []
+    ContentType = StringCodec.ContentType
 
     @classmethod
     def _find_encode_codecs(
-        cls, payload: Dict[AnyStr, Any]
-    ) -> Dict[str, Union[Type[InputCodecTy], None]]:
-        field_codec = {}
+        cls, payload: Dict[str, Any]
+    ) -> Dict[str, Union[Type["InputCodecTy"], "InputCodecTy", None]]:
+        field_codec: Dict[str, Union[Type["InputCodecTy"], "InputCodecTy", None]] = {}
         for field, value in payload.items():
             for codec in cls.InputCodecsWithPriority:
                 if codec.can_encode(value):
                     field_codec[field] = codec
                     break
             if field not in field_codec:
                 field_codec[field] = find_input_codec_by_payload(value)
+
         return field_codec
 
     @classmethod
     def _find_decode_codecs(
         cls, data: Union[InferenceResponse, InferenceRequest]
-    ) -> Dict[str, Union[Type[InputCodecTy], None]]:
+    ) -> Dict[str, Union[Type[InputCodecTy], InputCodecTy, None]]:
         field_codec = {}
-        fields = []
+        fields = []  # type: ignore
         if data.parameters:
             default_codec = find_input_codec(data.parameters.content_type)
         else:
             default_codec = cls.DefaultCodec
         if isinstance(data, InferenceRequest):
             fields = data.inputs
         else:
-            fields = data.outputs
+            fields = data.outputs  # type: ignore
         for field in fields:
             if not field.parameters:
                 field_codec[field.name] = default_codec
@@ -87,12 +88,12 @@ def _find_decode_codecs(
         return field_codec
 
     @classmethod
-    def _can_encode_request(cls, payload: Dict[AnyStr, Any]) -> bool:
+    def _can_encode_request(cls, payload: Dict[str, Any]) -> bool:
         field_codecs = cls._find_encode_codecs(payload)
         return bool(all(field_codecs.values()))
 
     @classmethod
-    def can_encode(cls, payload: Dict[AnyStr, Any]) -> bool:
+    def can_encode(cls, payload: Dict[str, Any]) -> bool:
         """
         Inputs always is Dict, Outputs always is list
         """
@@ -128,17 +129,21 @@ def encode_response(
         )
 
     @classmethod
-    def decode_response(cls, response: InferenceResponse) -> List[Any]:
+    def decode_response(
+        cls, response: InferenceResponse
+    ) -> Union[List[Any], Dict[Any, Any]]:
         """
         Always use HuggingfaceJSONCodec
         """
         data = {}
         is_list = True
         field_codecs = cls._find_decode_codecs(response)
         for item in response.outputs:
-            if not has_decoded(item) and field_codecs.get(item.name):
-                decoded_payload = field_codecs[item.name].decode_input(item)
-                _save_decoded(item, decoded_payload)
+            if not has_decoded(item):
+                codec = field_codecs[item.name]
+                if codec is not None:
+                    decoded_payload = codec.decode_output(item)
+                    _save_decoded(item, decoded_payload)
 
             value = get_decoded_or_raw(item)
             data[item.name] = value
@@ -154,6 +159,10 @@ def encode_request(cls, payload: Dict[str, Any], **kwargs) -> InferenceRequest:
         inputs = []
         for key, value in payload.items():
             codec = field_codecs[key]
+            if codec is None:
+                raise Exception(
+                    f"codec for key {key} value not found, value is {value}"
+                )
             input_v = codec.encode_input(key, value, **kwargs)
             set_content_type(input_v, codec.ContentType)
             inputs.append(input_v)
@@ -169,9 +178,11 @@ def decode_request(cls, request: InferenceRequest) -> Dict[str, Any]:
         values = {}
         field_codecs = cls._find_decode_codecs(request)
         for item in request.inputs:
-            if not has_decoded(item) and field_codecs.get(item.name):
-                decoded_payload = field_codecs[item.name].decode_input(item)
-                _save_decoded(item, decoded_payload)
+            if not has_decoded(item):
+                codec = field_codecs[item.name]
+                if codec is not None:
+                    decoded_payload = codec.decode_input(item)
+                    _save_decoded(item, decoded_payload)
 
             value = get_decoded_or_raw(item)
             values[item.name] = value

runtimes/huggingface/mlserver_huggingface/codecs/base.py

@@ -1,4 +1,4 @@
-from typing import List, Any, Dict
+from typing import List, Any
 from mlserver.codecs.base import InputCodec, register_input_codec
 from mlserver.types import RequestInput, ResponseOutput, Parameters
 from transformers.pipelines import Conversation
@@ -22,7 +22,7 @@ def can_encode(cls, payload: Any) -> bool:
     def encode_output(
         cls, name: str, payload: List[Conversation], use_bytes: bool = True, **kwargs
     ) -> ResponseOutput:
-        encoded = [json_encode(item) for item in payload]
+        encoded = [json_encode(item, use_bytes=use_bytes) for item in payload]
         shape = [len(encoded), 1]
         return ResponseOutput(
             name=name,
@@ -35,7 +35,7 @@ def encode_output(
         )
 
     @classmethod
-    def decode_output(cls, response_output: ResponseOutput) -> Dict[Any, Any]:
+    def decode_output(cls, response_output: ResponseOutput) -> List[Any]:
         packed = response_output.data.__root__
         return [json_decode(item) for item in packed]
 

runtimes/huggingface/mlserver_huggingface/codecs/conversation.py

@@ -1,22 +1,24 @@
 import io
 import base64
-from typing import List, Any
+from typing import List, Any, Union
 from PIL import Image
 from mlserver.codecs.base import InputCodec, register_input_codec
 from mlserver.codecs.lists import as_list, is_list_of
 from mlserver.types import RequestInput, ResponseOutput, Parameters
 from functools import partial
 
 
-def _pil_base64encode(img: "Image.Image", use_bytes: bool = False) -> bytes:
+def _pil_base64encode(img: "Image.Image", use_bytes: bool = False) -> Union[bytes, str]:
     buf = io.BytesIO()
-    img.save(buf, format="png")
+    img.save(buf, format=img.format)
     if use_bytes:
         return base64.b64encode(buf.getvalue())
     return base64.b64encode(buf.getvalue()).decode()
 
 
-def _pil_base64decode(imgbytes: bytes) -> "Image.Image":
+def _pil_base64decode(imgbytes: Union[bytes, str]) -> "Image.Image":
+    if isinstance(imgbytes, bytes):
+        imgbytes = imgbytes.decode()
     buf = io.BytesIO(base64.b64decode(imgbytes))
     return Image.open(buf)
 

runtimes/huggingface/mlserver_huggingface/codecs/image.py

@@ -21,14 +21,13 @@ def encode_output(
         cls, name: str, payload: Dict[Any, Any], use_bytes: bool = True, **kwargs
     ) -> ResponseOutput:
         encoded = json_encode(payload, use_bytes)
-        shape = [len(encoded), 1]
         return ResponseOutput(
             name=name,
             parameters=Parameters(
                 content_type=cls.ContentType,
             ),
             datatype="BYTES",
-            shape=shape,
+            shape=[1],
             data=[encoded],
         )
 

runtimes/huggingface/mlserver_huggingface/codecs/json.py

@@ -2,6 +2,7 @@
 from mlserver.codecs.base import InputCodec, register_input_codec
 from mlserver.types import RequestInput, ResponseOutput, Parameters
 from mlserver.codecs.lists import is_list_of
+from functools import partial
 from .utils import json_decode, json_encode
 
 
@@ -21,7 +22,7 @@ def can_encode(cls, payload: Any) -> bool:
     def encode_output(
         cls, name: str, payload: List[Dict[Any, Any]], use_bytes: bool = True, **kwargs
     ) -> ResponseOutput:
-        packed = map(json_encode, payload)
+        packed = map(partial(json_encode, use_bytes=use_bytes), payload)
         shape = [len(payload), 1]
         return ResponseOutput(
             name=name,
@@ -40,7 +41,7 @@ def decode_output(cls, response_output: ResponseOutput) -> List[Dict[Any, Any]]:
 
     @classmethod
     def encode_input(
-        cls, name: str, payload: Dict[Any, Any], use_bytes: bool = True, **kwargs
+        cls, name: str, payload: List[Dict[Any, Any]], use_bytes: bool = True, **kwargs
     ) -> RequestInput:
         output = cls.encode_output(name, payload, use_bytes)
         return RequestInput(

runtimes/huggingface/mlserver_huggingface/codecs/jsonlist.py

@@ -20,7 +20,10 @@ class NumpyListCodec(InputCodec):
 
     @classmethod
     def can_encode(csl, payload: Any) -> bool:
-        return is_list_of(payload, np.ndarray)
+        if not is_list_of(payload, np.ndarray):
+            return False
+        # only the support same shaped ndarray
+        return len(set([matrix.shape for matrix in payload])) == 1
 
     @classmethod
     def encode_output(
@@ -36,6 +39,7 @@ def encode_output(
             datatype=datatype,
             shape=shape,
             data=_encode_data(composed, datatype),
+            parameters=Parameters(content_type=cls.ContentType),
         )
 
     @classmethod

runtimes/huggingface/mlserver_huggingface/codecs/numpylist.py

@@ -0,0 +1,54 @@
+from typing import Any, Union
+from mlserver.types import RequestInput, ResponseOutput, Parameters
+from mlserver.codecs.base import InputCodec, register_input_codec
+
+
+@register_input_codec
+class RawCodec(InputCodec):
+    """
+    Encode/Decode raw python datatypes
+    """
+
+    ContentType = "raw"
+
+    @classmethod
+    def can_encode(cls, payload: Any) -> bool:
+        return (
+            isinstance(payload, int)
+            or isinstance(payload, str)
+            or isinstance(payload, float)
+        )
+
+    @classmethod
+    def encode_output(
+        cls, name: str, payload: Union[int, str, float], **kwargs
+    ) -> ResponseOutput:
+        return ResponseOutput(
+            name=name,
+            datatype="BYTES",
+            shape=[1],
+            data=[payload],
+            parameters=Parameters(content_type=cls.ContentType),
+        )
+
+    @classmethod
+    def decode_output(cls, response_output: ResponseOutput) -> Union[int, str, float]:
+        return cls.decode_input(response_output)  # type: ignore
+
+    @classmethod
+    def encode_input(
+        cls, name: str, payload: Union[int, str, float], **kwargs
+    ) -> RequestInput:
+        output = cls.encode_output(name=name, payload=payload)
+
+        return RequestInput(
+            name=output.name,
+            datatype=output.datatype,
+            shape=output.shape,
+            data=output.data,
+            parameters=output.parameters,
+        )
+
+    @classmethod
+    def decode_input(cls, request_input: RequestInput) -> Union[int, str, float]:
+        return request_input.data[0]