clarifai 11.3.0rc2__py3-none-any.whl → 11.4.0__py3-none-any.whl

clarifai/client/model_client.py

@@ -9,20 +9,24 @@ from clarifai.constants.model import MAX_MODEL_PREDICT_INPUTS
 from clarifai.errors import UserError
 from clarifai.runners.utils import code_script, method_signatures
 from clarifai.runners.utils.data_utils import is_openai_chat_format
-from clarifai.runners.utils.method_signatures import (CompatibilitySerializer, deserialize,
-                                                      get_stream_from_signature, serialize,
-                                                      signatures_from_json)
+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
 
 
 class ModelClient:
-  '''
-  Client for calling model predict, generate, and stream methods.
-  '''
-
-  def __init__(self, stub, request_template: service_pb2.PostModelOutputsRequest = None):
     '''
+    Client for calling model predict, generate, and stream methods.
+    '''
+
+    def __init__(self, stub, request_template: service_pb2.PostModelOutputsRequest = None):
+        '''
         Initialize the model client.
 
         Args:
@@ -30,501 +34,520 @@ class ModelClient:
             request_template: The template for the request to send to the model, including
             common fields like model_id, model_version, cluster, etc.
         '''
-    self.STUB = stub
-    self.request_template = request_template or service_pb2.PostModelOutputsRequest()
-    self._method_signatures = None
-    self._defined = False
+        self.STUB = stub
+        self.request_template = request_template or service_pb2.PostModelOutputsRequest()
+        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):
-    '''
-      Fetch the method signatures from the model.
-
-      Returns:
-          Dict: The method signatures.
-      '''
-    try:
-      response = self.STUB.GetModelVersion(
-          service_pb2.GetModelVersionRequest(
-              user_app_id=self.request_template.user_app_id,
-              model_id=self.request_template.model_id,
-              version_id=self.request_template.version_id,
-          ))
-
-      method_signatures = None
-      if response.status.code == status_code_pb2.SUCCESS:
-        method_signatures = response.model_version.method_signatures
-      if response.status.code != status_code_pb2.SUCCESS:
-        raise Exception(f"Model failed with response {response!r}")
-      self._method_signatures = {}
-      for method_signature in method_signatures:
-        method_name = method_signature.name
-        # check for duplicate method names
-        if method_name in self._method_signatures:
-          raise ValueError(f"Duplicate method name {method_name}")
-        self._method_signatures[method_name] = method_signature
-      if not self._method_signatures:  # if no method signatures, try to fetch from the model
-        self._fetch_signatures_backup()
-    except Exception:
-      # try to fetch from the model
-      self._fetch_signatures_backup()
-      if not self._method_signatures:
-        raise ValueError("Failed to fetch method signatures from model and backup method")
-
-  def _fetch_signatures_backup(self):
-    '''
-      This is a temporary method of fetching the method signatures from the model.
-
-      Returns:
-          Dict: The method signatures.
-      '''
-
-    request = service_pb2.PostModelOutputsRequest()
-    request.CopyFrom(self.request_template)
-    # request.model.model_version.output_info.params['_method_name'] = '_GET_SIGNATURES'
-    inp = request.inputs.add()  # empty input for this method
-    inp.data.parts.add()  # empty part for this input
-    inp.data.metadata['_method_name'] = '_GET_SIGNATURES'
-    start_time = time.time()
-    backoff_iterator = BackoffIterator(10)
-    while True:
-      response = self.STUB.PostModelOutputs(request)
-      if status_is_retryable(
-          response.status.code) and time.time() - start_time < 60 * 10:  # 10 minutes
-        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 or
-        (response.status.code == status_code_pb2.SUCCESS and
-         response.outputs[0].data.text.raw == '')):
-      # return codes/values 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)
-
-  def _define_functions(self):
-    '''
-    Define the functions based on the method signatures.
-    '''
-    for method_name, method_signature in self._method_signatures.items():
-      # define the function in this client instance
-      if resources_pb2.RunnerMethodType.Name(method_signature.method_type) == 'UNARY_UNARY':
-        call_func = self._predict
-      elif resources_pb2.RunnerMethodType.Name(method_signature.method_type) == 'UNARY_STREAMING':
-        call_func = self._generate
-      elif resources_pb2.RunnerMethodType.Name(
-          method_signature.method_type) == 'STREAMING_STREAMING':
-        call_func = self._stream
-      else:
-        raise ValueError(f"Unknown method type {method_signature.method_type}")
-
-      # method argnames, in order, collapsing nested keys to corresponding user function args
-      method_argnames = []
-      for var in method_signature.input_fields:
-        outer = var.name.split('.', 1)[0]
-        if outer in method_argnames:
-          continue
-        method_argnames.append(outer)
-
-      def bind_f(method_name, method_argnames, call_func):
-
-        def f(*args, **kwargs):
-          if len(args) > len(method_argnames):
-            raise TypeError(
-                f"{method_name}() takes {len(method_argnames)} positional arguments but {len(args)} were given"
+    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):
+        '''
+        Fetch the method signatures from the model.
+
+        Returns:
+            Dict: The method signatures.
+        '''
+        try:
+            response = self.STUB.GetModelVersion(
+                service_pb2.GetModelVersionRequest(
+                    user_app_id=self.request_template.user_app_id,
+                    model_id=self.request_template.model_id,
+                    version_id=self.request_template.version_id,
+                )
             )
 
-          if len(args) + len(kwargs) > len(method_argnames):
-            raise TypeError(
-                f"{method_name}() got an unexpected keyword argument {next(iter(kwargs))}")
-          if len(args) == 1 and (not kwargs) and isinstance(args[0], list):
-            batch_inputs = args[0]
-            # Validate each input is a dictionary
-            is_batch_input_valid = all(isinstance(input, dict) for input in batch_inputs)
-            if is_batch_input_valid and (not is_openai_chat_format(batch_inputs)):
-              # If the batch input is valid, call the function with the batch inputs and the method name
-              return call_func(batch_inputs, method_name)
-
-          for name, arg in zip(method_argnames, args):  # handle positional with zip shortest
-            if name in kwargs:
-              raise TypeError(f"Multiple values for argument {name}")
-            kwargs[name] = arg
-          return call_func(kwargs, method_name)
-
-        return f
-
-      # need to bind method_name to the value, not the mutating loop variable
-      f = bind_f(method_name, method_argnames, call_func)
-
-      # set names, annotations and docstrings
-      f.__name__ = method_name
-      f.__qualname__ = f'{self.__class__.__name__}.{method_name}'
-      input_annotations = code_script._get_annotations_source(method_signature)
-      return_annotation = input_annotations.pop('return', (None, None))[0]
-      sig = inspect.signature(f).replace(
-          parameters=[
-              inspect.Parameter(k, inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=v[0])
-              for k, v in input_annotations.items()
-          ],
-          return_annotation=return_annotation,
-      )
-      f.__signature__ = sig
-      f.__doc__ = method_signature.description
-      setattr(self, method_name, f)
-
-  def available_methods(self) -> List[str]:
-    """Get the available methods for this model.
-
-    Returns:
-        List[str]: The available methods.
-    """
-    if not self._defined:
-      self.fetch()
-    return self._method_signatures.keys()
-
-  def method_signature(self, method_name: str) -> str:
-    """Get the method signature for a method.
-
-    Args:
-        method_name (str): The name of the method.
-
-    Returns:
-        str: The method signature.
-    """
-    if not self._defined:
-      self.fetch()
-    return method_signatures.get_method_signature(self._method_signatures[method_name])
-
-  def generate_client_script(self) -> str:
-    """Generate a client script for this model.
-
-    Returns:
-        str: The client script.
-    """
-    if not self._defined:
-      self.fetch()
-    method_signatures = []
-    for _, method_signature in self._method_signatures.items():
-      method_signatures.append(method_signature)
-    return code_script.generate_client_script(
-        method_signatures,
-        user_id=self.request_template.user_app_id.user_id,
-        app_id=self.request_template.user_app_id.app_id,
-        model_id=self.request_template.model_id)
-
-  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?
-      method_name: str = 'predict',
-  ) -> Any:
-    input_signature = self._method_signatures[method_name].input_fields
-    output_signature = self._method_signatures[method_name].output_fields
-
-    batch_input = True
-    if isinstance(inputs, dict):
-      inputs = [inputs]
-      batch_input = False
-
-    proto_inputs = []
-    for input in inputs:
-      proto = resources_pb2.Input()
-
-      serialize(input, input_signature, proto.data)
-      proto_inputs.append(proto)
-
-    response = self._predict_by_proto(proto_inputs, method_name)
-
-    outputs = []
-    for output in response.outputs:
-      outputs.append(deserialize(output.data, output_signature, is_output=True))
-    if batch_input:
-      return outputs
-    return outputs[0]
-
-  def _predict_by_proto(
-      self,
-      inputs: List[resources_pb2.Input],
-      method_name: str = None,
-      inference_params: Dict = None,
-      output_config: Dict = None,
-  ) -> service_pb2.MultiOutputResponse:
-    """Predicts the model based on the given inputs.
-
-      Args:
-          inputs (List[resources_pb2.Input]): The inputs to predict.
-          method_name (str): The remote method name to call.
-          inference_params (Dict): Inference parameters to override.
-          output_config (Dict): Output configuration to override.
-
-      Returns:
-          service_pb2.MultiOutputResponse: The prediction response(s).
-      """
-    if not isinstance(inputs, list):
-      raise UserError('Invalid inputs, inputs must be a list of Input objects.')
-    if len(inputs) > MAX_MODEL_PREDICT_INPUTS:
-      raise UserError(f"Too many inputs. Max is {MAX_MODEL_PREDICT_INPUTS}.")
-
-    request = service_pb2.PostModelOutputsRequest()
-    request.CopyFrom(self.request_template)
-
-    request.inputs.extend(inputs)
-
-    if method_name:
-      # TODO put in new proto field?
-      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:
-      request.model.model_version.output_info.output_config.MergeFrom(
-          resources_pb2.OutputConfig(**output_config))
-
-    start_time = time.time()
-    backoff_iterator = BackoffIterator(10)
-    while True:
-      response = self.STUB.PostModelOutputs(request)
-      if status_is_retryable(
-          response.status.code) and time.time() - start_time < 60 * 10:  # 10 minutes
-        logger.info("Model is still deploying, please wait...")
-        time.sleep(next(backoff_iterator))
-        continue
-
-      if response.status.code != status_code_pb2.SUCCESS:
-        raise Exception(f"Model predict failed with response {response!r}")
-      break
-    return response
-
-  def _generate(
-      self,
-      inputs,  # TODO set up functions according to fetched signatures?
-      method_name: str = 'generate',
-  ) -> Any:
-    input_signature = self._method_signatures[method_name].input_fields
-    output_signature = self._method_signatures[method_name].output_fields
-
-    batch_input = True
-    if isinstance(inputs, dict):
-      inputs = [inputs]
-      batch_input = False
-
-    proto_inputs = []
-    for input in inputs:
-      proto = resources_pb2.Input()
-      serialize(input, input_signature, proto.data)
-      proto_inputs.append(proto)
-
-    response_stream = self._generate_by_proto(proto_inputs, method_name)
-
-    for response in response_stream:
-      outputs = []
-      for output in response.outputs:
-        outputs.append(deserialize(output.data, output_signature, is_output=True))
-      if batch_input:
-        yield outputs
-      else:
-        yield outputs[0]
-
-  def _generate_by_proto(
-      self,
-      inputs: List[resources_pb2.Input],
-      method_name: str = None,
-      inference_params: Dict = {},
-      output_config: Dict = {},
-  ):
-    """Generate the stream output on model based on the given inputs.
-
-    Args:
-        inputs (list[Input]): The inputs to generate, must be less than 128.
-        method_name (str): The remote method name to call.
-        inference_params (dict): The inference params to override.
-        output_config (dict): The output config to override.
-    """
-    if not isinstance(inputs, list):
-      raise UserError('Invalid inputs, inputs must be a list of Input objects.')
-    if len(inputs) > MAX_MODEL_PREDICT_INPUTS:
-      raise UserError(f"Too many inputs. Max is {MAX_MODEL_PREDICT_INPUTS}."
-                     )  # TODO Use Chunker for inputs len > 128
-
-    request = service_pb2.PostModelOutputsRequest()
-    request.CopyFrom(self.request_template)
-
-    request.inputs.extend(inputs)
-
-    if method_name:
-      # TODO put in new proto field?
-      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:
-      request.model.model_version.output_info.output_config.MergeFromDict(output_config)
-
-    start_time = time.time()
-    backoff_iterator = BackoffIterator(10)
-    started = False
-    while not started:
-      stream_response = self.STUB.GenerateModelOutputs(request)
-      try:
-        response = next(stream_response)  # get the first response
-      except StopIteration:
-        raise Exception("Model Generate failed with no response")
-      if status_is_retryable(response.status.code) and \
-              time.time() - start_time < 60 * 10:
-        logger.info("Model is still deploying, please wait...")
-        time.sleep(next(backoff_iterator))
-        continue
-      if response.status.code != status_code_pb2.SUCCESS:
-        raise Exception(f"Model Generate failed with response {response.status!r}")
-      started = True
-
-    yield response  # yield the first response
-
-    for response in stream_response:
-      if response.status.code != status_code_pb2.SUCCESS:
-        raise Exception(f"Model Generate failed with response {response.status!r}")
-      yield response
-
-  def _stream(
-      self,
-      inputs,
-      method_name: str = 'stream',
-  ) -> Any:
-    input_signature = self._method_signatures[method_name].input_fields
-    output_signature = self._method_signatures[method_name].output_fields
-
-    if isinstance(inputs, list):
-      assert len(inputs) == 1, 'streaming methods do not support batched calls'
-      inputs = inputs[0]
-    assert isinstance(inputs, dict)
-    kwargs = inputs
-
-    # find the streaming vars in the input signature, and the streaming input python param
-    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)
-
-    def _input_proto_stream():
-      # first item contains all the inputs and the first stream item
-      proto = resources_pb2.Input()
-      try:
-        item = next(user_inputs_generator)
-      except StopIteration:
-        return  # no items to stream
-      kwargs[stream_argname] = item
-      serialize(kwargs, input_signature, proto.data)
-
-      yield proto
-
-      # subsequent items are just the stream items
-      for item in user_inputs_generator:
-        proto = resources_pb2.Input()
-        serialize({stream_argname: item}, [stream_sig], proto.data)
-        yield proto
-
-    response_stream = self._stream_by_proto(_input_proto_stream(), method_name)
-
-    for response in response_stream:
-      assert len(response.outputs) == 1, 'streaming methods must have exactly one output'
-      yield deserialize(response.outputs[0].data, output_signature, is_output=True)
-
-  def _req_iterator(self,
-                    input_iterator: Iterator[List[resources_pb2.Input]],
-                    method_name: str = None,
-                    inference_params: Dict = {},
-                    output_config: Dict = {}):
-    request = service_pb2.PostModelOutputsRequest()
-    request.CopyFrom(self.request_template)
-    if inference_params:
-      request.model.model_version.output_info.params.update(inference_params)
-    if output_config:
-      request.model.model_version.output_info.output_config.MergeFromDict(output_config)
-    for inputs in input_iterator:
-      req = service_pb2.PostModelOutputsRequest()
-      req.CopyFrom(request)
-      if isinstance(inputs, list):
-        req.inputs.extend(inputs)
-      else:
-        req.inputs.append(inputs)
-      # TODO: put into new proto field?
-      if method_name:
-        for inp in req.inputs:
-          inp.data.metadata['_method_name'] = method_name
-      yield req
-
-  def _stream_by_proto(self,
-                       inputs: Iterator[List[resources_pb2.Input]],
-                       method_name: str = None,
-                       inference_params: Dict = {},
-                       output_config: Dict = {}):
-    """Generate the stream output on model based on the given stream of inputs.
-    """
-    # if not isinstance(inputs, Iterator[List[Input]]):
-    #   raise UserError('Invalid inputs, inputs must be a iterator of list of Input objects.')
-
-    request = self._req_iterator(inputs, method_name, inference_params, output_config)
-
-    start_time = time.time()
-    backoff_iterator = BackoffIterator(10)
-    generation_started = False
-    while True:
-      if generation_started:
-        break
-      stream_response = self.STUB.StreamModelOutputs(request)
-      for response in stream_response:
-        if status_is_retryable(response.status.code) and \
-                time.time() - start_time < 60 * 10:
-          logger.info("Model is still deploying, please wait...")
-          time.sleep(next(backoff_iterator))
-          break
+            method_signatures = None
+            if response.status.code == status_code_pb2.SUCCESS:
+                method_signatures = response.model_version.method_signatures
+            if response.status.code != status_code_pb2.SUCCESS:
+                raise Exception(f"Model failed with response {response!r}")
+            self._method_signatures = {}
+            for method_signature in method_signatures:
+                method_name = method_signature.name
+                # check for duplicate method names
+                if method_name in self._method_signatures:
+                    raise ValueError(f"Duplicate method name {method_name}")
+                self._method_signatures[method_name] = method_signature
+            if not self._method_signatures:  # if no method signatures, try to fetch from the model
+                self._fetch_signatures_backup()
+        except Exception:
+            # try to fetch from the model
+            self._fetch_signatures_backup()
+            if not self._method_signatures:
+                raise ValueError("Failed to fetch method signatures from model and backup method")
+
+    def _fetch_signatures_backup(self):
+        '''
+        This is a temporary method of fetching the method signatures from the model.
+
+        Returns:
+            Dict: The method signatures.
+        '''
+
+        request = service_pb2.PostModelOutputsRequest()
+        request.CopyFrom(self.request_template)
+        # request.model.model_version.output_info.params['_method_name'] = '_GET_SIGNATURES'
+        inp = request.inputs.add()  # empty input for this method
+        inp.data.parts.add()  # empty part for this input
+        inp.data.metadata['_method_name'] = '_GET_SIGNATURES'
+        start_time = time.time()
+        backoff_iterator = BackoffIterator(10)
+        while True:
+            response = self.STUB.PostModelOutputs(request)
+            if (
+                status_is_retryable(response.status.code) and time.time() - start_time < 60 * 10
+            ):  # 10 minutes
+                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 or (
+            response.status.code == status_code_pb2.SUCCESS
+            and response.outputs[0].data.text.raw == ''
+        ):
+            # return codes/values 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 Predict failed with response {response.status!r}")
-        else:
-          if not generation_started:
-            generation_started = True
-          yield response
+            raise Exception(f"Model failed with response {response!r}")
+        self._method_signatures = signatures_from_json(response.outputs[0].data.text.raw)
+
+    def _define_functions(self):
+        '''
+        Define the functions based on the method signatures.
+        '''
+        for method_name, method_signature in self._method_signatures.items():
+            # define the function in this client instance
+            if resources_pb2.RunnerMethodType.Name(method_signature.method_type) == 'UNARY_UNARY':
+                call_func = self._predict
+            elif (
+                resources_pb2.RunnerMethodType.Name(method_signature.method_type)
+                == 'UNARY_STREAMING'
+            ):
+                call_func = self._generate
+            elif (
+                resources_pb2.RunnerMethodType.Name(method_signature.method_type)
+                == 'STREAMING_STREAMING'
+            ):
+                call_func = self._stream
+            else:
+                raise ValueError(f"Unknown method type {method_signature.method_type}")
+
+            # method argnames, in order, collapsing nested keys to corresponding user function args
+            method_argnames = []
+            for var in method_signature.input_fields:
+                outer = var.name.split('.', 1)[0]
+                if outer in method_argnames:
+                    continue
+                method_argnames.append(outer)
+
+            def bind_f(method_name, method_argnames, call_func):
+                def f(*args, **kwargs):
+                    if len(args) > len(method_argnames):
+                        raise TypeError(
+                            f"{method_name}() takes {len(method_argnames)} positional arguments but {len(args)} were given"
+                        )
+
+                    if len(args) + len(kwargs) > len(method_argnames):
+                        raise TypeError(
+                            f"{method_name}() got an unexpected keyword argument {next(iter(kwargs))}"
+                        )
+                    if len(args) == 1 and (not kwargs) and isinstance(args[0], list):
+                        batch_inputs = args[0]
+                        # Validate each input is a dictionary
+                        is_batch_input_valid = all(
+                            isinstance(input, dict) for input in batch_inputs
+                        )
+                        if is_batch_input_valid and (not is_openai_chat_format(batch_inputs)):
+                            # If the batch input is valid, call the function with the batch inputs and the method name
+                            return call_func(batch_inputs, method_name)
+
+                    for name, arg in zip(
+                        method_argnames, args
+                    ):  # handle positional with zip shortest
+                        if name in kwargs:
+                            raise TypeError(f"Multiple values for argument {name}")
+                        kwargs[name] = arg
+                    return call_func(kwargs, method_name)
+
+                return f
+
+            # need to bind method_name to the value, not the mutating loop variable
+            f = bind_f(method_name, method_argnames, call_func)
+
+            # set names, annotations and docstrings
+            f.__name__ = method_name
+            f.__qualname__ = f'{self.__class__.__name__}.{method_name}'
+            input_annotations = code_script._get_annotations_source(method_signature)
+            return_annotation = input_annotations.pop('return', (None, None))[0]
+            sig = inspect.signature(f).replace(
+                parameters=[
+                    inspect.Parameter(k, inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=v[0])
+                    for k, v in input_annotations.items()
+                ],
+                return_annotation=return_annotation,
+            )
+            f.__signature__ = sig
+            f.__doc__ = method_signature.description
+            setattr(self, method_name, f)
+
+    def available_methods(self) -> List[str]:
+        """Get the available methods for this model.
+
+        Returns:
+            List[str]: The available methods.
+        """
+        if not self._defined:
+            self.fetch()
+        return self._method_signatures.keys()
+
+    def method_signature(self, method_name: str) -> str:
+        """Get the method signature for a method.
+
+        Args:
+            method_name (str): The name of the method.
+
+        Returns:
+            str: The method signature.
+        """
+        if not self._defined:
+            self.fetch()
+        return method_signatures.get_method_signature(self._method_signatures[method_name])
+
+    def generate_client_script(self) -> str:
+        """Generate a client script for this model.
+
+        Returns:
+            str: The client script.
+        """
+        if not self._defined:
+            self.fetch()
+        method_signatures = []
+        for _, method_signature in self._method_signatures.items():
+            method_signatures.append(method_signature)
+        return code_script.generate_client_script(
+            method_signatures,
+            user_id=self.request_template.user_app_id.user_id,
+            app_id=self.request_template.user_app_id.app_id,
+            model_id=self.request_template.model_id,
+        )
+
+    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?
+        method_name: str = 'predict',
+    ) -> Any:
+        input_signature = self._method_signatures[method_name].input_fields
+        output_signature = self._method_signatures[method_name].output_fields
+
+        batch_input = True
+        if isinstance(inputs, dict):
+            inputs = [inputs]
+            batch_input = False
+
+        proto_inputs = []
+        for input in inputs:
+            proto = resources_pb2.Input()
+
+            serialize(input, input_signature, proto.data)
+            proto_inputs.append(proto)
+
+        response = self._predict_by_proto(proto_inputs, method_name)
+
+        outputs = []
+        for output in response.outputs:
+            outputs.append(deserialize(output.data, output_signature, is_output=True))
+        if batch_input:
+            return outputs
+        return outputs[0]
+
+    def _predict_by_proto(
+        self,
+        inputs: List[resources_pb2.Input],
+        method_name: str = None,
+        inference_params: Dict = None,
+        output_config: Dict = None,
+    ) -> service_pb2.MultiOutputResponse:
+        """Predicts the model based on the given inputs.
+
+        Args:
+            inputs (List[resources_pb2.Input]): The inputs to predict.
+            method_name (str): The remote method name to call.
+            inference_params (Dict): Inference parameters to override.
+            output_config (Dict): Output configuration to override.
+
+        Returns:
+            service_pb2.MultiOutputResponse: The prediction response(s).
+        """
+        if not isinstance(inputs, list):
+            raise UserError('Invalid inputs, inputs must be a list of Input objects.')
+        if len(inputs) > MAX_MODEL_PREDICT_INPUTS:
+            raise UserError(f"Too many inputs. Max is {MAX_MODEL_PREDICT_INPUTS}.")
+
+        request = service_pb2.PostModelOutputsRequest()
+        request.CopyFrom(self.request_template)
+
+        request.inputs.extend(inputs)
+
+        if method_name:
+            # TODO put in new proto field?
+            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:
+            request.model.model_version.output_info.output_config.MergeFrom(
+                resources_pb2.OutputConfig(**output_config)
+            )
+
+        start_time = time.time()
+        backoff_iterator = BackoffIterator(10)
+        while True:
+            response = self.STUB.PostModelOutputs(request)
+            if (
+                status_is_retryable(response.status.code) and time.time() - start_time < 60 * 10
+            ):  # 10 minutes
+                logger.info("Model is still deploying, please wait...")
+                time.sleep(next(backoff_iterator))
+                continue
+
+            if response.status.code != status_code_pb2.SUCCESS:
+                raise Exception(f"Model predict failed with response {response!r}")
+            break
+        return response
+
+    def _generate(
+        self,
+        inputs,  # TODO set up functions according to fetched signatures?
+        method_name: str = 'generate',
+    ) -> Any:
+        input_signature = self._method_signatures[method_name].input_fields
+        output_signature = self._method_signatures[method_name].output_fields
+
+        batch_input = True
+        if isinstance(inputs, dict):
+            inputs = [inputs]
+            batch_input = False
+
+        proto_inputs = []
+        for input in inputs:
+            proto = resources_pb2.Input()
+            serialize(input, input_signature, proto.data)
+            proto_inputs.append(proto)
+
+        response_stream = self._generate_by_proto(proto_inputs, method_name)
+
+        for response in response_stream:
+            outputs = []
+            for output in response.outputs:
+                outputs.append(deserialize(output.data, output_signature, is_output=True))
+            if batch_input:
+                yield outputs
+            else:
+                yield outputs[0]
+
+    def _generate_by_proto(
+        self,
+        inputs: List[resources_pb2.Input],
+        method_name: str = None,
+        inference_params: Dict = {},
+        output_config: Dict = {},
+    ):
+        """Generate the stream output on model based on the given inputs.
+
+        Args:
+            inputs (list[Input]): The inputs to generate, must be less than 128.
+            method_name (str): The remote method name to call.
+            inference_params (dict): The inference params to override.
+            output_config (dict): The output config to override.
+        """
+        if not isinstance(inputs, list):
+            raise UserError('Invalid inputs, inputs must be a list of Input objects.')
+        if len(inputs) > MAX_MODEL_PREDICT_INPUTS:
+            raise UserError(
+                f"Too many inputs. Max is {MAX_MODEL_PREDICT_INPUTS}."
+            )  # TODO Use Chunker for inputs len > 128
+
+        request = service_pb2.PostModelOutputsRequest()
+        request.CopyFrom(self.request_template)
+
+        request.inputs.extend(inputs)
+
+        if method_name:
+            # TODO put in new proto field?
+            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:
+            request.model.model_version.output_info.output_config.MergeFromDict(output_config)
+
+        start_time = time.time()
+        backoff_iterator = BackoffIterator(10)
+        started = False
+        while not started:
+            stream_response = self.STUB.GenerateModelOutputs(request)
+            try:
+                response = next(stream_response)  # get the first response
+            except StopIteration:
+                raise Exception("Model Generate failed with no response")
+            if status_is_retryable(response.status.code) and time.time() - start_time < 60 * 10:
+                logger.info("Model is still deploying, please wait...")
+                time.sleep(next(backoff_iterator))
+                continue
+            if response.status.code != status_code_pb2.SUCCESS:
+                raise Exception(f"Model Generate failed with response {response.status!r}")
+            started = True
+
+        yield response  # yield the first response
+
+        for response in stream_response:
+            if response.status.code != status_code_pb2.SUCCESS:
+                raise Exception(f"Model Generate failed with response {response.status!r}")
+            yield response
+
+    def _stream(
+        self,
+        inputs,
+        method_name: str = 'stream',
+    ) -> Any:
+        input_signature = self._method_signatures[method_name].input_fields
+        output_signature = self._method_signatures[method_name].output_fields
+
+        if isinstance(inputs, list):
+            assert len(inputs) == 1, 'streaming methods do not support batched calls'
+            inputs = inputs[0]
+        assert isinstance(inputs, dict)
+        kwargs = inputs
+
+        # find the streaming vars in the input signature, and the streaming input python param
+        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)
+
+        def _input_proto_stream():
+            # first item contains all the inputs and the first stream item
+            proto = resources_pb2.Input()
+            try:
+                item = next(user_inputs_generator)
+            except StopIteration:
+                return  # no items to stream
+            kwargs[stream_argname] = item
+            serialize(kwargs, input_signature, proto.data)
+
+            yield proto
+
+            # subsequent items are just the stream items
+            for item in user_inputs_generator:
+                proto = resources_pb2.Input()
+                serialize({stream_argname: item}, [stream_sig], proto.data)
+                yield proto
+
+        response_stream = self._stream_by_proto(_input_proto_stream(), method_name)
+
+        for response in response_stream:
+            assert len(response.outputs) == 1, 'streaming methods must have exactly one output'
+            yield deserialize(response.outputs[0].data, output_signature, is_output=True)
+
+    def _req_iterator(
+        self,
+        input_iterator: Iterator[List[resources_pb2.Input]],
+        method_name: str = None,
+        inference_params: Dict = {},
+        output_config: Dict = {},
+    ):
+        request = service_pb2.PostModelOutputsRequest()
+        request.CopyFrom(self.request_template)
+        if inference_params:
+            request.model.model_version.output_info.params.update(inference_params)
+        if output_config:
+            request.model.model_version.output_info.output_config.MergeFromDict(output_config)
+        for inputs in input_iterator:
+            req = service_pb2.PostModelOutputsRequest()
+            req.CopyFrom(request)
+            if isinstance(inputs, list):
+                req.inputs.extend(inputs)
+            else:
+                req.inputs.append(inputs)
+            # TODO: put into new proto field?
+            if method_name:
+                for inp in req.inputs:
+                    inp.data.metadata['_method_name'] = method_name
+            yield req
+
+    def _stream_by_proto(
+        self,
+        inputs: Iterator[List[resources_pb2.Input]],
+        method_name: str = None,
+        inference_params: Dict = {},
+        output_config: Dict = {},
+    ):
+        """Generate the stream output on model based on the given stream of inputs."""
+        # if not isinstance(inputs, Iterator[List[Input]]):
+        #   raise UserError('Invalid inputs, inputs must be a iterator of list of Input objects.')
+
+        request = self._req_iterator(inputs, method_name, inference_params, output_config)
+
+        start_time = time.time()
+        backoff_iterator = BackoffIterator(10)
+        generation_started = False
+        while True:
+            if generation_started:
+                break
+            stream_response = self.STUB.StreamModelOutputs(request)
+            for response in stream_response:
+                if (
+                    status_is_retryable(response.status.code)
+                    and time.time() - start_time < 60 * 10
+                ):
+                    logger.info("Model is still deploying, please wait...")
+                    time.sleep(next(backoff_iterator))
+                    break
+                if response.status.code != status_code_pb2.SUCCESS:
+                    raise Exception(f"Model Predict failed with response {response.status!r}")
+                else:
+                    if not generation_started:
+                        generation_started = True
+                    yield response