basic_predictor.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from abc import abstractmethod
  15. from ....utils.subclass_register import AutoRegisterABCMetaClass
  16. from ....utils import logging
  17. from ...components.base import BaseComponent, ComponentsEngine
  18. from ...utils.pp_option import PaddlePredictorOption
  19. from ...utils.process_hook import generatorable_method
  20. from ..utils.predict_set import DeviceSetMixin, PPOptionSetMixin, BatchSizeSetMixin
  21. from .base_predictor import BasePredictor
  22. class BasicPredictor(
  23. BasePredictor,
  24. DeviceSetMixin,
  25. PPOptionSetMixin,
  26. BatchSizeSetMixin,
  27. metaclass=AutoRegisterABCMetaClass,
  28. ):
  29. __is_base = True
  30. def __init__(self, model_dir, config=None, device=None, pp_option=None):
  31. super().__init__(model_dir=model_dir, config=config)
  32. if not pp_option:
  33. pp_option = PaddlePredictorOption(model_name=self.model_name)
  34. if device:
  35. pp_option.device = device
  36. self._pp_option = pp_option
  37. self.components = {}
  38. self._build_components()
  39. self.engine = ComponentsEngine(self.components)
  40. logging.debug(f"{self.__class__.__name__}: {self.model_dir}")
  41. def apply(self, x):
  42. """predict"""
  43. yield from self._generate_res(self.engine(x))
  44. @generatorable_method
  45. def _generate_res(self, batch_data):
  46. return [{"result": self._pack_res(data)} for data in batch_data]
  47. def _add_component(self, cmps):
  48. if not isinstance(cmps, list):
  49. cmps = [cmps]
  50. for cmp in cmps:
  51. if not isinstance(cmp, (list, tuple)):
  52. key = cmp.name
  53. else:
  54. assert len(cmp) == 2
  55. key = cmp[0]
  56. cmp = cmp[1]
  57. assert isinstance(key, str)
  58. assert isinstance(cmp, BaseComponent)
  59. assert (
  60. key not in self.components
  61. ), f"The key ({key}) has been used: {self.components}!"
  62. self.components[key] = cmp
  63. def set_predictor(self, **kwargs):
  64. for k in kwargs:
  65. if self._has_setter(k):
  66. setattr(self, k, kwargs[k])
  67. else:
  68. raise Exception(
  69. f"The arg({k}) is not supported to specify in predict() func! Only supports: {self._get_settable_attributes}"
  70. )
  71. def _has_setter(self, attr):
  72. prop = getattr(self.__class__, attr, None)
  73. return isinstance(prop, property) and prop.fset is not None
  74. def _get_settable_attributes(self):
  75. return [
  76. name
  77. for name, prop in vars(self.__class__).items()
  78. if isinstance(prop, property) and prop.fset is not None
  79. ]
  80. @abstractmethod
  81. def _build_components(self):
  82. raise NotImplementedError
  83. @abstractmethod
  84. def _pack_res(self, data):
  85. raise NotImplementedError