sparknlp.annotator.audio.whisper_for_ctc#

Classes#

class WhisperForCTC(classname='com.johnsnowlabs.nlp.annotators.audio.WhisperForCTC', java_model=None)[source]#

Whisper Model with a language modeling head on top for Connectionist Temporal Classification (CTC).

Whisper is an automatic speech recognition (ASR) system trained on 680,000 hours of multilingual and multitask supervised data collected from the web. It transcribe in multiple languages, as well as translate from those languages into English.

The audio needs to be provided pre-processed an array of floats.

Note that at the moment, this annotator only supports greedy search and only Spark Versions 3.4 and up are supported.

For multilingual models, the language and the task (transcribe or translate) can be set with setLanguage and setTask.

Pretrained models can be loaded with pretrained of the companion object:

speechToText = WhisperForCTC.pretrained() \
    .setInputCols(["audio_assembler"]) \
    .setOutputCol("text")

The default model is "asr_whisper_tiny_opt", if no name is provided.

For available pretrained models please see the Models Hub.

To see which models are compatible and how to import them see JohnSnowLabs/spark-nlp#5669 and to see more extended examples, see WhisperForCTCTestSpec.

References:

Robust Speech Recognition via Large-Scale Weak Supervision

Paper Abstract:

We study the capabilities of speech processing systems trained simply to predict large amounts of transcripts of audio on the internet. When scaled to 680,000 hours of multilingual and multitask supervision, the resulting models generalize well to standard benchmarks and are often competitive with prior fully supervised results but in a zero- shot transfer setting without the need for any fine- tuning. When compared to humans, the models approach their accuracy and robustness. We are releasing models and inference code to serve as a foundation for further work on robust speech processing.

Input Annotation types	Output Annotation type
AUDIO	DOCUMENT

Parameters:

task

The formatted task for the audio. Either <|translate|> or <|transcribe|>.

language

The language for the audio, formatted to e.g. <|en|>. Check the model description for supported languages.

isMultilingual

Whether the model is multilingual

minOutputLength

Minimum length of the sequence to be generated

maxOutputLength

Maximum length of output text

doSample

Whether or not to use sampling; use greedy decoding otherwise

temperature

The value used to module the next token probabilities

topK

The number of highest probability vocabulary tokens to keep for top-k-filtering

topP

If set to float < 1, only the most probable tokens with probabilities that add up to top_p or higher are kept for generation

repetitionPenalty

The parameter for repetition penalty. 1.0 means no penalty. See this paper for more details

noRepeatNgramSize

If set to int > 0, all ngrams of that size can only occur once

beamSize

The Number of beams for beam search

Examples

>>> import sparknlp
>>> from sparknlp.base import *
>>> from sparknlp.annotator import *
>>> from pyspark.ml import Pipeline
>>> audioAssembler = AudioAssembler() \
...     .setInputCol("audio_content") \
...     .setOutputCol("audio_assembler")
>>> speechToText = WhisperForCTC.pretrained() \
...     .setInputCols(["audio_assembler"]) \
...     .setOutputCol("text")
>>> pipeline = Pipeline().setStages([audioAssembler, speechToText])
>>> processedAudioFloats = spark.createDataFrame([[rawFloats]]).toDF("audio_content")
>>> result = pipeline.fit(processedAudioFloats).transform(processedAudioFloats)
>>> result.select("text.result").show(truncate = False)
+------------------------------------------------------------------------------------------+
|result                                                                                    |
+------------------------------------------------------------------------------------------+
|[ Mr. Quilter is the apostle of the middle classes and we are glad to welcome his gospel.]|
+------------------------------------------------------------------------------------------+

setConfigProtoBytes(b)[source]#

Sets configProto from tensorflow, serialized into byte array.

Parameters:

bList[int]

ConfigProto from tensorflow, serialized into byte array

getLanguage()[source]#

Gets the langauge for the transcription.

getIsMultilingual()[source]#

Gets whether the model is multilingual.

setLanguage(value)[source]#

Sets the language for the audio, formatted to e.g. <|en|>. Check the model description for supported languages.

Parameters:

valueString

Formatted language code

setTask(value)[source]#

Sets the formatted task for the audio. Either <|translate|> or <|transcribe|>.

Only multilingual models can do translation.

Parameters:

valueString

Formatted task

static loadSavedModel(folder, spark_session)[source]#

Loads a locally saved model.

Parameters:

folderstr

Folder of the saved model

spark_sessionpyspark.sql.SparkSession

The current SparkSession

Returns:

WhisperForCTC

The restored model

static pretrained(name='asr_whisper_tiny_opt', lang='xx', remote_loc=None)[source]#

Downloads and loads a pretrained model.

Parameters:

namestr, optional

Name of the pretrained model, by default “asr_hubert_large_ls960”

langstr, optional

Language of the pretrained model, by default “en”

remote_locstr, optional

Optional remote address of the resource, by default None. Will use Spark NLPs repositories otherwise.

Returns:

WhisperForCTC

The restored model