predictor.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123
  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. import lazy_paddle as paddle
  15. from ....utils.func_register import FuncRegister
  16. from ...common.batch_sampler import AudioBatchSampler
  17. from ..base import BasicPredictor
  18. from .result import WhisperResult
  19. from ...utils.io import AudioReader
  20. from ....modules.multilingual_speech_recognition.model_list import MODELS
  21. from ....utils.download import download_and_extract
  22. class WhisperPredictor(BasicPredictor):
  23. entities = MODELS
  24. def __init__(self, *args, **kwargs):
  25. """Initializes WhisperPredictor.
  26. Args:
  27. *args: Arbitrary positional arguments passed to the superclass.
  28. **kwargs: Arbitrary keyword arguments passed to the superclass.
  29. """
  30. super().__init__(*args, **kwargs)
  31. self.audio_reader = self._build()
  32. download_and_extract(
  33. self.config["resource_path"], self.config["resource_dir"], "assets"
  34. )
  35. def _build_batch_sampler(self):
  36. """Builds and returns an AudioBatchSampler instance.
  37. Returns:
  38. AudioBatchSampler: An instance of AudioBatchSampler.
  39. """
  40. return AudioBatchSampler()
  41. def _get_result_class(self):
  42. """Returns the result class, WhisperResult.
  43. Returns:
  44. type: The WhisperResult class.
  45. """
  46. return WhisperResult
  47. def _build(self):
  48. """Build the model, audio reader based on the configuration.
  49. Returns:
  50. AudioReader: An instance of AudioReader.
  51. """
  52. from .processors import (
  53. ModelDimensions,
  54. Whisper,
  55. LANGUAGES,
  56. TO_LANGUAGE_CODE,
  57. )
  58. # build model
  59. model_dict = paddle.load(self.config["model_file"])
  60. dims = ModelDimensions(**model_dict["dims"])
  61. self.model = Whisper(dims)
  62. self.model.load_dict(model_dict)
  63. self.model.eval()
  64. # build audio reader
  65. audio_reader = AudioReader(backend="wav")
  66. return audio_reader
  67. def process(self, batch_data):
  68. """
  69. Process a batch of data through the preprocessing, inference, and postprocessing.
  70. Args:
  71. batch_data (List[Union[str, np.ndarray], ...]): A batch of input data (e.g., audio file paths).
  72. Returns:
  73. dict: A dictionary containing the input path and result. The result include 'text', 'segments' and 'language'.
  74. """
  75. from .processors import log_mel_spectrogram
  76. # load mel_filters from resource_dir and extract feature for audio
  77. audio, sample_rate = self.audio_reader.read(batch_data[0])
  78. audio = paddle.to_tensor(audio)
  79. audio = audio[:, 0]
  80. audio = log_mel_spectrogram(audio, resource_path=self.config["resource_dir"])
  81. # model inference
  82. result = self.model.transcribe(
  83. audio,
  84. verbose=self.config["verbose"],
  85. task=self.config["task"],
  86. language=self.config["language"],
  87. resource_path=self.config["resource_dir"],
  88. temperature=self.config["temperature"],
  89. compression_ratio_threshold=self.config["compression_ratio_threshold"],
  90. logprob_threshold=self.config["logprob_threshold"],
  91. best_of=self.config["best_of"],
  92. beam_size=self.config["beam_size"],
  93. patience=self.config["patience"],
  94. length_penalty=self.config["length_penalty"],
  95. initial_prompt=self.config["initial_prompt"],
  96. condition_on_previous_text=self.config["condition_on_previous_text"],
  97. no_speech_threshold=self.config["no_speech_threshold"],
  98. )
  99. return {
  100. "input_path": batch_data,
  101. "result": [result],
  102. }