defmodule Bumblebee.Audio do @moduledoc """ High-level tasks related to audio processing. """ @typedoc """ A term representing audio. Can be either of: * a 1-dimensional `Nx.Tensor` with audio samples * an enumerable of 1-dimensional `Nx.Tensor`s, represending a continuous stream of input. The tensors are not required to have the same number of samples, the serving takes care of accumulating and chunking the input as needed. This input is only supported when chunking is enabled with the `:chunk_num_seconds` option * `{:file, path}` with path to an audio file (note that this requires `ffmpeg` installed) """ @type audio :: Nx.t() | Enumerable.t(Nx.t()) | {:file, String.t()} @type speech_to_text_whisper_input :: audio() | %{:audio => audio(), optional(:seed) => integer() | nil} @type speech_to_text_whisper_output :: %{ chunks: list(speech_to_text_whisper_chunk()) } @type speech_to_text_whisper_chunk :: %{ text: String.t(), start_timestamp_seconds: number() | nil, end_timestamp_seconds: number() | nil } @doc """ Builds serving for speech-to-text generation with Whisper models. The serving accepts `t:speech_to_text_whisper_input/0` and returns `t:speech_to_text_whisper_output/0`. This serving always accepts a single input. A list of tensors is interpreted as continuous chunks. To transcribe multiple inputs concurrently use `Nx.Serving.batched_run/2`. ## Options * `:chunk_num_seconds` - enables long-form transcription by splitting the input into chunks of the given length. Models generally have a limit on the input length, so by chunking we can feed smaller bits into the model, then merge the individual outputs into a single result at the end. By default chunking is disabled * `:context_num_seconds` - specifies the amount of overlap between chunks on both sides of split points. The context is effectively discarded when merging the chunks at the end, but it improves the results at the chunk edges. Note that the context is included in the total `:chunk_num_seconds`. Defaults to 1/6 of `:chunk_num_seconds` * `:client_batch_size` - the number of input chunks the client should gather and run at once. When streaming an input that is already available (such as a file), you can use any value without introducing a delay. A good default is to use the same value as `opts[:compile][:batch_size]`. When streaming an input that is being produced live, you may want to process chunks as soon as they are available, in which case set this option to `1`. Defaults to `opts[:compile][:batch_size]` if present, otherwise to `1` * `:language` - the language of the speech, when known upfront. Should be given as ISO alpha-2 code as string. By default no language is assumed and it is inferred from the input * `:task` - either of: * `:transcribe` (default) - generate audio transcription in the same language as the speech * `:translate` - generate translation of the given speech in English * `:timestamps` - when set, the model predicts timestamps and each annotated segment becomes an output chunk. Currently the only supported value is `:segments`, the length of each segment is up to the model * `:compile` - compiles all computations for predefined input shapes during serving initialization. Should be a keyword list with the following keys: * `:batch_size` - the maximum batch size of the input. Inputs are optionally padded to always match this batch size It is advised to set this option in production and also configure a defn compiler using `:defn_options` to maximally reduce inference time. * `:defn_options` - the options for JIT compilation. Defaults to `[]` * `:preallocate_params` - when `true`, explicitly allocates params on the device configured by `:defn_options`. You may want to set this option when using partitioned serving, to allocate params on each of the devices. When using this option, you should first load the parameters into the host. This can be done by passing `backend: {EXLA.Backend, client: :host}` to `load_model/1` and friends. Defaults to `false` * `:stream` - when `true`, the serving immediately returns a stream that emits chunks as they are generated. Note that when using streaming, only a single input can be given to the serving. To process a batch, call the serving with each input separately. Defaults to `false` ## Examples {:ok, whisper} = Bumblebee.load_model({:hf, "openai/whisper-tiny"}) {:ok, featurizer} = Bumblebee.load_featurizer({:hf, "openai/whisper-tiny"}) {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "openai/whisper-tiny"}) {:ok, generation_config} = Bumblebee.load_generation_config({:hf, "openai/whisper-tiny"}) serving = Bumblebee.Audio.speech_to_text_whisper(whisper, featurizer, tokenizer, generation_config, defn_options: [compiler: EXLA] ) output = Nx.Serving.run(serving, {:file, "/path/to/audio.wav"}) #=> %{ #=> chunks: [ #=> %{ #=> text: " There is a cat outside the window.", #=> start_timestamp_seconds: nil, #=> end_timestamp_seconds: nil #=> } #=> ] #=> } text = output.chunks |> Enum.map_join(& &1.text) |> String.trim() #=> "There is a cat outside the window." And with timestamps: serving = Bumblebee.Audio.speech_to_text_whisper(whisper, featurizer, tokenizer, generation_config, defn_options: [compiler: EXLA], chunk_num_seconds: 30, timestamps: :segments ) Nx.Serving.run(serving, {:file, "/path/to/colouredstars_08_mathers_128kb.mp3"}) #=> %{ #=> chunks: [ #=> %{ #=> text: " Such an eight of colored stars, versions of fifty isiatic love poems by Edward Powis-Mathers.", #=> start_timestamp_seconds: 0.0, #=> end_timestamp_seconds: 7.0 #=> }, #=> %{ #=> text: " This the revocs recording is in the public domain. Doubt. From the Japanese of Hori-Kawa,", #=> start_timestamp_seconds: 7.0, #=> end_timestamp_seconds: 14.0 #=> }, #=> %{ #=> text: " will he be true to me that I do not know. But since the dawn, I have had as much disorder in my thoughts as in my black hair, and of doubt.", #=> start_timestamp_seconds: 14.0, #=> end_timestamp_seconds: 27.0 #=> } #=> ] #=> } """ @spec speech_to_text_whisper( Bumblebee.model_info(), Bumblebee.Featurizer.t(), Bumblebee.Tokenizer.t(), Bumblebee.Text.GenerationConfig.t(), keyword() ) :: Nx.Serving.t() defdelegate speech_to_text_whisper( model_info, featurizer, tokenizer, generation_config, opts \\ [] ), to: Bumblebee.Audio.SpeechToTextWhisper end