################################################################################
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
################################################################################
"""Speech model operators for decoded audio waveforms.
These operators are model-backed pandas UDFs. They accept standardized
``AUDIO_WAVEFORM`` values only: callers should decode boundary inputs and
normalize to 16 kHz mono before invoking Whisper-based operators.
"""
from typing import TYPE_CHECKING
from pyflink.common import Row
from pyflink.dataframe import DataType, udf
from pyflink.model.cache_manager import prepare_and_load_model_handle
from pyflink.multimodal.codec.audio import (
ensure_supported_audio_input,
waveform_to_ndarray,
)
from pyflink.multimodal.types.audio import AUDIO_WAVEFORM_FIELDS
from pyflink.multimodal.utils import (
_build_or_apply_udf,
_udf_runtime_kwargs,
scatter_results,
)
from pyflink.table.udf import ScalarFunction
if TYPE_CHECKING:
import pandas as pd
__all__ = [
"audio_asr_whisper",
"audio_detect_language",
]
_MODEL_DEPENDENCIES = ("transformers", "torch", "numpy")
_WHISPER_SAMPLE_RATE = 16000
_WHISPER_CHANNELS = 1
_DEFAULT_WHISPER_MODEL = "openai/whisper-tiny"
_TIMESTAMP_TYPE = DataType.struct([
("start_ms", DataType.int64()),
("end_ms", DataType.int64()),
])
_ASR_SEGMENT_TYPE = DataType.struct([
("text", DataType.string()),
("start_ms", DataType.int64()),
("end_ms", DataType.int64()),
])
_ASR_RESULT_TYPE = DataType.struct([
("asr_result", DataType.string()),
("timestamps", DataType.list(_TIMESTAMP_TYPE)),
("segments", DataType.list(_ASR_SEGMENT_TYPE)),
])
_DETECTED_LANGUAGE_TYPE = DataType.struct([
("language_tag", DataType.string()),
("language_name", DataType.string()),
("confidence", DataType.float64()),
])
_DETECTED_LANGUAGE_LIST_TYPE = DataType.list(_DETECTED_LANGUAGE_TYPE)
_AUDIO_WAVEFORM_INT_FIELDS = frozenset(
("sample_rate", "channels", "frames", "duration_ms",
"start_time_ms", "end_time_ms")
)
def _hf_whisper_speech_adapter_cls():
from pyflink.model.backends.huggingface import HfWhisperSpeechAdapter
return HfWhisperSpeechAdapter
def _validate_whisper_model(model, operator_name):
if not isinstance(model, str) or not model:
raise ValueError(f"{operator_name}.model must be a string, got {model!r}")
return model
def _valid_audio_payloads(audio_series, operator_name):
payloads = []
valid_indices = []
for index, audio_input in enumerate(_iter_audio_inputs(audio_series, operator_name)):
if audio_input is None:
continue
ensure_supported_audio_input(
operator_name,
("AUDIO_WAVEFORM",),
audio_input,
)
payloads.append(_waveform_payload(audio_input, operator_name))
valid_indices.append(index)
return payloads, valid_indices
def _iter_audio_inputs(audio_batch, operator_name):
"""Yield AUDIO_WAVEFORM rows from a pandas UDF input batch.
PyFlink flattens a ROW argument for pandas UDFs into a pandas.DataFrame:
each ROW field is one column, and each DataFrame row is one input row.
Reconstruct the audio waveform row before applying the normal codec layer
validation.
"""
import pandas as pd
if isinstance(audio_batch, pd.DataFrame):
field_positions = _flattened_audio_waveform_field_positions(
audio_batch,
operator_name,
)
for values in audio_batch.itertuples(index=False, name=None):
row_values = {
field: _audio_waveform_field_value(field, values[position], pd)
for field, position in field_positions
}
if all(value is None for value in row_values.values()):
yield None
continue
yield Row(**row_values)
return
yield from audio_batch
def _flattened_audio_waveform_field_positions(audio_frame, operator_name):
columns = list(audio_frame.columns)
if all(field in columns for field in AUDIO_WAVEFORM_FIELDS):
return [
(field, columns.index(field))
for field in AUDIO_WAVEFORM_FIELDS
]
if len(columns) == len(AUDIO_WAVEFORM_FIELDS):
return [
(field, index)
for index, field in enumerate(AUDIO_WAVEFORM_FIELDS)
]
raise ValueError(
f"{operator_name} expects flattened AUDIO_WAVEFORM DataFrame columns "
f"{list(AUDIO_WAVEFORM_FIELDS)}, got {columns}"
)
def _audio_waveform_field_value(field, value, pd):
if value is None:
return None
if field != "data":
try:
if pd.isna(value):
return None
except (TypeError, ValueError):
pass
if field in _AUDIO_WAVEFORM_INT_FIELDS and value is not None:
return int(value)
return value
def _audio_batch_length(audio_batch):
import pandas as pd
if isinstance(audio_batch, pd.DataFrame):
return len(audio_batch.index)
return len(audio_batch)
def _waveform_payload(waveform, operator_name):
sample_rate = int(waveform.sample_rate)
channels = int(waveform.channels)
if sample_rate != _WHISPER_SAMPLE_RATE or channels != _WHISPER_CHANNELS:
raise ValueError(
f"{operator_name} requires 16 kHz mono AUDIO_WAVEFORM; call "
"audio_standardize(sample_rate=16000, channels=1) first, got "
f"sample_rate={sample_rate}, channels={channels}"
)
samples = waveform_to_ndarray(waveform)
return {
"samples": samples[:, 0],
"sample_rate": sample_rate,
}
def _scatter_series(results, valid_indices, total_len):
import pandas as pd
return pd.Series(scatter_results(results, valid_indices, total_len))
class _AudioAsrWhisper(ScalarFunction):
"""Run Whisper ASR over AUDIO_WAVEFORM batches."""
def __init__(self, language=None, task="transcribe",
model="openai/whisper-tiny", model_sharing=None, num_gpus=None,
gpu_type=None):
super().__init__()
if language is not None and not isinstance(language, str):
raise ValueError(f"language must be a string or None, got {language!r}")
if task not in ("transcribe", "translate"):
raise ValueError("task must be 'transcribe' or 'translate'")
self.model = _validate_whisper_model(model, "audio_asr_whisper")
self.language = language
self.task = task
self.model_sharing = model_sharing
self._num_gpus = num_gpus
self._gpu_type = gpu_type
@staticmethod
def _predict_batch(model, payloads, language, task):
return model.transcribe(payloads, language=language, task=task)
def open(self, function_context):
self._model_handle = prepare_and_load_model_handle(
adapter_cls=_hf_whisper_speech_adapter_cls(),
config={},
function_context=function_context,
model_sharing=self.model_sharing,
dependencies=_MODEL_DEPENDENCIES,
model_id=self.model,
requested_num_gpus=self._num_gpus,
requested_gpu_type=self._gpu_type,
)
self._model_handle.register_operation(
"audio_asr_whisper", _AudioAsrWhisper._predict_batch
)
def close(self):
if getattr(self, "_model_handle", None) is not None:
self._model_handle.release()
self._model_handle = None
def eval(self, audio_series: "pd.Series") -> "pd.Series":
total_len = _audio_batch_length(audio_series)
payloads, valid_indices = _valid_audio_payloads(
audio_series,
"audio_asr_whisper",
)
if not payloads:
return _scatter_series([], [], total_len)
results = self._model_handle.call(
"audio_asr_whisper",
payloads,
self.language,
self.task,
)
return _scatter_series(
[_normalize_asr_result(result) for result in results],
valid_indices,
total_len,
)
[docs]def audio_asr_whisper(
*columns,
language=None,
task="transcribe",
model="openai/whisper-tiny",
model_sharing=None,
concurrency=None,
batch_size=None,
num_gpus=None,
gpu_type=None,
):
"""
Run Whisper automatic speech recognition over audio waveforms.
The input contract is intentionally narrow: each row must be a waveform row
with 16 kHz sample rate and one channel. Use ``audio_decode`` and
``audio_standardize`` before calling this operator. Rows that do not match
the Whisper contract fail early instead of being implicitly resampled inside
the model operator.
The output is a row with these fields:
* ``asr_result``: full transcription text.
* ``timestamps``: ``ARRAY<ROW<start_ms BIGINT, end_ms BIGINT>>``.
* ``segments``: ``ARRAY<ROW<text STRING, start_ms BIGINT, end_ms BIGINT>>``.
The operator uses HuggingFace Whisper-compatible checkpoints through the
PyFlink model runtime. The selected model is resolved by the PyFlink model
store. Depending on the model store policy, the model may be loaded from a
local path/cache or resolved by model id.
Args:
*columns: Optional audio columns. Supported input is a 16 kHz mono
waveform row.
language: Optional source language hint passed to Whisper. ``None``
lets the model infer the language.
task: ``"transcribe"`` or ``"translate"``.
model: HuggingFace Whisper or Whisper-compatible model id. The default
is ``"openai/whisper-tiny"``. Other checkpoints can be used when
they are available in the Python worker environment or model cache.
model_sharing: Model handle sharing mode understood by the PyFlink
model runtime.
concurrency: UDF concurrency. ``None`` uses the framework default.
batch_size: Pandas UDF batch size.
num_gpus: GPU resource amount requested by the UDF runtime.
gpu_type: GPU type requested for model inference. The DataFrame UDF
runtime requires this when ``num_gpus`` is set.
Returns:
A pandas UDF producing ``ROW<asr_result STRING, timestamps
ARRAY<ROW<start_ms BIGINT, end_ms BIGINT>>, segments
ARRAY<ROW<text STRING, start_ms BIGINT, end_ms BIGINT>>>``. In SQL,
downstream code can access fields with dot notation, for example
``asr.asr_result`` or ``asr.segments`` after aliasing.
Raises:
ValueError: If a row is not a 16 kHz mono waveform row or if ``task``
is unsupported.
Notes:
In SQL, positional configuration arguments follow the same order:
``audio_asr_whisper(audio, language, task, model)``. ``model`` is last
because most jobs keep the checkpoint fixed and tune language/task more
often.
Examples::
>>> decode = audio_decode()
>>> standardize = audio_standardize(sample_rate=16000, channels=1)
>>> speech_ready = standardize(decode(col("audio_bytes")))
>>> # Usage 1: create a reusable UDF and apply it to a column.
>>> asr = audio_asr_whisper()
>>> df = df.with_column(
... "asr",
... asr(speech_ready),
... )
>>>
>>> # Usage 2: pass the column directly when building the expression.
>>> df = df.with_column(
... "asr",
... audio_asr_whisper(speech_ready),
... )
"""
wrapper = udf(
_AudioAsrWhisper(
language=language,
task=task,
model=model,
model_sharing=model_sharing,
num_gpus=num_gpus,
gpu_type=gpu_type,
),
func_type="pandas",
return_dtype=_ASR_RESULT_TYPE,
**_udf_runtime_kwargs(
concurrency=concurrency,
batch_size=batch_size,
num_gpus=num_gpus,
gpu_type=gpu_type,
),
)
return _build_or_apply_udf(wrapper, *columns)
class _AudioDetectLanguage(ScalarFunction):
"""Detect spoken language over AUDIO_WAVEFORM batches."""
def __init__(self, top_k=1, model=_DEFAULT_WHISPER_MODEL,
model_sharing=None, num_gpus=None, gpu_type=None):
super().__init__()
if not isinstance(top_k, int) or isinstance(top_k, bool) or top_k <= 0:
raise ValueError(f"top_k must be a positive integer, got {top_k!r}")
self.model = _validate_whisper_model(model, "audio_detect_language")
self.model_sharing = model_sharing
self.top_k = top_k
self._num_gpus = num_gpus
self._gpu_type = gpu_type
@staticmethod
def _predict_batch(model, payloads, top_k):
return model.detect_language(payloads, top_k=top_k)
def open(self, function_context):
self._model_handle = prepare_and_load_model_handle(
adapter_cls=_hf_whisper_speech_adapter_cls(),
config={},
function_context=function_context,
model_sharing=self.model_sharing,
dependencies=_MODEL_DEPENDENCIES,
model_id=self.model,
requested_num_gpus=self._num_gpus,
requested_gpu_type=self._gpu_type,
)
self._model_handle.register_operation(
"audio_detect_language", _AudioDetectLanguage._predict_batch
)
def close(self):
if getattr(self, "_model_handle", None) is not None:
self._model_handle.release()
self._model_handle = None
def eval(self, audio_series: "pd.Series") -> "pd.Series":
total_len = _audio_batch_length(audio_series)
payloads, valid_indices = _valid_audio_payloads(
audio_series,
"audio_detect_language",
)
if not payloads:
return _scatter_series([], [], total_len)
results = self._model_handle.call(
"audio_detect_language",
payloads,
self.top_k,
)
return _scatter_series(
[_normalize_detected_languages(result) for result in results],
valid_indices,
total_len,
)
[docs]def audio_detect_language(
*columns,
top_k=1,
model="openai/whisper-tiny",
model_sharing=None,
concurrency=None,
batch_size=None,
num_gpus=None,
gpu_type=None,
):
"""
Detect likely spoken languages in an audio waveform.
The operator uses the selected Whisper-compatible checkpoint to rank likely
languages for the whole audio row. It returns the top ``top_k`` language
candidates; it does not split code-switched audio into time ranges.
The input contract is the same as ``audio_asr_whisper``: a 16 kHz mono
waveform row.
Each result item contains these fields:
* ``language_tag``: Model language tag, such as ``en`` or ``zh``. Some
checkpoints may return longer tags such as ``yue``. The value is named
``language_tag`` because it comes from the model vocabulary and is not
guaranteed to be a normalized BCP 47 language tag.
* ``language_name``: English display name for the language.
* ``confidence``: Model confidence for the candidate language.
Args:
*columns: Optional audio columns. Supported input is a 16 kHz mono
waveform row.
top_k: Number of language candidates to return per audio row.
model: HuggingFace Whisper or Whisper-compatible model id. The default
is ``"openai/whisper-tiny"``.
model_sharing: Model handle sharing mode understood by the PyFlink
model runtime.
concurrency: UDF concurrency. ``None`` uses the framework default.
batch_size: Pandas UDF batch size.
num_gpus: GPU resource amount requested by the UDF runtime.
gpu_type: GPU type requested for model inference. The DataFrame UDF
runtime requires this when ``num_gpus`` is set.
Returns:
A pandas UDF producing ``ARRAY<ROW<language_tag STRING,
language_name STRING, confidence DOUBLE>>`` sorted by confidence.
In SQL, unnest or index the array before consuming item fields.
Raises:
ValueError: If a row is not a 16 kHz mono waveform row.
Notes:
In SQL, positional configuration arguments follow the same order:
``audio_detect_language(audio, top_k, model)``.
Examples::
>>> decode = audio_decode()
>>> standardize = audio_standardize(sample_rate=16000, channels=1)
>>> speech_ready = standardize(decode(col("audio_bytes")))
>>> # Usage 1: create a reusable UDF and apply it to a column.
>>> detect_language = audio_detect_language(top_k=3)
>>> df = df.with_column(
... "language",
... detect_language(speech_ready),
... )
>>>
>>> # Usage 2: pass the column directly when building the expression.
>>> df = df.with_column(
... "language",
... audio_detect_language(speech_ready, top_k=3),
... )
"""
wrapper = udf(
_AudioDetectLanguage(
top_k=top_k,
model=model,
model_sharing=model_sharing,
num_gpus=num_gpus,
gpu_type=gpu_type,
),
func_type="pandas",
return_dtype=_DETECTED_LANGUAGE_LIST_TYPE,
**_udf_runtime_kwargs(
concurrency=concurrency,
batch_size=batch_size,
num_gpus=num_gpus,
gpu_type=gpu_type,
),
)
return _build_or_apply_udf(wrapper, *columns)
def _normalize_asr_result(result):
if result is None:
return None
asr_result = (
_optional_str(_get_value(result, "asr_result"))
or _optional_str(_get_value(result, "text"))
or ""
)
raw_segments = (
_get_value(result, "segments")
or _get_value(result, "chunks")
or []
)
segments = [_normalize_asr_segment(segment) for segment in raw_segments]
timestamps = [
{"start_ms": segment["start_ms"], "end_ms": segment["end_ms"]}
for segment in segments
]
return {
"asr_result": asr_result,
"timestamps": timestamps,
"segments": segments,
}
def _normalize_asr_segment(segment):
text = _optional_str(_get_value(segment, "text")) or ""
start_ms, end_ms = _segment_bounds_ms(segment)
return {"text": text, "start_ms": start_ms, "end_ms": end_ms}
def _segment_bounds_ms(segment):
start_ms = _get_value(segment, "start_ms")
end_ms = _get_value(segment, "end_ms")
if start_ms is not None or end_ms is not None:
return _time_value_ms(start_ms), _time_value_ms(end_ms)
timestamp = _get_value(segment, "timestamp")
if timestamp is not None:
return _time_value_ms(timestamp[0]), _time_value_ms(timestamp[1])
return (
_seconds_to_ms(_get_value(segment, "start")),
_seconds_to_ms(_get_value(segment, "end")),
)
def _normalize_detected_language(result):
if result is None:
return None
language_tag = (
_optional_str(_get_value(result, "language_tag"))
or _optional_str(_get_value(result, "language_code"))
or _optional_str(_get_value(result, "language"))
or _optional_str(_get_value(result, "label"))
or "unknown"
)
language_name = (
_optional_str(_get_value(result, "language_name"))
or _optional_str(_get_value(result, "name"))
or language_tag
)
confidence = _get_value(result, "confidence")
if confidence is None:
confidence = _get_value(result, "score")
return {
"language_tag": language_tag,
"language_name": language_name,
"confidence": None if confidence is None else float(confidence),
}
def _normalize_detected_languages(result):
if result is None:
return None
if isinstance(result, (list, tuple)):
return [_normalize_detected_language(item) for item in result]
raise ValueError(
"audio_detect_language backend must return a list of language "
f"candidates per input row, got {type(result).__name__}"
)
def _get_value(value, key):
if isinstance(value, dict):
return value.get(key)
return getattr(value, key, None)
def _optional_str(value):
if value is None:
return None
return str(value)
def _time_value_ms(value):
if value is None:
return None
if isinstance(value, float):
return _seconds_to_ms(value)
return int(value)
def _seconds_to_ms(value):
if value is None:
return None
return int(round(float(value) * 1000.0))