ts_fc.py 2.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465
  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 typing import Any, Dict, List, Union
  15. import ultra_infer as ui
  16. import pandas as pd
  17. from paddlex.inference.common.batch_sampler import TSBatchSampler
  18. from paddlex.inference.models.ts_forecasting.result import TSFcResult
  19. from paddlex.modules.ts_forecast.model_list import MODELS
  20. from paddlex_hpi.models.base import TSPredictor
  21. class TSFcPredictor(TSPredictor):
  22. entities = MODELS
  23. def _build_batch_sampler(self) -> TSBatchSampler:
  24. return TSBatchSampler()
  25. def _get_result_class(self) -> type:
  26. return TSFcResult
  27. def _build_ui_model(
  28. self, option: ui.RuntimeOption
  29. ) -> ui.ts.forecasting.PyOnlyForecastingModel:
  30. model = ui.ts.forecasting.PyOnlyForecastingModel(
  31. str(self.model_path),
  32. str(self.params_path),
  33. str(self.config_path),
  34. runtime_option=option,
  35. )
  36. return model
  37. def process(self, batch_data: List[Union[str, pd.DataFrame]]) -> Dict[str, Any]:
  38. batch_raw_ts = self._data_reader(ts_list=batch_data)
  39. ui_results = self._ui_model.batch_predict(batch_raw_ts)
  40. forecast_list = []
  41. for ui_result in ui_results:
  42. data_dict = {
  43. ui_result.col_names[i]: ui_result.data[i]
  44. for i in range(len(ui_result.col_names))
  45. }
  46. forecast = pd.DataFrame.from_dict(data_dict)
  47. forecast.index = ui_result.dates
  48. forecast.index.name = "date"
  49. forecast_list.append(forecast)
  50. return {
  51. "input_path": batch_data,
  52. "input_ts": batch_raw_ts,
  53. "forecast": forecast_list,
  54. }