wallaroo.ModelConversion

 1from enum import Enum
 2from typing import Any, Dict, List, NamedTuple, Union
 3
 4
 5class ModelConversionInputType(Enum):
 6    Float16 = "float16"
 7    Float32 = "float32"
 8    Float64 = "float64"
 9    Int16 = "int16"
10    Int32 = "int32"
11    Int64 = "int64"
12    UInt8 = "uint8"
13    UInt16 = "uint16"
14    UInt32 = "uint32"
15    UInt64 = "uint64"
16    Boolean = "bool"
17    Double = "double"
18
19
20class ConvertKerasArguments(NamedTuple):
21    name: str
22    comment: Union[str, None]
23    input_type: ModelConversionInputType
24    dimensions: List[Union[None, int, float]]
25
26    def to_dict(self) -> Dict[str, Any]:
27        base = self._asdict()
28        base["input_type"] = self.input_type.value
29        return base
30
31
32class ConvertSKLearnArguments(NamedTuple):
33    name: str
34    number_of_columns: int
35    input_type: ModelConversionInputType
36    comment: Union[str, None]
37
38    def to_dict(self) -> Dict[str, Any]:
39        base = self._asdict()
40        base["input_type"] = self.input_type.value
41        return base
42
43
44class ConvertXGBoostArgs(NamedTuple):
45    name: str
46    number_of_columns: int
47    input_type: ModelConversionInputType
48    comment: Union[str, None]
49
50    def to_dict(self) -> Dict[str, Any]:
51        base = self._asdict()
52        base["input_type"] = self.input_type.value
53        return base
54
55
56ModelConversionArguments = Union[
57    ConvertKerasArguments, ConvertSKLearnArguments, ConvertXGBoostArgs
58]
59
60
61class ModelConversionSource(Enum):
62    KERAS = "keras"
63    XGBOOST = "xgboost"
64    SKLEARN = "sklearn"
65
66
67class ModelConversionGenericException(Exception):
68    pass
69
70
71class ModelConversionFailure(Exception):
72    pass
73
74
75class ModelConversionUnsupportedType(Exception):
76    pass
77
78
79class ModelConversionSourceFileNotPresent(Exception):
80    pass