clarifai 11.1.7rc1__py3-none-any.whl → 11.1.7rc2__py3-none-any.whl

clarifai/runners/models/model_class.py

@@ -94,15 +94,17 @@ class ModelClass(ABC):
       method_info = method._cf_method_info
       signature = method_info.signature
       python_param_types = method_info.python_param_types
-      inputs = self._convert_input_protos_to_python(request.inputs, signature.inputs,
+      inputs = self._convert_input_protos_to_python(request.inputs, signature.input_fields,
                                                     python_param_types)
       if len(inputs) == 1:
         inputs = inputs[0]
         output = method(**inputs)
-        outputs.append(self._convert_output_to_proto(output, signature.outputs))
+        outputs.append(self._convert_output_to_proto(output, signature.output_fields))
       else:
         outputs = self._batch_predict(method, inputs)
-        outputs = [self._convert_output_to_proto(output, signature.outputs) for output in outputs]
+        outputs = [
+            self._convert_output_to_proto(output, signature.output_fields) for output in outputs
+        ]
 
       return service_pb2.MultiOutputResponse(
           outputs=outputs, status=status_pb2.Status(code=status_code_pb2.SUCCESS))
@@ -126,20 +128,21 @@ class ModelClass(ABC):
       signature = method_info.signature
       python_param_types = method_info.python_param_types
 
-      inputs = self._convert_input_protos_to_python(request.inputs, signature.inputs,
+      inputs = self._convert_input_protos_to_python(request.inputs, signature.input_fields,
                                                     python_param_types)
       if len(inputs) == 1:
         inputs = inputs[0]
         for output in method(**inputs):
           resp = service_pb2.MultiOutputResponse()
-          self._convert_output_to_proto(output, signature.outputs, proto=resp.outputs.add())
+          self._convert_output_to_proto(output, signature.output_fields, proto=resp.outputs.add())
           resp.status.code = status_code_pb2.SUCCESS
           yield resp
       else:
         for outputs in self._batch_generate(method, inputs):
           resp = service_pb2.MultiOutputResponse()
           for output in outputs:
-            self._convert_output_to_proto(output, signature.outputs, proto=resp.outputs.add())
+            self._convert_output_to_proto(
+                output, signature.output_fields, proto=resp.outputs.add())
           resp.status.code = status_code_pb2.SUCCESS
           yield resp
     except Exception as e:
@@ -166,13 +169,13 @@ class ModelClass(ABC):
       python_param_types = method_info.python_param_types
 
       # find the streaming vars in the signature
-      stream_sig = get_stream_from_signature(signature.inputs)
+      stream_sig = get_stream_from_signature(signature.input_fields)
       if stream_sig is None:
         raise ValueError("Streaming method must have a Stream input")
       stream_argname = stream_sig.name
 
       # convert all inputs for the first request, including the first stream value
-      inputs = self._convert_input_protos_to_python(request.inputs, signature.inputs,
+      inputs = self._convert_input_protos_to_python(request.inputs, signature.input_fields,
                                                     python_param_types)
       kwargs = inputs[0]
 
@@ -184,7 +187,7 @@ class ModelClass(ABC):
         yield first_item
         # subsequent streaming items contain only the streaming input
         for request in request_iterator:
-          item = self._convert_input_protos_to_python(request.inputs, stream_sig,
+          item = self._convert_input_protos_to_python(request.inputs, [stream_sig],
                                                       python_param_types)
           item = item[0][stream_argname]
           yield item
@@ -194,7 +197,7 @@ class ModelClass(ABC):
 
       for output in method(**kwargs):
         resp = service_pb2.MultiOutputResponse()
-        self._convert_output_to_proto(output, signature.outputs, proto=resp.outputs.add())
+        self._convert_output_to_proto(output, signature.output_fields, proto=resp.outputs.add())
         resp.status.code = status_code_pb2.SUCCESS
         yield resp
     except Exception as e:
@@ -206,7 +209,8 @@ class ModelClass(ABC):
           details=str(e),
           stack_trace=traceback.format_exc().split('\n')))
 
-  def _convert_input_protos_to_python(self, inputs: List[resources_pb2.Input], variables_signature,
+  def _convert_input_protos_to_python(self, inputs: List[resources_pb2.Input],
+                                      variables_signature: List[resources_pb2.ModelTypeField],
                                       python_param_types) -> List[Dict[str, Any]]:
     result = []
     for input in inputs:
@@ -219,7 +223,9 @@ class ModelClass(ABC):
       result.append(kwargs)
     return result
 
-  def _convert_output_to_proto(self, output: Any, variables_signature,
+  def _convert_output_to_proto(self,
+                               output: Any,
+                               variables_signature: List[resources_pb2.ModelTypeField],
                                proto=None) -> resources_pb2.Output:
     if proto is None:
       proto = resources_pb2.Output()

clarifai/runners/models/model_class_refract.py

@@ -0,0 +1,80 @@
+import inspect
+from abc import ABC, abstractmethod
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Dict, Iterator, List, get_type_hints
+
+from clarifai_grpc.grpc.api import resources_pb2, service_pb2
+
+from clarifai.runners.utils.data_handler import Output, kwargs_to_proto, proto_to_kwargs
+
+
+class ModelClass(ABC):
+
+  @abstractmethod
+  def load_model(self):
+    raise NotImplementedError("load_model() not implemented")
+
+  @abstractmethod
+  def predict(self, *args, **kwargs) -> Output:
+    raise NotImplementedError("predict() not implemented")
+
+  @abstractmethod
+  def generate(self, *args, **kwargs) -> Iterator[Output]:
+    raise NotImplementedError("generate() not implemented")
+
+  @abstractmethod
+  def stream(self, *args, **kwargs) -> Iterator[Output]:
+    raise NotImplementedError("stream() not implemented")
+
+  def batch_predict(self, inputs: List[Dict[str, Any]]) -> List[Output]:
+    with ThreadPoolExecutor() as executor:
+      return list(executor.map(lambda x: self.predict(**x), inputs))
+
+  def _process_request(self, request, process_func, is_stream=False):
+    inputs = self._convert_proto_to_python(request.inputs)
+    if len(inputs) == 1:
+      result = process_func(**inputs[0])
+      if is_stream:
+        return (self._convert_output_to_proto(output) for output in result)
+      else:
+        return [self._convert_output_to_proto(result)]
+    else:
+      results = self.batch_predict(inputs) if not is_stream else []
+      return [self._convert_output_to_proto(output) for output in results]
+
+  def predict_wrapper(
+      self, request: service_pb2.PostModelOutputsRequest) -> service_pb2.MultiOutputResponse:
+    outputs = self._process_request(request, self.predict)
+    return service_pb2.MultiOutputResponse(outputs=outputs)
+
+  def generate_wrapper(self, request: service_pb2.PostModelOutputsRequest
+                      ) -> Iterator[service_pb2.MultiOutputResponse]:
+    outputs = self._process_request(request, self.generate, is_stream=True)
+    for output in outputs:
+      yield service_pb2.MultiOutputResponse(outputs=[output])
+
+  def stream_wrapper(self, requests: Iterator[service_pb2.PostModelOutputsRequest]
+                    ) -> Iterator[service_pb2.MultiOutputResponse]:
+    for request in requests:
+      outputs = self._process_request(request, self.stream, is_stream=True)
+      yield service_pb2.MultiOutputResponse(outputs=outputs)
+
+  def _convert_proto_to_python(self, inputs: List[resources_pb2.Input]) -> List[Dict[str, Any]]:
+    get_type_hints(self.predict)
+    required_params = [
+        name for name, param in inspect.signature(self.predict).parameters.items()
+        if param.default == inspect.Parameter.empty
+    ]
+    kwargs_list = []
+    for input_proto in inputs:
+      kwargs = proto_to_kwargs(input_proto.data)
+      missing = [name for name in required_params if name not in kwargs]
+      if missing:
+        raise ValueError(f"Missing required parameters: {missing}")
+      kwargs_list.append(kwargs)
+    return kwargs_list
+
+  def _convert_output_to_proto(self, output: Any) -> resources_pb2.Output:
+    if isinstance(output, Output):
+      return output.to_proto()
+    return kwargs_to_proto(**output).outputs.add()