clarifai 11.1.5rc7__py3-none-any.whl → 11.1.5rc8__py3-none-any.whl

clarifai/client/model_client.py

@@ -1,4 +1,5 @@
 import inspect
+import json
 import time
 from typing import Any, Dict, Iterator, List
 
@@ -7,9 +8,10 @@ from clarifai_grpc.grpc.api.status import status_code_pb2
 
 from clarifai.constants.model import MAX_MODEL_PREDICT_INPUTS
 from clarifai.errors import UserError
-from clarifai.runners.utils.method_signatures import (deserialize, get_stream_from_signature,
-                                                      serialize, signatures_from_json,
-                                                      unflatten_nested_keys)
+from clarifai.runners.utils.method_signatures import (CompatibilitySerializer, deserialize,
+                                                      get_stream_from_signature, serialize,
+                                                      signatures_from_json)
+from clarifai.utils.logging import logger
 from clarifai.utils.misc import BackoffIterator, status_is_retryable
 
 
@@ -29,8 +31,25 @@ class ModelClient:
         '''
     self.STUB = stub
     self.request_template = request_template or service_pb2.PostModelOutputsRequest()
-    self._fetch_signatures()
-    self._define_functions()
+    self._method_signatures = None
+    self._defined = False
+
+  def fetch(self):
+    '''
+    Fetch function signature definitions from the model and define the functions in the client
+    '''
+    if self._defined:
+      return
+    try:
+      self._fetch_signatures()
+      self._define_functions()
+    finally:
+      self._defined = True
+
+  def __getattr__(self, name):
+    if not self._defined:
+      self.fetch()
+    return self.__getattribute__(name)
 
   def _fetch_signatures(self):
     '''
@@ -58,19 +77,18 @@ class ModelClient:
       response = self.STUB.PostModelOutputs(request)
       if status_is_retryable(
           response.status.code) and time.time() - start_time < 60 * 10:  # 10 minutes
-        self.logger.info(f"Retrying model info fetch with response {response.status!r}")
+        logger.info(f"Retrying model info fetch with response {response.status!r}")
         time.sleep(next(backoff_iterator))
         continue
       break
     if response.status.code == status_code_pb2.INPUT_UNSUPPORTED_FORMAT:
       # return code from older models that don't support _GET_SIGNATURES
       self._method_signatures = {}
+      self._define_compatability_functions()
       return
     if response.status.code != status_code_pb2.SUCCESS:
       raise Exception(f"Model failed with response {response!r}")
     self._method_signatures = signatures_from_json(response.outputs[0].data.text.raw)
-    import pdb
-    pdb.set_trace()
 
   def _define_functions(self):
     '''
@@ -116,39 +134,13 @@ class ModelClient:
       # set names, annotations and docstrings
       f.__name__ = method_name
       f.__qualname__ = f'{self.__class__.__name__}.{method_name}'
-      input_annos = {var.name: var.data_type for var in method_signature.inputs}
-      output_annos = {var.name: var.data_type for var in method_signature.outputs}
-      # unflatten nested keys to match the user function args for docs
-      input_annos = unflatten_nested_keys(input_annos, method_signature.inputs, is_output=False)
-      output_annos = unflatten_nested_keys(output_annos, method_signature.outputs, is_output=True)
-
-      # add Stream[] to the stream input annotations for docs
-      input_stream_argname, _ = get_stream_from_signature(method_signature.inputs)
-      if input_stream_argname:
-        input_annos[input_stream_argname] = 'Stream[' + str(
-            input_annos[input_stream_argname]) + ']'
-
-      # handle multiple outputs in the return annotation
-      return_annotation = output_annos
-      name = next(iter(output_annos.keys()))
-      if len(output_annos) == 1 and name == 'return':
-        # single output
-        return_annotation = output_annos[name]
-      elif name.startswith('return.') and name.split('.', 1)[1].isnumeric():
-        # tuple output
-        return_annotation = '(' + ", ".join(output_annos[f'return.{i}']
-                                            for i in range(len(output_annos))) + ')'
-      else:
-        # named output
-        return_annotation = f'Output({", ".join(f"{k}={t}" for k, t in output_annos.items())})'
-      if method_signature.method_type in ['generate', 'stream']:
-        return_annotation = f'Stream[{return_annotation}]'
-
-      # set annotations and docstrings
+      f.__doc__ = method_signature.docstring
+      input_annotations = json.loads(method_signature.annotations_json)
+      return_annotation = input_annotations.pop('return', None)
       sig = inspect.signature(f).replace(
           parameters=[
               inspect.Parameter(k, inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=v)
-              for k, v in input_annos.items()
+              for k, v in input_annotations.items()
           ],
           return_annotation=return_annotation,
       )
@@ -156,6 +148,28 @@ class ModelClient:
       f.__doc__ = method_signature.docstring
       setattr(self, method_name, f)
 
+  def _define_compatability_functions(self):
+
+    serializer = CompatibilitySerializer()
+
+    def predict(input: Any) -> Any:
+      proto = resources_pb2.Input()
+      serializer.serialize(proto.data, input)
+      # always use text.raw for compat
+      if proto.data.string_value:
+        proto.data.text.raw = proto.data.string_value
+        proto.data.string_value = ''
+      response = self._predict_by_proto([proto])
+      if response.status.code != status_code_pb2.SUCCESS:
+        raise Exception(f"Model predict failed with response {response!r}")
+      response_data = response.outputs[0].data
+      if response_data.text.raw:
+        response_data.string_value = response_data.text.raw
+        response_data.text.raw = ''
+      return serializer.deserialize(response_data)
+
+    self.predict = predict
+
   def _predict(
       self,
       inputs,  # TODO set up functions according to fetched signatures?
@@ -214,10 +228,9 @@ class ModelClient:
     request.inputs.extend(inputs)
 
     if method_name:
-      for inp in inputs:
-        inp.data.metadata['_method_name'] = method_name
       # TODO put in new proto field?
-      # request.model.model_version.output_info.params['_method_name'] = method_name
+      for inp in request.inputs:
+        inp.data.metadata['_method_name'] = method_name
     if inference_params:
       request.model.model_version.output_info.params.update(inference_params)
     if output_config:
@@ -230,7 +243,7 @@ class ModelClient:
       response = self.STUB.PostModelOutputs(request)
       if status_is_retryable(
           response.status.code) and time.time() - start_time < 60 * 10:  # 10 minutes
-        self.logger.info(f"Model predict failed with response {response!r}")
+        logger.info(f"Model predict failed with response {response!r}")
         time.sleep(next(backoff_iterator))
         continue
 
@@ -298,9 +311,8 @@ class ModelClient:
 
     if method_name:
       # TODO put in new proto field?
-      for inp in inputs:
+      for inp in request.inputs:
         inp.data.metadata['_method_name'] = method_name
-      # request.model.model_version.output_info.params['_method_name'] = method_name
     if inference_params:
       request.model.model_version.output_info.params.update(inference_params)
     if output_config:
@@ -317,7 +329,7 @@ class ModelClient:
         raise Exception("Model Generate failed with no response")
       if status_is_retryable(response.status.code) and \
               time.time() - start_time < 60 * 10:
-        self.logger.info("Model is still deploying, please wait...")
+        logger.info("Model is still deploying, please wait...")
         time.sleep(next(backoff_iterator))
         continue
       if response.status.code != status_code_pb2.SUCCESS:
@@ -346,7 +358,10 @@ class ModelClient:
     kwargs = inputs
 
     # find the streaming vars in the input signature, and the streaming input python param
-    stream_argname, streaming_var_signatures = get_stream_from_signature(input_signature)
+    stream_sig = get_stream_from_signature(input_signature)
+    if stream_sig is None:
+      raise ValueError("Streaming method must have a Stream input")
+    stream_argname = stream_sig.name
 
     # get the streaming input generator from the user-provided function arg values
     user_inputs_generator = kwargs.pop(stream_argname)
@@ -366,7 +381,7 @@ class ModelClient:
       # subsequent items are just the stream items
       for item in user_inputs_generator:
         proto = resources_pb2.Input()
-        serialize({stream_argname: item}, streaming_var_signatures, proto.data)
+        serialize({stream_argname: item}, [stream_sig], proto.data)
         yield proto
 
     response_stream = self._stream_by_proto(_input_proto_stream(), method_name)
@@ -383,7 +398,6 @@ class ModelClient:
                     output_config: Dict = {}):
     request = service_pb2.PostModelOutputsRequest()
     request.CopyFrom(self.request_template)
-    # request.model.model_version.output_info.params['_method_name'] = method_name
     if inference_params:
       request.model.model_version.output_info.params.update(inference_params)
     if output_config:
@@ -395,6 +409,7 @@ class ModelClient:
         req.inputs.extend(inputs)
       else:
         req.inputs.append(inputs)
+      # TODO: put into new proto field?
       for inp in req.inputs:
         inp.data.metadata['_method_name'] = method_name
       yield req
@@ -421,7 +436,7 @@ class ModelClient:
       for response in stream_response:
         if status_is_retryable(response.status.code) and \
                 time.time() - start_time < 60 * 10:
-          self.logger.info("Model is still deploying, please wait...")
+          logger.info("Model is still deploying, please wait...")
           time.sleep(next(backoff_iterator))
           break
         if response.status.code != status_code_pb2.SUCCESS:

clarifai/runners/models/model_class.py

@@ -19,32 +19,11 @@ _METHOD_INFO_ATTR = '_cf_method_info'
 _RAISE_EXCEPTIONS = os.getenv("RAISE_EXCEPTIONS", "false").lower() == "true"
 
 
-class methods:
-  '''
-  Decorators to mark methods as predict, generate, or stream methods.
-  '''
-
-  @staticmethod
-  def predict(method):
-    setattr(method, _METHOD_INFO_ATTR, _MethodInfo(method, 'predict'))
-    return method
-
-  @staticmethod
-  def generate(method):
-    setattr(method, _METHOD_INFO_ATTR, _MethodInfo(method, 'generate'))
-    return method
-
-  @staticmethod
-  def stream(method):
-    setattr(method, _METHOD_INFO_ATTR, _MethodInfo(method, 'stream'))
-    return method
-
-
 class ModelClass(ABC):
   '''
   Base class for model classes that can be run as a service.
 
-  Define methods as predict, generate, or stream methods using the @methods decorators.
+  Define predict, generate, or stream methods using the @ModelClass.method decorator.
 
   Example:
 
@@ -53,21 +32,26 @@ class ModelClass(ABC):
 
     class MyModel(ModelClass):
 
-      @methods.predict
+      @ModelClass.method
       def predict(self, x: str, y: int) -> List[str]:
         return [x] * y
 
-      @methods.generate
+      @ModelClass.method
       def generate(self, x: str, y: int) -> Stream[str]:
         for i in range(y):
           yield x + str(i)
 
-      @methods.stream
+      @ModelClass.method
       def stream(self, input_stream: Stream[Input(x=str, y=int)]) -> Stream[str]:
         for item in input_stream:
           yield item.x + ' ' + str(item.y)
   '''
 
+  @staticmethod
+  def method(func):
+    setattr(func, _METHOD_INFO_ATTR, _MethodInfo(func))
+    return func
+
   def load_model(self):
     """Load the model."""
 
@@ -80,7 +64,7 @@ class ModelClass(ABC):
     output.data.text.raw = signatures_to_json(signatures)
     return resp
 
-  def batch_predict(self, method, inputs: List[Dict[str, Any]]) -> List[Any]:
+  def _batch_predict(self, method, inputs: List[Dict[str, Any]]) -> List[Any]:
     """Batch predict method for multiple inputs."""
     outputs = []
     for input in inputs:
@@ -88,7 +72,7 @@ class ModelClass(ABC):
       outputs.append(output)
     return outputs
 
-  def batch_generate(self, method, inputs: List[Dict[str, Any]]) -> Iterator[List[Any]]:
+  def _batch_generate(self, method, inputs: List[Dict[str, Any]]) -> Iterator[List[Any]]:
     """Batch generate method for multiple inputs."""
     generators = [method(**input) for input in inputs]
     for outputs in itertools.zip_longest(*generators):
@@ -99,11 +83,9 @@ class ModelClass(ABC):
     outputs = []
     try:
       # TODO add method name field to proto
-      method_name = None
+      method_name = 'predict'
       if len(request.inputs) > 0 and '_method_name' in request.inputs[0].data.metadata:
         method_name = request.inputs[0].data.metadata['_method_name']
-      # call_params = dict(request.model.model_version.output_info.params)
-      # method_name = call_params.get('_method_name', 'predict')
       if method_name == '_GET_SIGNATURES':  # special case to fetch signatures, TODO add endpoint for this
         return self._handle_get_signatures_request()
       if method_name not in self._get_method_info():
@@ -119,7 +101,7 @@ class ModelClass(ABC):
         output = method(**inputs)
         outputs.append(self._convert_output_to_proto(output, signature.outputs))
       else:
-        outputs = self.batch_predict(method, inputs)
+        outputs = self._batch_predict(method, inputs)
         outputs = [self._convert_output_to_proto(output, signature.outputs) for output in outputs]
 
       return service_pb2.MultiOutputResponse(
@@ -136,8 +118,9 @@ class ModelClass(ABC):
   def generate_wrapper(self, request: service_pb2.PostModelOutputsRequest
                       ) -> Iterator[service_pb2.MultiOutputResponse]:
     try:
-      call_params = dict(request.model.model_version.output_info.params)
-      method_name = call_params.get('_method_name', 'generate')
+      method_name = 'generate'
+      if len(request.inputs) > 0 and '_method_name' in request.inputs[0].data.metadata:
+        method_name = request.inputs[0].data.metadata['_method_name']
       method = getattr(self, method_name)
       method_info = method._cf_method_info
       signature = method_info.signature
@@ -153,7 +136,7 @@ class ModelClass(ABC):
           resp.status.code = status_code_pb2.SUCCESS
           yield resp
       else:
-        for outputs in self.batch_generate(method, inputs):
+        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())
@@ -174,15 +157,19 @@ class ModelClass(ABC):
       request = next(request_iterator)  # get first request to determine method
       assert len(request.inputs) == 1, "Streaming requires exactly one input"
 
-      call_params = dict(request.model.model_version.output_info.params)
-      method_name = call_params.get('_method_name', 'stream')
+      method_name = 'generate'
+      if len(request.inputs) > 0 and '_method_name' in request.inputs[0].data.metadata:
+        method_name = request.inputs[0].data.metadata['_method_name']
       method = getattr(self, method_name)
       method_info = method._cf_method_info
       signature = method_info.signature
       python_param_types = method_info.python_param_types
 
       # find the streaming vars in the signature
-      stream_argname, streaming_var_signatures = get_stream_from_signature(signature.inputs)
+      stream_sig = get_stream_from_signature(signature.inputs)
+      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,
@@ -197,7 +184,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, streaming_var_signatures,
+          item = self._convert_input_protos_to_python(request.inputs, stream_sig,
                                                       python_param_types)
           item = item[0][stream_argname]
           yield item
@@ -236,11 +223,7 @@ class ModelClass(ABC):
                                proto=None) -> resources_pb2.Output:
     if proto is None:
       proto = resources_pb2.Output()
-    if isinstance(output, tuple):
-      output = {f'return.{i}': item for i, item in enumerate(output)}
-    if not isinstance(output, dict):  # TODO Output type, not just dict
-      output = {'return': output}
-    serialize(output, variables_signature, proto.data, is_output=True)
+    serialize({'return': output}, [variables_signature], proto.data, is_output=True)
     proto.status.code = status_code_pb2.SUCCESS
     return proto
 
@@ -259,7 +242,7 @@ class ModelClass(ABC):
     #  if hasattr(cls, name):
     #    method = getattr(cls, name)
     #    if not hasattr(method, _METHOD_INFO_ATTR):  # not already put in registry
-    #      methods[name] = _MethodInfo(method, method_type=name)
+    #      methods[name] = _MethodInfo(method)
     # set method table for this class in the registry
     return methods
 
@@ -275,9 +258,9 @@ class ModelClass(ABC):
 
 class _MethodInfo:
 
-  def __init__(self, method, method_type):
+  def __init__(self, method):
     self.name = method.__name__
-    self.signature = build_function_signature(method, method_type)
+    self.signature = build_function_signature(method)
     self.python_param_types = {
         p.name: p.annotation
         for p in inspect.signature(method).parameters.values()

clarifai/runners/utils/data_types.py

@@ -1,4 +1,5 @@
 import io
+import json
 from typing import Iterable, List, get_args, get_origin
 
 import numpy as np
@@ -21,13 +22,19 @@ class MessageData:
   def from_proto(cls, proto):
     raise NotImplementedError
 
+  @classmethod
+  def from_value(cls, value):
+    if isinstance(value, cls):
+      return value
+    return cls(value)
+
   def cast(self, python_type):
     if python_type == self.__class__:
       return self
     raise TypeError(f'Incompatible type for {self.__class__.__name__}: {python_type}')
 
 
-class Output(dict):
+class NamedFields(dict):
   __getattr__ = dict.__getitem__
   __setattr__ = dict.__setitem__
 
@@ -38,19 +45,36 @@ class Output(dict):
     return list(self.keys())
 
 
-class Input(dict):
-  __getattr__ = dict.__getitem__
-  __setattr__ = dict.__setitem__
+class Stream(Iterable):
+  pass
 
-  def __origin__(self):
-    return self
 
-  def __args__(self):
-    return list(self.keys())
+class JSON:
 
+  def __init__(self, value):
+    self.value = value
 
-class Stream(Iterable):
-  pass
+  def __eq__(self, other):
+    return self.value == other
+
+  def __bool__(self):
+    return bool(self.value)
+
+  def to_json(self):
+    return json.dumps(self.value)
+
+  @classmethod
+  def from_json(cls, json_str):
+    return cls(json.loads(json_str))
+
+  @classmethod
+  def from_value(cls, value):
+    return cls(value)
+
+  def cast(self, python_type):
+    if not isinstance(self.value, python_type):
+      raise TypeError(f'Incompatible type {type(self.value)} for {python_type}')
+    return self.value
 
 
 class Text(MessageData):
@@ -59,6 +83,16 @@ class Text(MessageData):
     self.text = text
     self.url = url
 
+  def __eq__(self, other):
+    if isinstance(other, Text):
+      return self.text == other.text and self.url == other.url
+    if isinstance(other, str):
+      return self.text == other
+    return False
+
+  def __bool__(self):
+    return bool(self.text) or bool(self.url)
+
   def to_proto(self) -> TextProto:
     return TextProto(raw=self.text or '', url=self.url or '')
 
@@ -66,6 +100,16 @@ class Text(MessageData):
   def from_proto(cls, proto: TextProto) -> "Text":
     return cls(proto.raw, proto.url or None)
 
+  @classmethod
+  def from_value(cls, value):
+    if isinstance(value, str):
+      return cls(value)
+    if isinstance(value, Text):
+      return value
+    if isinstance(value, dict):
+      return cls(value.get('text'), value.get('url'))
+    raise TypeError(f'Incompatible type for Text: {type(value)}')
+
   def cast(self, python_type):
     if python_type == str:
       return self.text
@@ -189,6 +233,14 @@ class Image(MessageData):
   def from_proto(cls, proto: ImageProto) -> "Image":
     return cls(proto)
 
+  @classmethod
+  def from_value(cls, value):
+    if isinstance(value, PILImage.Image):
+      return cls.from_pil(value)
+    if isinstance(value, Image):
+      return value
+    raise TypeError(f'Incompatible type for Image: {type(value)}')
+
   def cast(self, python_type):
     if python_type == Image:
       return self