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

clarifai/runners/models/model_class.py

@@ -1,10 +1,10 @@
-import collections.abc as abc
 import inspect
 import itertools
 import logging
 import os
 import traceback
 from abc import ABC
+from collections import abc
 from typing import Any, Dict, Iterator, List
 
 from clarifai_grpc.grpc.api import resources_pb2, service_pb2
@@ -13,9 +13,13 @@ from google.protobuf import json_format
 
 from clarifai.runners.utils import data_types
 from clarifai.runners.utils.data_utils import DataConverter
-from clarifai.runners.utils.method_signatures import (build_function_signature, deserialize,
-                                                      get_stream_from_signature, serialize,
-                                                      signatures_to_json)
+from clarifai.runners.utils.method_signatures import (
+    build_function_signature,
+    deserialize,
+    get_stream_from_signature,
+    serialize,
+    signatures_to_json,
+)
 
 _METHOD_INFO_ATTR = '_cf_method_info'
 
@@ -23,337 +27,372 @@ _RAISE_EXCEPTIONS = os.getenv("RAISE_EXCEPTIONS", "false").lower() in ("true", "
 
 
 class ModelClass(ABC):
-  '''
-  Base class for model classes that can be run as a service.
-
-  Define predict, generate, or stream methods using the @ModelClass.method decorator.
-
-  Example:
-
-    from clarifai.runners.model_class import ModelClass
-    from clarifai.runners.utils.data_types import NamedFields
-    from typing import List, Iterator
-
-    class MyModel(ModelClass):
-
-      @ModelClass.method
-      def predict(self, x: str, y: int) -> List[str]:
-        return [x] * y
-
-      @ModelClass.method
-      def generate(self, x: str, y: int) -> Iterator[str]:
-        for i in range(y):
-          yield x + str(i)
-
-      @ModelClass.method
-      def stream(self, input_stream: Iterator[NamedFields(x=str, y=int)]) -> Iterator[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 set_output_context(self, prompt_tokens=None, completion_tokens=None):
-    """This is used to set the prompt and completion tokens in the Output proto"""
-    self._prompt_tokens = prompt_tokens
-    self._completion_tokens = completion_tokens
-
-  def load_model(self):
-    """Load the model."""
-
-  def _handle_get_signatures_request(self) -> service_pb2.MultiOutputResponse:
-    methods = self._get_method_info()
-    signatures = {method.name: method.signature for method in methods.values()}
-    resp = service_pb2.MultiOutputResponse(status=status_pb2.Status(code=status_code_pb2.SUCCESS))
-    output = resp.outputs.add()
-    output.status.code = status_code_pb2.SUCCESS
-    output.data.text.raw = signatures_to_json(signatures)
-    return resp
-
-  def _batch_predict(self, method, inputs: List[Dict[str, Any]]) -> List[Any]:
-    """Batch predict method for multiple inputs."""
-    outputs = []
-    for input in inputs:
-      output = method(**input)
-      outputs.append(output)
-    return outputs
-
-  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):
-      yield outputs
-
-  def predict_wrapper(
-      self, request: service_pb2.PostModelOutputsRequest) -> service_pb2.MultiOutputResponse:
-    outputs = []
-    try:
-      # TODO add method name field to proto
-      method_name = 'predict'
-      inference_params = get_inference_params(request)
-      if len(request.inputs) > 0 and '_method_name' in request.inputs[0].data.metadata:
-        method_name = request.inputs[0].data.metadata['_method_name']
-      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():
-        raise ValueError(f"Method {method_name} not found in model class")
-      method = getattr(self, method_name)
-      method_info = method._cf_method_info
-      signature = method_info.signature
-      python_param_types = method_info.python_param_types
-      for input in request.inputs:
-        # check if input is in old format
-        is_convert = DataConverter.is_old_format(input.data)
-        if is_convert:
-          # convert to new format
-          new_data = DataConverter.convert_input_data_to_new_format(input.data,
-                                                                    signature.input_fields)
-          input.data.CopyFrom(new_data)
-      # convert inputs to python types
-      inputs = self._convert_input_protos_to_python(request.inputs, inference_params,
-                                                    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.output_fields, convert_old_format=is_convert))
-      else:
-        outputs = self._batch_predict(method, inputs)
-        outputs = [
-            self._convert_output_to_proto(
-                output, signature.output_fields, convert_old_format=is_convert)
-            for output in outputs
-        ]
-
-      return service_pb2.MultiOutputResponse(
-          outputs=outputs, status=status_pb2.Status(code=status_code_pb2.SUCCESS))
-    except Exception as e:
-      if _RAISE_EXCEPTIONS:
-        raise
-      logging.exception("Error in predict")
-      return service_pb2.MultiOutputResponse(status=status_pb2.Status(
-          code=status_code_pb2.FAILURE,
-          details=str(e),
-          stack_trace=traceback.format_exc().split('\n')))
-
-  def generate_wrapper(self, request: service_pb2.PostModelOutputsRequest
-                      ) -> Iterator[service_pb2.MultiOutputResponse]:
-    try:
-      method_name = 'generate'
-      inference_params = get_inference_params(request)
-      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
-      for input in request.inputs:
-        # check if input is in old format
-        is_convert = DataConverter.is_old_format(input.data)
-        if is_convert:
-          # convert to new format
-          new_data = DataConverter.convert_input_data_to_new_format(input.data,
-                                                                    signature.input_fields)
-          input.data.CopyFrom(new_data)
-      inputs = self._convert_input_protos_to_python(request.inputs, inference_params,
-                                                    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.output_fields,
-              proto=resp.outputs.add(),
-              convert_old_format=is_convert)
-          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.output_fields,
-                proto=resp.outputs.add(),
-                convert_old_format=is_convert)
-          resp.status.code = status_code_pb2.SUCCESS
-          yield resp
-    except Exception as e:
-      if _RAISE_EXCEPTIONS:
-        raise
-      logging.exception("Error in generate")
-      yield service_pb2.MultiOutputResponse(status=status_pb2.Status(
-          code=status_code_pb2.FAILURE,
-          details=str(e),
-          stack_trace=traceback.format_exc().split('\n')))
-
-  def stream_wrapper(self, request_iterator: Iterator[service_pb2.PostModelOutputsRequest]
-                    ) -> Iterator[service_pb2.MultiOutputResponse]:
-    try:
-      request = next(request_iterator)  # get first request to determine method
-      assert len(request.inputs) == 1, "Streaming requires exactly one input"
-
-      method_name = 'stream'
-      inference_params = get_inference_params(request)
-      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_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
-
-      for input in request.inputs:
-        # check if input is in old format
-        is_convert = DataConverter.is_old_format(input.data)
-        if is_convert:
-          # convert to new format
-          new_data = DataConverter.convert_input_data_to_new_format(input.data,
-                                                                    signature.input_fields)
-          input.data.CopyFrom(new_data)
-      # convert all inputs for the first request, including the first stream value
-      inputs = self._convert_input_protos_to_python(request.inputs, inference_params,
-                                                    signature.input_fields, python_param_types)
-      kwargs = inputs[0]
-
-      # first streaming item
-      first_item = kwargs.pop(stream_argname)
-
-      # streaming generator
-      def InputStream():
-        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, inference_params,
-                                                      [stream_sig], python_param_types)
-          item = item[0][stream_argname]
-          yield item
-
-      # add stream generator back to the input kwargs
-      kwargs[stream_argname] = InputStream()
-
-      for output in method(**kwargs):
-        resp = service_pb2.MultiOutputResponse()
-        self._convert_output_to_proto(
-            output,
-            signature.output_fields,
-            proto=resp.outputs.add(),
-            convert_old_format=is_convert)
-        resp.status.code = status_code_pb2.SUCCESS
-        yield resp
-    except Exception as e:
-      if _RAISE_EXCEPTIONS:
-        raise
-      logging.exception("Error in stream")
-      yield service_pb2.MultiOutputResponse(status=status_pb2.Status(
-          code=status_code_pb2.FAILURE,
-          details=str(e),
-          stack_trace=traceback.format_exc().split('\n')))
-
-  def _convert_input_protos_to_python(self, inputs: List[resources_pb2.Input],
-                                      inference_params: dict,
-                                      variables_signature: List[resources_pb2.ModelTypeField],
-                                      python_param_types) -> List[Dict[str, Any]]:
-    result = []
-    for input in inputs:
-      kwargs = deserialize(input.data, variables_signature, inference_params)
-      # dynamic cast to annotated types
-      for k, v in kwargs.items():
-        if k not in python_param_types:
-          continue
-
-        if hasattr(python_param_types[k],
-                   "__args__") and (getattr(python_param_types[k], "__origin__",
-                                            None) in [abc.Iterator, abc.Generator, abc.Iterable]):
-          # get the type of the items in the stream
-          stream_type = python_param_types[k].__args__[0]
-
-          kwargs[k] = data_types.cast(v, stream_type)
-        else:
-          kwargs[k] = data_types.cast(v, python_param_types[k])
-      result.append(kwargs)
-    return result
-
-  def _convert_output_to_proto(self,
-                               output: Any,
-                               variables_signature: List[resources_pb2.ModelTypeField],
-                               proto=None,
-                               convert_old_format=False) -> resources_pb2.Output:
-    if proto is None:
-      proto = resources_pb2.Output()
-    serialize({'return': output}, variables_signature, proto.data, is_output=True)
-    if convert_old_format:
-      # convert to old format
-      data = DataConverter.convert_output_data_to_old_format(proto.data)
-      proto.data.CopyFrom(data)
-    proto.status.code = status_code_pb2.SUCCESS
-    if hasattr(self, "_prompt_tokens") and self._prompt_tokens is not None:
-      proto.prompt_tokens = self._prompt_tokens
-    if hasattr(self, "_completion_tokens") and self._completion_tokens is not None:
-      proto.completion_tokens = self._completion_tokens
-    self._prompt_tokens = None
-    self._completion_tokens = None
-    return proto
-
-  @classmethod
-  def _register_model_methods(cls):
-    # go up the class hierarchy to find all decorated methods, and add to registry of current class
-    methods = {}
-    for base in reversed(cls.__mro__):
-      for name, method in base.__dict__.items():
-        method_info = getattr(method, _METHOD_INFO_ATTR, None)
-        if not method_info:  # regular function, not a model method
-          continue
-        methods[name] = method_info
-    # check for generic predict(request) -> response, etc. methods
-    #for name in ('predict', 'generate', 'stream'):
-    #  if hasattr(cls, name):
-    #    method = getattr(cls, name)
-    #    if not hasattr(method, _METHOD_INFO_ATTR):  # not already put in registry
-    #      methods[name] = _MethodInfo(method)
-    # set method table for this class in the registry
-    return methods
-
-  @classmethod
-  def _get_method_info(cls, func_name=None):
-    if not hasattr(cls, _METHOD_INFO_ATTR):
-      setattr(cls, _METHOD_INFO_ATTR, cls._register_model_methods())
-    method_info = getattr(cls, _METHOD_INFO_ATTR)
-    if func_name:
-      return method_info[func_name]
-    return method_info
+    '''
+    Base class for model classes that can be run as a service.
+
+    Define predict, generate, or stream methods using the @ModelClass.method decorator.
+
+    Example:
+
+      from clarifai.runners.model_class import ModelClass
+      from clarifai.runners.utils.data_types import NamedFields
+      from typing import List, Iterator
+
+      class MyModel(ModelClass):
+
+        @ModelClass.method
+        def predict(self, x: str, y: int) -> List[str]:
+          return [x] * y
+
+        @ModelClass.method
+        def generate(self, x: str, y: int) -> Iterator[str]:
+          for i in range(y):
+            yield x + str(i)
+
+        @ModelClass.method
+        def stream(self, input_stream: Iterator[NamedFields(x=str, y=int)]) -> Iterator[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 set_output_context(self, prompt_tokens=None, completion_tokens=None):
+        """This is used to set the prompt and completion tokens in the Output proto"""
+        self._prompt_tokens = prompt_tokens
+        self._completion_tokens = completion_tokens
+
+    def load_model(self):
+        """Load the model."""
+
+    def _handle_get_signatures_request(self) -> service_pb2.MultiOutputResponse:
+        methods = self._get_method_info()
+        signatures = {method.name: method.signature for method in methods.values()}
+        resp = service_pb2.MultiOutputResponse(
+            status=status_pb2.Status(code=status_code_pb2.SUCCESS)
+        )
+        output = resp.outputs.add()
+        output.status.code = status_code_pb2.SUCCESS
+        output.data.text.raw = signatures_to_json(signatures)
+        return resp
+
+    def _batch_predict(self, method, inputs: List[Dict[str, Any]]) -> List[Any]:
+        """Batch predict method for multiple inputs."""
+        outputs = []
+        for input in inputs:
+            output = method(**input)
+            outputs.append(output)
+        return outputs
+
+    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):
+            yield outputs
+
+    def predict_wrapper(
+        self, request: service_pb2.PostModelOutputsRequest
+    ) -> service_pb2.MultiOutputResponse:
+        outputs = []
+        try:
+            # TODO add method name field to proto
+            method_name = 'predict'
+            inference_params = get_inference_params(request)
+            if len(request.inputs) > 0 and '_method_name' in request.inputs[0].data.metadata:
+                method_name = request.inputs[0].data.metadata['_method_name']
+            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():
+                raise ValueError(f"Method {method_name} not found in model class")
+            method = getattr(self, method_name)
+            method_info = method._cf_method_info
+            signature = method_info.signature
+            python_param_types = method_info.python_param_types
+            for input in request.inputs:
+                # check if input is in old format
+                is_convert = DataConverter.is_old_format(input.data)
+                if is_convert:
+                    # convert to new format
+                    new_data = DataConverter.convert_input_data_to_new_format(
+                        input.data, signature.input_fields
+                    )
+                    input.data.CopyFrom(new_data)
+            # convert inputs to python types
+            inputs = self._convert_input_protos_to_python(
+                request.inputs, inference_params, 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.output_fields, convert_old_format=is_convert
+                    )
+                )
+            else:
+                outputs = self._batch_predict(method, inputs)
+                outputs = [
+                    self._convert_output_to_proto(
+                        output, signature.output_fields, convert_old_format=is_convert
+                    )
+                    for output in outputs
+                ]
+
+            return service_pb2.MultiOutputResponse(
+                outputs=outputs, status=status_pb2.Status(code=status_code_pb2.SUCCESS)
+            )
+        except Exception as e:
+            if _RAISE_EXCEPTIONS:
+                raise
+            logging.exception("Error in predict")
+            return service_pb2.MultiOutputResponse(
+                status=status_pb2.Status(
+                    code=status_code_pb2.FAILURE,
+                    details=str(e),
+                    stack_trace=traceback.format_exc().split('\n'),
+                )
+            )
+
+    def generate_wrapper(
+        self, request: service_pb2.PostModelOutputsRequest
+    ) -> Iterator[service_pb2.MultiOutputResponse]:
+        try:
+            method_name = 'generate'
+            inference_params = get_inference_params(request)
+            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
+            for input in request.inputs:
+                # check if input is in old format
+                is_convert = DataConverter.is_old_format(input.data)
+                if is_convert:
+                    # convert to new format
+                    new_data = DataConverter.convert_input_data_to_new_format(
+                        input.data, signature.input_fields
+                    )
+                    input.data.CopyFrom(new_data)
+            inputs = self._convert_input_protos_to_python(
+                request.inputs, inference_params, 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.output_fields,
+                        proto=resp.outputs.add(),
+                        convert_old_format=is_convert,
+                    )
+                    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.output_fields,
+                            proto=resp.outputs.add(),
+                            convert_old_format=is_convert,
+                        )
+                    resp.status.code = status_code_pb2.SUCCESS
+                    yield resp
+        except Exception as e:
+            if _RAISE_EXCEPTIONS:
+                raise
+            logging.exception("Error in generate")
+            yield service_pb2.MultiOutputResponse(
+                status=status_pb2.Status(
+                    code=status_code_pb2.FAILURE,
+                    details=str(e),
+                    stack_trace=traceback.format_exc().split('\n'),
+                )
+            )
+
+    def stream_wrapper(
+        self, request_iterator: Iterator[service_pb2.PostModelOutputsRequest]
+    ) -> Iterator[service_pb2.MultiOutputResponse]:
+        try:
+            request = next(request_iterator)  # get first request to determine method
+            assert len(request.inputs) == 1, "Streaming requires exactly one input"
+
+            method_name = 'stream'
+            inference_params = get_inference_params(request)
+            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_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
+
+            for input in request.inputs:
+                # check if input is in old format
+                is_convert = DataConverter.is_old_format(input.data)
+                if is_convert:
+                    # convert to new format
+                    new_data = DataConverter.convert_input_data_to_new_format(
+                        input.data, signature.input_fields
+                    )
+                    input.data.CopyFrom(new_data)
+            # convert all inputs for the first request, including the first stream value
+            inputs = self._convert_input_protos_to_python(
+                request.inputs, inference_params, signature.input_fields, python_param_types
+            )
+            kwargs = inputs[0]
+
+            # first streaming item
+            first_item = kwargs.pop(stream_argname)
+
+            # streaming generator
+            def InputStream():
+                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, inference_params, [stream_sig], python_param_types
+                    )
+                    item = item[0][stream_argname]
+                    yield item
+
+            # add stream generator back to the input kwargs
+            kwargs[stream_argname] = InputStream()
+
+            for output in method(**kwargs):
+                resp = service_pb2.MultiOutputResponse()
+                self._convert_output_to_proto(
+                    output,
+                    signature.output_fields,
+                    proto=resp.outputs.add(),
+                    convert_old_format=is_convert,
+                )
+                resp.status.code = status_code_pb2.SUCCESS
+                yield resp
+        except Exception as e:
+            if _RAISE_EXCEPTIONS:
+                raise
+            logging.exception("Error in stream")
+            yield service_pb2.MultiOutputResponse(
+                status=status_pb2.Status(
+                    code=status_code_pb2.FAILURE,
+                    details=str(e),
+                    stack_trace=traceback.format_exc().split('\n'),
+                )
+            )
+
+    def _convert_input_protos_to_python(
+        self,
+        inputs: List[resources_pb2.Input],
+        inference_params: dict,
+        variables_signature: List[resources_pb2.ModelTypeField],
+        python_param_types,
+    ) -> List[Dict[str, Any]]:
+        result = []
+        for input in inputs:
+            kwargs = deserialize(input.data, variables_signature, inference_params)
+            # dynamic cast to annotated types
+            for k, v in kwargs.items():
+                if k not in python_param_types:
+                    continue
+
+                if hasattr(python_param_types[k], "__args__") and (
+                    getattr(python_param_types[k], "__origin__", None)
+                    in [abc.Iterator, abc.Generator, abc.Iterable]
+                ):
+                    # get the type of the items in the stream
+                    stream_type = python_param_types[k].__args__[0]
+
+                    kwargs[k] = data_types.cast(v, stream_type)
+                else:
+                    kwargs[k] = data_types.cast(v, python_param_types[k])
+            result.append(kwargs)
+        return result
+
+    def _convert_output_to_proto(
+        self,
+        output: Any,
+        variables_signature: List[resources_pb2.ModelTypeField],
+        proto=None,
+        convert_old_format=False,
+    ) -> resources_pb2.Output:
+        if proto is None:
+            proto = resources_pb2.Output()
+        serialize({'return': output}, variables_signature, proto.data, is_output=True)
+        if convert_old_format:
+            # convert to old format
+            data = DataConverter.convert_output_data_to_old_format(proto.data)
+            proto.data.CopyFrom(data)
+        proto.status.code = status_code_pb2.SUCCESS
+        if hasattr(self, "_prompt_tokens") and self._prompt_tokens is not None:
+            proto.prompt_tokens = self._prompt_tokens
+        if hasattr(self, "_completion_tokens") and self._completion_tokens is not None:
+            proto.completion_tokens = self._completion_tokens
+        self._prompt_tokens = None
+        self._completion_tokens = None
+        return proto
+
+    @classmethod
+    def _register_model_methods(cls):
+        # go up the class hierarchy to find all decorated methods, and add to registry of current class
+        methods = {}
+        for base in reversed(cls.__mro__):
+            for name, method in base.__dict__.items():
+                method_info = getattr(method, _METHOD_INFO_ATTR, None)
+                if not method_info:  # regular function, not a model method
+                    continue
+                methods[name] = method_info
+        # check for generic predict(request) -> response, etc. methods
+        # for name in ('predict', 'generate', 'stream'):
+        #  if hasattr(cls, name):
+        #    method = getattr(cls, name)
+        #    if not hasattr(method, _METHOD_INFO_ATTR):  # not already put in registry
+        #      methods[name] = _MethodInfo(method)
+        # set method table for this class in the registry
+        return methods
+
+    @classmethod
+    def _get_method_info(cls, func_name=None):
+        if not hasattr(cls, _METHOD_INFO_ATTR):
+            setattr(cls, _METHOD_INFO_ATTR, cls._register_model_methods())
+        method_info = getattr(cls, _METHOD_INFO_ATTR)
+        if func_name:
+            return method_info[func_name]
+        return method_info
 
 
 # Helper function to get the inference params
 def get_inference_params(request) -> dict:
-  """Get the inference params from the request."""
-  inference_params = {}
-  if request.model.model_version.id != "":
-    output_info = request.model.model_version.output_info
-    output_info = json_format.MessageToDict(output_info, preserving_proto_field_name=True)
-    if "params" in output_info:
-      inference_params = output_info["params"]
-  return inference_params
+    """Get the inference params from the request."""
+    inference_params = {}
+    if request.model.model_version.id != "":
+        output_info = request.model.model_version.output_info
+        output_info = json_format.MessageToDict(output_info, preserving_proto_field_name=True)
+        if "params" in output_info:
+            inference_params = output_info["params"]
+    return inference_params
 
 
 class _MethodInfo:
-
-  def __init__(self, method):
-    self.name = method.__name__
-    self.signature = build_function_signature(method)
-    self.python_param_types = {
-        p.name: p.annotation
-        for p in inspect.signature(method).parameters.values()
-        if p.annotation != inspect.Parameter.empty
-    }
-    self.python_param_types.pop('self', None)
+    def __init__(self, method):
+        self.name = method.__name__
+        self.signature = build_function_signature(method)
+        self.python_param_types = {
+            p.name: p.annotation
+            for p in inspect.signature(method).parameters.values()
+            if p.annotation != inspect.Parameter.empty
+        }
+        self.python_param_types.pop('self', None)