wx44wx/pipecat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add language support to WhisperSTTService

Browse Source
This commit is contained in:
Mark Backman
2025-02-09 10:51:23 -05:00
parent 081abcedb3
commit c9d8c572c7
2 changed files with 261 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
CHANGELOG.md
View File

@@ -20,6 +20,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
### Changed
- Enhanced `WhisperSTTService` with full language support and improved model
documentation.
- Updated foundation example `14f-function-calling-groq.py` to use
`GroqSTTService` for transcription.

265
src/pipecat/services/whisper.py
View File

@@ -15,6 +15,7 @@ from loguru import logger
from pipecat.frames.frames import ErrorFrame, Frame, TranscriptionFrame
from pipecat.services.ai_services import SegmentedSTTService
from pipecat.transcriptions.language import Language
from pipecat.utils.time import time_now_iso8601
try:
@@ -26,18 +27,219 @@ except ModuleNotFoundError as e:
class Model(Enum):
"""Class of basic Whisper model selection options"""
"""Class of basic Whisper model selection options.
Available models:
Multilingual models:
TINY: Smallest multilingual model
BASE: Basic multilingual model
MEDIUM: Good balance for multilingual
LARGE: Best quality multilingual
DISTIL_LARGE_V2: Fast multilingual
English-only models:
DISTIL_MEDIUM_EN: Fast English-only
"""
# Multilingual models
TINY = "tiny"
BASE = "base"
MEDIUM = "medium"
LARGE = "large-v3"
DISTIL_LARGE_V2 = "Systran/faster-distil-whisper-large-v2"
# English-only models
DISTIL_MEDIUM_EN = "Systran/faster-distil-whisper-medium.en"
def language_to_whisper_language(language: Language) -> Optional[str]:
"""Maps pipecat Language enum to Whisper language codes.
Args:
language: A Language enum value representing the input language.
Returns:
str or None: The corresponding Whisper language code, or None if not supported.
Note:
Only includes languages officially supported by Whisper.
"""
language_map = {
# Arabic
Language.AR: "ar",
Language.AR_AE: "ar",
Language.AR_BH: "ar",
Language.AR_DZ: "ar",
Language.AR_EG: "ar",
Language.AR_IQ: "ar",
Language.AR_JO: "ar",
Language.AR_KW: "ar",
Language.AR_LB: "ar",
Language.AR_LY: "ar",
Language.AR_MA: "ar",
Language.AR_OM: "ar",
Language.AR_QA: "ar",
Language.AR_SA: "ar",
Language.AR_SY: "ar",
Language.AR_TN: "ar",
Language.AR_YE: "ar",
# Bengali
Language.BN: "bn",
Language.BN_BD: "bn",
Language.BN_IN: "bn",
# Czech
Language.CS: "cs",
Language.CS_CZ: "cs",
# Danish
Language.DA: "da",
Language.DA_DK: "da",
# German
Language.DE: "de",
Language.DE_AT: "de",
Language.DE_CH: "de",
Language.DE_DE: "de",
# Greek
Language.EL: "el",
Language.EL_GR: "el",
# English
Language.EN: "en",
Language.EN_AU: "en",
Language.EN_CA: "en",
Language.EN_GB: "en",
Language.EN_HK: "en",
Language.EN_IE: "en",
Language.EN_IN: "en",
Language.EN_KE: "en",
Language.EN_NG: "en",
Language.EN_NZ: "en",
Language.EN_PH: "en",
Language.EN_SG: "en",
Language.EN_TZ: "en",
Language.EN_US: "en",
Language.EN_ZA: "en",
# Spanish
Language.ES: "es",
Language.ES_AR: "es",
Language.ES_BO: "es",
Language.ES_CL: "es",
Language.ES_CO: "es",
Language.ES_CR: "es",
Language.ES_CU: "es",
Language.ES_DO: "es",
Language.ES_EC: "es",
Language.ES_ES: "es",
Language.ES_GQ: "es",
Language.ES_GT: "es",
Language.ES_HN: "es",
Language.ES_MX: "es",
Language.ES_NI: "es",
Language.ES_PA: "es",
Language.ES_PE: "es",
Language.ES_PR: "es",
Language.ES_PY: "es",
Language.ES_SV: "es",
Language.ES_US: "es",
Language.ES_UY: "es",
Language.ES_VE: "es",
# Persian
Language.FA: "fa",
Language.FA_IR: "fa",
# Finnish
Language.FI: "fi",
Language.FI_FI: "fi",
# French
Language.FR: "fr",
Language.FR_BE: "fr",
Language.FR_CA: "fr",
Language.FR_CH: "fr",
Language.FR_FR: "fr",
# Hindi
Language.HI: "hi",
Language.HI_IN: "hi",
# Hungarian
Language.HU: "hu",
Language.HU_HU: "hu",
# Indonesian
Language.ID: "id",
Language.ID_ID: "id",
# Italian
Language.IT: "it",
Language.IT_IT: "it",
# Japanese
Language.JA: "ja",
Language.JA_JP: "ja",
# Korean
Language.KO: "ko",
Language.KO_KR: "ko",
# Dutch
Language.NL: "nl",
Language.NL_BE: "nl",
Language.NL_NL: "nl",
# Polish
Language.PL: "pl",
Language.PL_PL: "pl",
# Portuguese
Language.PT: "pt",
Language.PT_BR: "pt",
Language.PT_PT: "pt",
# Romanian
Language.RO: "ro",
Language.RO_RO: "ro",
# Russian
Language.RU: "ru",
Language.RU_RU: "ru",
# Slovak
Language.SK: "sk",
Language.SK_SK: "sk",
# Swedish
Language.SV: "sv",
Language.SV_SE: "sv",
# Thai
Language.TH: "th",
Language.TH_TH: "th",
# Turkish
Language.TR: "tr",
Language.TR_TR: "tr",
# Ukrainian
Language.UK: "uk",
Language.UK_UA: "uk",
# Urdu
Language.UR: "ur",
Language.UR_IN: "ur",
Language.UR_PK: "ur",
# Vietnamese
Language.VI: "vi",
Language.VI_VN: "vi",
# Chinese
Language.ZH: "zh",
Language.ZH_CN: "zh",
Language.ZH_HK: "zh",
Language.ZH_TW: "zh",
}
return language_map.get(language)
class WhisperSTTService(SegmentedSTTService):
"""Class to transcribe audio with a locally-downloaded Whisper model"""
"""Class to transcribe audio with a locally-downloaded Whisper model.
This service uses Faster Whisper to perform speech-to-text transcription on audio
segments. It supports multiple languages and various model sizes.
Args:
model: The Whisper model to use for transcription. Can be a Model enum or string.
device: The device to run inference on ('cpu', 'cuda', or 'auto').
compute_type: The compute type for inference ('default', 'int8', 'int8_float16', etc.).
no_speech_prob: Probability threshold for filtering out non-speech segments.
language: The default language for transcription.
**kwargs: Additional arguments passed to SegmentedSTTService.
Attributes:
_device: The device used for inference.
_compute_type: The compute type for inference.
_no_speech_prob: Threshold for non-speech filtering.
_model: The loaded Whisper model instance.
_settings: Dictionary containing service settings.
"""
def __init__(
self,
@@ -46,6 +248,7 @@ class WhisperSTTService(SegmentedSTTService):
device: str = "auto",
compute_type: str = "default",
no_speech_prob: float = 0.4,
language: Language = Language.EN,
**kwargs,
):
super().__init__(**kwargs)
@@ -54,14 +257,47 @@ class WhisperSTTService(SegmentedSTTService):
self.set_model_name(model if isinstance(model, str) else model.value)
self._no_speech_prob = no_speech_prob
self._model: Optional[WhisperModel] = None
self._settings = {
"language": language,
}
self._load()
def can_generate_metrics(self) -> bool:
"""Indicates whether this service can generate metrics.
Returns:
bool: True, as this service supports metric generation.
"""
return True
def language_to_service_language(self, language: Language) -> Optional[str]:
"""Convert from pipecat Language to Whisper language code.
Args:
language: The Language enum value to convert.
Returns:
str or None: The corresponding Whisper language code, or None if not supported.
"""
return language_to_whisper_language(language)
async def set_language(self, language: Language):
"""Set the language for transcription.
Args:
language: The Language enum value to use for transcription.
"""
logger.info(f"Switching STT language to: [{language}]")
self._settings["language"] = language
def _load(self):
"""Loads the Whisper model. Note that if this is the first time
this model is being run, it will take time to download.
"""Loads the Whisper model.
Note:
If this is the first time this model is being run,
it will take time to download from the Hugging Face model hub.
"""
logger.debug("Loading Whisper model...")
self._model = WhisperModel(
@@ -70,7 +306,19 @@ class WhisperSTTService(SegmentedSTTService):
logger.debug("Loaded Whisper model")
async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
"""Transcribes given audio using Whisper"""
"""Transcribes given audio using Whisper.
Args:
audio: Raw audio bytes in 16-bit PCM format.
Yields:
Frame: Either a TranscriptionFrame containing the transcribed text
or an ErrorFrame if transcription fails.
Note:
The audio is expected to be 16-bit signed PCM data.
The service will normalize it to float32 in the range [-1, 1].
"""
if not self._model:
logger.error(f"{self} error: Whisper model not available")
yield ErrorFrame("Whisper model not available")
@@ -82,7 +330,10 @@ class WhisperSTTService(SegmentedSTTService):
# Divide by 32768 because we have signed 16-bit data.
audio_float = np.frombuffer(audio, dtype=np.int16).astype(np.float32) / 32768.0
segments, _ = await asyncio.to_thread(self._model.transcribe, audio_float)
whisper_lang = self.language_to_service_language(self._settings["language"])
segments, _ = await asyncio.to_thread(
self._model.transcribe, audio_float, language=whisper_lang
)
text: str = ""
for segment in segments:
if segment.no_speech_prob < self._no_speech_prob:
@@ -93,4 +344,4 @@ class WhisperSTTService(SegmentedSTTService):
if text:
logger.debug(f"Transcription: [{text}]")
yield TranscriptionFrame(text, "", time_now_iso8601())
yield TranscriptionFrame(text, "", time_now_iso8601(), self._settings["language"])