predictor.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134
  1. # Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
  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. import numpy as np
  16. from ....modules.multilingual_speech_recognition.model_list import MODELS
  17. from ....utils.download import download_and_extract
  18. from ...common.batch_sampler import AudioBatchSampler
  19. from ...utils.io import AudioReader
  20. from ..base import BasePredictor
  21. from .result import WhisperResult
  22. class WhisperPredictor(BasePredictor):
  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(self.config["resource_path"], self.model_dir, "assets")
  33. def _build_batch_sampler(self):
  34. """Builds and returns an AudioBatchSampler instance.
  35. Returns:
  36. AudioBatchSampler: An instance of AudioBatchSampler.
  37. """
  38. return AudioBatchSampler()
  39. def _get_result_class(self):
  40. """Returns the result class, WhisperResult.
  41. Returns:
  42. type: The WhisperResult class.
  43. """
  44. return WhisperResult
  45. def _build(self):
  46. """Build the model, audio reader based on the configuration.
  47. Returns:
  48. AudioReader: An instance of AudioReader.
  49. """
  50. from .processors import ModelDimensions, Whisper
  51. # build model
  52. model_file = (self.model_dir / f"{self.MODEL_FILE_PREFIX}.pdparams").as_posix()
  53. model_dict = paddle.load(model_file)
  54. dims = ModelDimensions(**model_dict["dims"])
  55. self.model = Whisper(dims)
  56. self.model.load_dict(model_dict)
  57. self.model.eval()
  58. # build audio reader
  59. audio_reader = AudioReader(backend="wav")
  60. return audio_reader
  61. def process(self, batch_data):
  62. """
  63. Process a batch of data through the preprocessing, inference, and postprocessing.
  64. Args:
  65. batch_data (List[Union[str, np.ndarray], ...]): A batch of input data (e.g., audio file paths).
  66. Returns:
  67. dict: A dictionary containing the input path and result. The result include 'text', 'segments' and 'language'.
  68. """
  69. from .processors import log_mel_spectrogram
  70. # load mel_filters from resource_dir and extract feature for audio
  71. audio, sample_rate = self.audio_reader.read(batch_data[0])
  72. audio = paddle.to_tensor(audio)
  73. audio = audio[:, 0]
  74. audio = log_mel_spectrogram(audio, resource_path=self.model_dir)
  75. # adapt temperature
  76. temperature_increment_on_fallback = self.config[
  77. "temperature_increment_on_fallback"
  78. ]
  79. if (
  80. temperature_increment_on_fallback is not None
  81. and temperature_increment_on_fallback != "None"
  82. ):
  83. temperature = tuple(
  84. np.arange(
  85. self.config["temperature"],
  86. 1.0 + 1e-6,
  87. temperature_increment_on_fallback,
  88. )
  89. )
  90. else:
  91. temperature = [self.config["temperature"]]
  92. # model inference
  93. result = self.model.transcribe(
  94. audio,
  95. verbose=self.config["verbose"],
  96. task=self.config["task"],
  97. language=self.config["language"],
  98. resource_path=self.model_dir,
  99. temperature=temperature,
  100. compression_ratio_threshold=self.config["compression_ratio_threshold"],
  101. logprob_threshold=self.config["logprob_threshold"],
  102. best_of=self.config["best_of"],
  103. beam_size=self.config["beam_size"],
  104. patience=self.config["patience"],
  105. length_penalty=self.config["length_penalty"],
  106. initial_prompt=self.config["initial_prompt"],
  107. condition_on_previous_text=self.config["condition_on_previous_text"],
  108. no_speech_threshold=self.config["no_speech_threshold"],
  109. )
  110. return {
  111. "input_path": batch_data,
  112. "result": [result],
  113. }