class_docstr: A DoFn that loads tflite models and predicts.
class_docstr_tok: ["A","DoFn","that","loads","tflite","models","and","predicts","."]
docstr: Invokes the tflite model on the provided inputs and stores the result.
docstr_tok: ["Invokes","the","tflite","model","on","the","provided","inputs","and","stores","the","result","."]
code: def _batch_reducible_process(
      self, element: types.Extracts
  ) -> Sequence[types.Extracts]:
      """Invokes the tflite model on the provided inputs and stores the result."""
      result = copy.copy(element)
  
      batched_features = element[constants.FEATURES_KEY]
      batch_size = util.batch_size(batched_features)
  
      for spec in self._eval_config.model_specs:
          model_name = (
              spec.name
              if len(self._eval_config.model_specs) > 1
              else ""
          )
          if model_name not in self._loaded_models:
              raise ValueError(
                  'model for "{}" not found: eval_config={}'.format(
                      spec.name, self._eval_config
                  )
              )
  
          interpreter = self._interpreters[model_name]
  
          input_details = interpreter.get_input_details()
          output_details = interpreter.get_output_details()
  
          input_features = collections.defaultdict(list)
          for i in input_details:
              input_name = self._get_input_name_from_input_detail(i)
              # The batch dimension is the specific batch size of the last time the
              # model was invoked. Set it to 1 to "reset".
              input_shape = [1] + list(i["shape"])[1:]
              input_type = i[_DTYPE]
              for idx in range(batch_size):
                  if input_name in batched_features:
                      value = batched_features[input_name][idx]
                  else:
                      value = None
                  if value is None or np.any(np.equal(value, None)):
                      default = (
                          -1
                          if input_type in [np.float32, np.int64]
                          else ""
                      )
                      value = np.full(
                          input_shape, default, dtype=input_type
                      )
                      logging.log_every_n(
                          logging.WARNING,
                          "Feature %s not found. Setting default value.",
                          100,
                          input_name,
                      )
                  else:
                      value = np.reshape(value, input_shape)
                  input_features[input_name].append(value)
              # Concatenate with numpy to avoid implicit conversion to tf.Tensor
              # which causes byte-string inputs to fail the set_tensor call.
              input_features[input_name] = np.concatenate(
                  input_features[input_name], axis=0
              )
              if np.shape(input_features[input_name]) != tuple(
                  i["shape"]
              ):
                  interpreter.resize_tensor_input(
                      i[_INDEX], np.shape(input_features[input_name])
                  )
          interpreter.allocate_tensors()
  
          for i in input_details:
              input_name = self._get_input_name_from_input_detail(i)
              params = i[_QUANTIZATION_PARAMETERS]
              interpreter.set_tensor(
                  i[_INDEX],
                  self._quantize(
                      input_features[input_name],
                      params[_SCALES],
                      params[_ZERO_POINTS],
                      i[_DTYPE],
                  ),
              )
          interpreter.invoke()
  
          outputs = {}
          for o in output_details:
              tensor = interpreter.get_tensor(o[_INDEX])
              params = o[_QUANTIZATION_PARAMETERS]
              outputs[o["name"]] = self._dequantize(
                  tensor, params[_SCALES], params[_ZERO_POINTS]
              )
  
          for v in outputs.values():
              if len(v) != batch_size:
                  raise ValueError(
                      "Did not get the expected number of results."
                  )
  
          if len(outputs) == 1:
              outputs = list(outputs.values())[0]
  
          if len(self._eval_config.model_specs) == 1:
              result[constants.PREDICTIONS_KEY] = outputs
          else:
              if constants.PREDICTIONS_KEY not in result:
                  result[constants.PREDICTIONS_KEY] = {}
              result[constants.PREDICTIONS_KEY][spec.name] = outputs
      return [result]
  
imports: import collections
  import copy
  from typing import Callable, Dict, Sequence, Union
  from absl import logging
  import apache_beam
  import numpy
  import tensorflow
  from tensorflow_model_analysis import constants
  from tensorflow_model_analysis.api import types
  from tensorflow_model_analysis.extractors import extractor
  from tensorflow_model_analysis.proto import config_pb2
  from tensorflow_model_analysis.utils import model_util
  from tensorflow_model_analysis.utils import util
label_desc: Use image node_id 7 for calling the TFLitePredictionDoFn obj's underlying member method code with example usage: obj._batch_reducible_process(element) and returns: list