diff --git a/apps/speech/app/_layout.tsx b/apps/speech/app/_layout.tsx index bd4a75cc71..cb6e4c3b9b 100644 --- a/apps/speech/app/_layout.tsx +++ b/apps/speech/app/_layout.tsx @@ -27,6 +27,20 @@ export default function Layout() { title: 'Voice Activity Detection', }} /> + + ); } diff --git a/apps/speech/app/audio-file-transcription/index.tsx b/apps/speech/app/audio-file-transcription/index.tsx new file mode 100644 index 0000000000..f6361c7b74 --- /dev/null +++ b/apps/speech/app/audio-file-transcription/index.tsx @@ -0,0 +1,318 @@ +import React, { useEffect, useState } from 'react'; +import { + Platform, + View, + Text, + StyleSheet, + ScrollView, + TouchableOpacity, + TextInput, +} from 'react-native'; +import { useSpeechToText, models, WHISPER_SAMPLE_RATE_HZ } from 'react-native-executorch'; +import { decodeAudioData } from 'react-native-audio-api'; + +import ScreenWrapper from '../../components/ScreenWrapper'; +import { ModelPicker } from '../../components/ModelPicker'; +import { ModelStatus } from '../../components/ModelStatus'; +import { theme } from '../../theme'; + +const MODELS = [ + { + name: 'Tiny English (CPU)', + config: models.speechToText.WHISPER.EN.TINY.XNNPACK_FP32, + }, + { + name: 'Tiny English (CoreML)', + config: models.speechToText.WHISPER.EN.TINY.COREML_FP16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Tiny English (MLX)', + config: models.speechToText.WHISPER.EN.TINY.MLX_BF16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Base English (CPU)', + config: models.speechToText.WHISPER.EN.BASE.XNNPACK_FP32, + }, + { + name: 'Base English (CoreML)', + config: models.speechToText.WHISPER.EN.BASE.COREML_FP16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Base English (MLX)', + config: models.speechToText.WHISPER.EN.BASE.MLX_BF16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Small English (CPU)', + config: models.speechToText.WHISPER.EN.SMALL.XNNPACK_FP32, + }, + { + name: 'Small English (CoreML)', + config: models.speechToText.WHISPER.EN.SMALL.COREML_FP16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Small English (MLX)', + config: models.speechToText.WHISPER.EN.SMALL.MLX_BF16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Tiny Multilingual (CPU)', + config: models.speechToText.WHISPER.TINY.XNNPACK_FP32, + }, + { + name: 'Tiny Multilingual (CoreML)', + config: models.speechToText.WHISPER.TINY.COREML_FP16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Tiny Multilingual (MLX)', + config: models.speechToText.WHISPER.TINY.MLX_BF16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Base Multilingual (CPU)', + config: models.speechToText.WHISPER.BASE.XNNPACK_FP32, + }, + { + name: 'Base Multilingual (CoreML)', + config: models.speechToText.WHISPER.BASE.COREML_FP16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Base Multilingual (MLX)', + config: models.speechToText.WHISPER.BASE.MLX_BF16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Small Multilingual (CPU)', + config: models.speechToText.WHISPER.SMALL.XNNPACK_FP32, + }, + { + name: 'Small Multilingual (CoreML)', + config: models.speechToText.WHISPER.SMALL.COREML_FP16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Small Multilingual (MLX)', + config: models.speechToText.WHISPER.SMALL.MLX_BF16, + disabled: Platform.OS !== 'ios', + }, +]; + +function STTAudioContent() { + const [selectedModel, setSelectedModel] = useState(MODELS[0]!); + const [selectedLanguage, setSelectedLanguage] = useState('en'); + const [status, setStatus] = useState('Idle'); + const [url, setUrl] = useState('https://models.silero.ai/vad_models/en.wav'); + const [isTranscribing, setIsTranscribing] = useState(false); + const [audioText, setAudioText] = useState(''); + const [runError, setRunError] = useState(null); + + const { + isReady: isSttReady, + transcribe, + transcribeStop, + downloadProgress, + error: modelError, + } = useSpeechToText(selectedModel.config); + + const supportedLanguages = selectedModel.config.supportedLanguages; + + useEffect(() => { + return () => { + if (transcribeStop) transcribeStop(); + }; + }, [transcribeStop]); + + const startTranscribing = async () => { + if (!isSttReady || !transcribe || isTranscribing || !url) return; + setRunError(null); + setAudioText(''); + setStatus('Decoding URL...'); + setIsTranscribing(true); + + try { + const audioBuffer = await decodeAudioData(url, WHISPER_SAMPLE_RATE_HZ); + const samples = audioBuffer.getChannelData(0); + + setStatus('Transcribing...'); + + let currentText = ''; + const text = await transcribe(samples, { language: selectedLanguage as any }, (token) => { + currentText += token; + setAudioText(currentText); + }); + setAudioText(text); + setStatus('Done'); + } catch (err) { + console.error('STT Audio transcription error:', err); + const errMsg = err instanceof Error ? err.message : String(err); + setRunError(errMsg); + setStatus('Error'); + } finally { + setIsTranscribing(false); + } + }; + + const stopTranscribing = () => { + if (!isTranscribing) return; + if (transcribeStop) transcribeStop(); + setIsTranscribing(false); + setStatus('Cancelled'); + }; + + const isModelBusy = status.includes('...'); + + return ( + + + ({ + label: m.name, + value: m, + disabled: isModelBusy || !!m.disabled, + }))} + selectedValue={selectedModel} + onValueChange={(m) => { + setSelectedModel(m); + setSelectedLanguage(m.config.supportedLanguages[0]!); + }} + /> + + {supportedLanguages.length > 1 && ( + ({ + label: lang, + value: lang, + disabled: isModelBusy, + }))} + selectedValue={selectedLanguage} + onValueChange={setSelectedLanguage} + /> + )} + + + {runError && ( + + {runError} + + )} + + {/* Audio Link Panel */} + + Transcribe Audio File + + + {!isTranscribing ? ( + + Transcribe Audio File + + ) : ( + + Stop Transcription + + )} + + + Transcription Output: + + {audioText} + + + + ); +} + +export default function STTAudioScreen() { + return ( + + + + ); +} + +const styles = StyleSheet.create({ + container: { flex: 1, backgroundColor: theme.colors.background }, + content: { padding: 16 }, + input: { + height: 40, + borderColor: theme.colors.border, + borderWidth: 1, + borderRadius: 8, + paddingHorizontal: 12, + color: theme.colors.textSecondary, + backgroundColor: theme.colors.background, + marginBottom: 12, + }, + card: { + backgroundColor: theme.colors.cardBackground, + borderRadius: 12, + padding: 16, + marginBottom: 16, + borderWidth: 1, + borderColor: theme.colors.border, + }, + cardTitle: { + fontSize: 16, + fontWeight: 'bold', + color: theme.colors.strongPrimary, + marginBottom: 12, + }, + errorContainer: { + padding: 12, + backgroundColor: '#ffdddd', + borderRadius: 8, + marginBottom: 16, + }, + errorText: { color: '#cc0000', fontSize: 13 }, + buttonContainer: { alignItems: 'center', marginVertical: 12 }, + button: { + backgroundColor: theme.colors.accent, + paddingVertical: 12, + paddingHorizontal: 24, + borderRadius: 8, + alignItems: 'center', + width: '100%', + }, + buttonStop: { backgroundColor: '#ff3b30' }, + buttonDisabled: { backgroundColor: '#d1d1d6' }, + buttonText: { color: '#ffffff', fontWeight: 'bold', fontSize: 15 }, + resultHeader: { + fontSize: 13, + fontWeight: 'bold', + color: theme.colors.textSecondary, + marginTop: 12, + marginBottom: 6, + }, + textOutputContainer: { + minHeight: 100, + padding: 12, + backgroundColor: theme.colors.background, + borderRadius: 8, + borderWidth: 1, + borderColor: theme.colors.border, + }, + committedText: { fontSize: 14, color: theme.colors.textSecondary, lineHeight: 20 }, +}); diff --git a/apps/speech/app/index.tsx b/apps/speech/app/index.tsx index 453b8e3863..4297cf0103 100644 --- a/apps/speech/app/index.tsx +++ b/apps/speech/app/index.tsx @@ -14,6 +14,18 @@ export default function Home() { router.navigate('vad/')}> Voice Activity Detection + router.navigate('microphone-transcription/')} + > + Live Transcription (Mic) + + router.navigate('audio-file-transcription/')} + > + Transcribe Audio File + ); diff --git a/apps/speech/app/microphone-transcription/index.tsx b/apps/speech/app/microphone-transcription/index.tsx new file mode 100644 index 0000000000..f4c8da9704 --- /dev/null +++ b/apps/speech/app/microphone-transcription/index.tsx @@ -0,0 +1,326 @@ +import React, { useEffect, useRef, useState } from 'react'; +import { Platform, View, Text, StyleSheet, ScrollView, TouchableOpacity } from 'react-native'; +import { useSpeechToText, models, WHISPER_SAMPLE_RATE_HZ } from 'react-native-executorch'; +import { AudioManager, AudioRecorder } from 'react-native-audio-api'; +import DeviceInfo from 'react-native-device-info'; + +import ScreenWrapper from '../../components/ScreenWrapper'; +import { ModelPicker } from '../../components/ModelPicker'; +import { ModelStatus } from '../../components/ModelStatus'; +import { theme } from '../../theme'; + +const MODELS = [ + { + name: 'Tiny English (CPU)', + config: models.speechToText.WHISPER.EN.TINY.XNNPACK_FP32, + }, + { + name: 'Tiny English (CoreML)', + config: models.speechToText.WHISPER.EN.TINY.COREML_FP16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Tiny English (MLX)', + config: models.speechToText.WHISPER.EN.TINY.MLX_BF16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Base English (CPU)', + config: models.speechToText.WHISPER.EN.BASE.XNNPACK_FP32, + }, + { + name: 'Base English (CoreML)', + config: models.speechToText.WHISPER.EN.BASE.COREML_FP16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Base English (MLX)', + config: models.speechToText.WHISPER.EN.BASE.MLX_BF16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Small English (CPU)', + config: models.speechToText.WHISPER.EN.SMALL.XNNPACK_FP32, + }, + { + name: 'Small English (CoreML)', + config: models.speechToText.WHISPER.EN.SMALL.COREML_FP16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Small English (MLX)', + config: models.speechToText.WHISPER.EN.SMALL.MLX_BF16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Tiny Multilingual (CPU)', + config: models.speechToText.WHISPER.TINY.XNNPACK_FP32, + }, + { + name: 'Tiny Multilingual (CoreML)', + config: models.speechToText.WHISPER.TINY.COREML_FP16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Tiny Multilingual (MLX)', + config: models.speechToText.WHISPER.TINY.MLX_BF16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Base Multilingual (CPU)', + config: models.speechToText.WHISPER.BASE.XNNPACK_FP32, + }, + { + name: 'Base Multilingual (CoreML)', + config: models.speechToText.WHISPER.BASE.COREML_FP16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Base Multilingual (MLX)', + config: models.speechToText.WHISPER.BASE.MLX_BF16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Small Multilingual (CPU)', + config: models.speechToText.WHISPER.SMALL.XNNPACK_FP32, + }, + { + name: 'Small Multilingual (CoreML)', + config: models.speechToText.WHISPER.SMALL.COREML_FP16, + disabled: Platform.OS !== 'ios', + }, + { + name: 'Small Multilingual (MLX)', + config: models.speechToText.WHISPER.SMALL.MLX_BF16, + disabled: Platform.OS !== 'ios', + }, +]; + +const isSimulator = DeviceInfo.isEmulatorSync(); + +function STTContent() { + const [selectedModel, setSelectedModel] = useState(MODELS[0]!); + const [selectedLanguage, setSelectedLanguage] = useState('en'); + const [status, setStatus] = useState('Idle'); + const [committedText, setCommittedText] = useState(''); + const [nonCommittedText, setNonCommittedText] = useState(''); + const [isRecording, setIsRecording] = useState(false); + const [runError, setRunError] = useState(null); + + const { + isReady: isSttReady, + stream, + streamInsert, + streamStop, + downloadProgress, + error: modelError, + } = useSpeechToText(selectedModel.config); + + const supportedLanguages = selectedModel.config.supportedLanguages; + + const recorderRef = useRef(null); + + useEffect(() => { + AudioManager.setAudioSessionOptions({ + iosCategory: 'playAndRecord', + iosMode: 'spokenAudio', + iosOptions: ['allowBluetoothHFP', 'defaultToSpeaker'], + }); + AudioManager.requestRecordingPermissions(); + return () => { + recorderRef.current?.stop().catch(() => {}); + if (streamStop) streamStop(); + }; + }, [streamStop]); + + const startRecording = async () => { + if (!isSttReady || !stream || !streamInsert || isRecording) return; + setRunError(null); + setCommittedText(''); + setNonCommittedText(''); + setStatus('Streaming...'); + setIsRecording(true); + + (async () => { + try { + const textStream = stream({ language: selectedLanguage as any }); + for await (const result of textStream) { + setCommittedText(result.committed); + setNonCommittedText(result.nonCommitted); + } + } catch (err) { + const errMsg = err instanceof Error ? err.message : String(err); + setRunError(errMsg); + } + })(); + + const recorder = new AudioRecorder(); + recorderRef.current = recorder; + + recorder.onAudioReady( + { sampleRate: WHISPER_SAMPLE_RATE_HZ, bufferLength: 4096, channelCount: 1 }, + (event: any) => { + const samples = event.buffer.getChannelData(0); + streamInsert(new Float32Array(samples)); + } + ); + + const result = await recorder.start(); + if (result.status === 'error') { + setRunError(result.message); + setIsRecording(false); + setStatus('Error'); + } + }; + + const stopRecording = async () => { + if (!isRecording) return; + if (recorderRef.current) { + await recorderRef.current.stop().catch(() => {}); + recorderRef.current = null; + } + if (streamStop) streamStop(); + setIsRecording(false); + setStatus('Done'); + }; + + const isModelBusy = status.includes('...'); + const isMicDisabled = isSimulator || !isSttReady; + + return ( + + + ({ + label: m.name, + value: m, + disabled: isModelBusy || !!m.disabled, + }))} + selectedValue={selectedModel} + onValueChange={(m) => { + setSelectedModel(m); + setSelectedLanguage(m.config.supportedLanguages[0]!); + }} + /> + + {supportedLanguages.length > 1 && ( + ({ + label: lang, + value: lang, + disabled: isModelBusy, + }))} + selectedValue={selectedLanguage} + onValueChange={setSelectedLanguage} + /> + )} + + + {runError && ( + + {runError} + + )} + + {/* Live Mic Panel */} + + Live Microphone Gated STT + + {!isRecording ? ( + + Start Streaming + + ) : ( + + Stop Streaming + + )} + + + Transcription Output: + + + {committedText} + {nonCommittedText ? ( + {nonCommittedText} + ) : null} + + + + + ); +} + +export default function STTScreen() { + return ( + + + + ); +} + +const styles = StyleSheet.create({ + container: { flex: 1, backgroundColor: theme.colors.background }, + content: { padding: 16 }, + card: { + backgroundColor: theme.colors.cardBackground, + borderRadius: 12, + padding: 16, + marginBottom: 16, + borderWidth: 1, + borderColor: theme.colors.border, + }, + cardTitle: { + fontSize: 16, + fontWeight: 'bold', + color: theme.colors.strongPrimary, + marginBottom: 12, + }, + errorContainer: { + padding: 12, + backgroundColor: '#ffdddd', + borderRadius: 8, + marginBottom: 16, + }, + errorText: { color: '#cc0000', fontSize: 13 }, + buttonContainer: { alignItems: 'center', marginVertical: 12 }, + button: { + backgroundColor: theme.colors.accent, + paddingVertical: 12, + paddingHorizontal: 24, + borderRadius: 8, + alignItems: 'center', + width: '100%', + }, + buttonStop: { backgroundColor: '#ff3b30' }, + buttonDisabled: { backgroundColor: '#d1d1d6' }, + buttonText: { color: '#ffffff', fontWeight: 'bold', fontSize: 15 }, + resultHeader: { + fontSize: 13, + fontWeight: 'bold', + color: theme.colors.textSecondary, + marginTop: 12, + marginBottom: 6, + }, + textOutputContainer: { + minHeight: 100, + padding: 12, + backgroundColor: theme.colors.background, + borderRadius: 8, + borderWidth: 1, + borderColor: theme.colors.border, + }, + committedText: { fontSize: 14, color: theme.colors.textSecondary, lineHeight: 20 }, + nonCommittedText: { color: theme.colors.textMuted, fontStyle: 'italic' }, +}); diff --git a/apps/speech/components/ModelPicker.tsx b/apps/speech/components/ModelPicker.tsx new file mode 100644 index 0000000000..04051beca5 --- /dev/null +++ b/apps/speech/components/ModelPicker.tsx @@ -0,0 +1,105 @@ +import React from 'react'; +import { View, Text, ScrollView, TouchableOpacity, StyleSheet } from 'react-native'; + +import { theme } from '../theme'; + +export type ModelOption = { + label: string; + value: T; + disabled?: boolean; +}; + +interface ModelPickerProps { + label: string; + options: ModelOption[]; + selectedValue: T; + onValueChange: (value: T) => void; +} + +export function ModelPicker({ + label, + options, + selectedValue, + onValueChange, +}: ModelPickerProps) { + return ( + + {label} + + {options.map((option, index) => { + const isSelected = option.value === selectedValue; + const isDisabled = option.disabled; + return ( + onValueChange(option.value)} + disabled={isDisabled} + > + + {option.label} + + + ); + })} + + + ); +} + +const styles = StyleSheet.create({ + container: { + width: '100%', + marginBottom: 4, + }, + label: { + fontSize: 12, + fontWeight: '700', + color: theme.colors.textMuted, + textTransform: 'uppercase', + marginBottom: 8, + }, + scroll: { + flexDirection: 'row', + gap: 8, + }, + chip: { + backgroundColor: theme.colors.placeholderBackground, + paddingHorizontal: 14, + paddingVertical: 8, + borderRadius: 20, + }, + chipActive: { + backgroundColor: theme.colors.primary, + }, + chipDisabled: { + backgroundColor: theme.colors.placeholderBackground, + opacity: 0.5, + }, + text: { + fontSize: 13, + color: theme.colors.textSecondary, + fontWeight: '500', + }, + textActive: { + color: '#fff', + fontWeight: '600', + }, + textDisabled: { + color: theme.colors.textMuted, + }, +}); diff --git a/packages/react-native-executorch/src/extensions/speech/tasks/whisperSpeechToText.ts b/packages/react-native-executorch/src/extensions/speech/tasks/whisperSpeechToText.ts new file mode 100644 index 0000000000..e1f14bfbb1 --- /dev/null +++ b/packages/react-native-executorch/src/extensions/speech/tasks/whisperSpeechToText.ts @@ -0,0 +1,381 @@ +import type { WorkletRuntime } from 'react-native-worklets'; +import { scheduleOnRN, createSynchronizable } from 'react-native-worklets'; + +import { tensor } from '../../../core/tensor'; +import { loadModel } from '../../../core/model'; +import { validateModelSchema, SymbolicTensor } from '../../../core/modelSchema'; +import { wrapAsync } from '../../../core/runtime'; + +import { argmax } from '../../../extensions/math'; +import { loadTokenizer } from '../../nlp/tokenizer'; +import { + createFsmnVoiceActivityDetector, + type FsmnVadModel, + type VadStreamOptions, +} from './fsmnVoiceActivityDetection'; + +/** + * Sample rate (Hz) Whisper models expect their input waveform to be at. + * @category Constants + */ +export const WHISPER_SAMPLE_RATE_HZ = 16000; + +const MAX_SEQ_LEN = 128; // maximum decoder output tokens per chunk +const MIN_CHUNK_SIZE = 201; // shortest audio slice (samples) the model was exported for +const CHUNK_LENGTH_SECONDS = 29; // Whisper's fixed context window length +const STRIDE_SIZE = 1 * WHISPER_SAMPLE_RATE_HZ; // minimum new samples before processing and retain-window size during silence (1s) +const BUFFER_SIZE = CHUNK_LENGTH_SECONDS * WHISPER_SAMPLE_RATE_HZ; // samples per full chunk + +/** + * Language codes supported by Whisper multilingual models. English-only + * model variants only accept `'en'`. + * @category Constants + */ +// prettier-ignore +export const WHISPER_LANGUAGES = [ + 'en', 'zh', 'de', 'es', 'ru', 'ko', 'fr', 'ja', 'pt', 'tr', + 'pl', 'ca', 'nl', 'ar', 'sv', 'it', 'id', 'hi', 'fi', 'vi', + 'he', 'uk', 'el', 'ms', 'cs', 'ro', 'da', 'hu', 'ta', 'no', + 'th', 'ur', 'hr', 'bg', 'lt', 'la', 'mi', 'ml', 'cy', 'sk', + 'te', 'fa', 'lv', 'bn', 'sr', 'az', 'sl', 'kn', 'et', 'mk', + 'br', 'eu', 'is', 'hy', 'ne', 'mn', 'bs', 'kk', 'sq', 'sw', + 'gl', 'mr', 'pa', 'si', 'km', 'sn', 'yo', 'so', 'af', 'oc', + 'ka', 'be', 'tg', 'sd', 'gu', 'am', 'yi', 'lo', 'uz', 'fo', + 'ht', 'ps', 'tk', 'nn', 'mt', 'sa', 'lb', 'my', 'bo', 'tl', + 'mg', 'as', 'tt', 'haw', 'ln', 'ha', 'ba', 'jw', 'su', 'yue' +] as const; + +/** + * Union type of all language codes supported by Whisper. Derived from + * {@link WHISPER_LANGUAGES}. + * @category Types + */ +export type WhisperLanguage = (typeof WHISPER_LANGUAGES)[number]; + +/** + * Options passed to a single transcription call. + * @category Types + * @property language - Whisper language code of the spoken audio. Must be one of + * the {@link WhisperLanguage} values declared in the model's + * `supportedLanguages` list. + */ +export type WhisperSttOptions = { + readonly language: L; +}; + +/** + * Options for the live-streaming transcription API. + * Extends {@link WhisperSttOptions} with optional VAD tuning. + * @category Types + * @property vadOptions - Fine-tuning knobs forwarded to the voice-activity + * detector. Omit to use the detector's built-in defaults. + */ +export type WhisperStreamOptions = + WhisperSttOptions & { readonly vadOptions?: VadStreamOptions }; + +/** + * Paths and metadata required to instantiate a Whisper speech-to-text model. + * @category Types + */ +export type WhisperSttModel = { + readonly modelPath: string; + readonly tokenizerPath: string; + readonly supportedLanguages: readonly L[]; + readonly vadModel: FsmnVadModel; +}; + +/** + * Loads a Whisper model and returns a set of transcription helpers. + * @category Typescript API + * @param config Model paths and supported-language metadata. See {@link + * WhisperSttModel}. + * @param runtime Optional worklet runtime thread on which to run the model + * execution. + * @returns A promise resolving to an object containing transcription and + * disposal controls. + */ +export async function createWhisperSpeechToText( + config: WhisperSttModel, + runtime?: WorkletRuntime +): Promise<{ + /** + * Releases all allocated native resources. + */ + dispose: () => void; + + /** + * Asynchronously transcribes a pre-recorded mono waveform sampled at + * {@link WHISPER_SAMPLE_RATE_HZ}. + * @param audio Raw 16 kHz mono PCM samples. + * @param options Transcription options. + * @param onToken Optional callback fired on the RN thread for each decoded + * token. + * @returns A promise resolving to the full transcript string. + */ + transcribe: ( + audio: Float32Array, + options: WhisperSttOptions, + onToken?: (token: string) => void + ) => Promise; + + /** + * Synchronous version of {@link transcribe} to be executed directly on the + * caller or worklet thread. + */ + transcribeWorklet: ( + audio: Float32Array, + options: WhisperSttOptions, + onToken?: (token: string) => void + ) => string; + + /** + * Interrupts and stops any active transcription call. + */ + transcribeStop: () => void; + + /** + * Async generator for real-time microphone transcription. Feed audio with + * {@link streamInsert} and stop with {@link streamStop}. Yields `{ committed, + * nonCommitted }` on every VAD or transcription event: `committed` is the + * finalized transcript so far; `nonCommitted` is the in-progress text that + * may still change. + * @param options Stream options (language and optional VAD tuning). + */ + stream: ( + options: WhisperStreamOptions + ) => AsyncGenerator<{ committed: string; nonCommitted: string }>; + + /** + * Signals the {@link stream} generator to finalize the current segment and + * return. Safe to call even when streaming is not active. + */ + streamStop: () => void; + + /** + * Appends a new PCM chunk to the live streaming buffer consumed by + * {@link stream}. Ignored when streaming is not active. + * @param audioChunk The newly captured audio samples. + */ + streamInsert: (audioChunk: Float32Array) => void; +}> { + const { modelPath, tokenizerPath, supportedLanguages, vadModel } = config; + const model = await wrapAsync(loadModel, runtime)(modelPath); + const tokenizer = await wrapAsync(loadTokenizer, runtime)(tokenizerPath); + const voiceDetector = await createFsmnVoiceActivityDetector(vadModel, runtime); + + const eotToken = tokenizer.tokenToId('<|endoftext|>')!; + const isEnglishOnly = supportedLanguages.length === 1 && supportedLanguages[0] === 'en'; + + const encMeta = validateModelSchema( + model, + 'encode', + [SymbolicTensor('float32', ['T_audio'])], + [SymbolicTensor('float32', [1, 'Seq', 'State'])] + ); + + const encSeqLen = encMeta.outputTensorMeta[0]!.shape[1]!; + const encStateDim = encMeta.outputTensorMeta[0]!.shape[2]!; + + validateModelSchema( + model, + 'decode', + [ + SymbolicTensor('int64', [1, 'Tokens']), + SymbolicTensor('int64', ['Tokens']), + SymbolicTensor('float32', [1, encSeqLen, encStateDim]), + ], + [SymbolicTensor('float32', [1, 'Tokens', 'Vocab'])] + ); + + const tensors = [ + tensor('int64', [1]), // tPosition + tensor('int64', [1, 1]), // tToken + tensor('int32', [1, 1, 1]), // tArgmax + tensor('float32', [1, encSeqLen, encStateDim]), //tEncodings + tensor('float32', [1, 1, tokenizer.getVocabSize()]), // tLogits + ] as const; + + const [tPosition, tToken, tArgmax, tEncodings, tLogits] = tensors; + const isCancelled = createSynchronizable(false); + + const dispose = () => { + tensors.forEach((t) => t.dispose()); + voiceDetector.dispose(); + tokenizer.dispose(); + model.dispose(); + }; + + const decode = (token: number, position: number): number => { + 'worklet'; + tPosition.setData(new BigInt64Array([BigInt(position)])); + tToken.setData(new BigInt64Array([BigInt(token)])); + model.execute('decode', [tToken, tPosition, tEncodings], [tLogits]); + return tLogits.through(argmax, tArgmax).getData(new Int32Array(1))[0]!; + }; + + const transcribeWorklet = ( + audio: Float32Array, + options: WhisperSttOptions, + onToken?: (token: string) => void + ): string => { + 'worklet'; + isCancelled.setBlocking(false); + + const promptTokenStrings = isEnglishOnly + ? ['<|startoftranscript|>', '<|notimestamps|>'] + : ['<|startoftranscript|>', `<|${options.language}|>`, '<|transcribe|>', '<|notimestamps|>']; + + if (!isEnglishOnly && tokenizer.tokenToId(`<|${options.language}|>`) === undefined) { + throw new Error(`Language "${options.language}" is not recognized.`); + } + const promptTokens = promptTokenStrings.map((token) => tokenizer.tokenToId(token)); + const maxNewTokens = MAX_SEQ_LEN - promptTokens.length; + + let text = ''; + let offset = 0; + + while (offset < audio.length) { + if (isCancelled.getBlocking()) break; + + const audioChunk = audio.slice(offset, Math.min(offset + BUFFER_SIZE, audio.length)); + if (audioChunk.length < MIN_CHUNK_SIZE) { + break; + } + + const tAudioInput = tensor('float32', [audioChunk.length], audioChunk); + try { + model.execute('encode', [tAudioInput], [tEncodings]); + } finally { + tAudioInput.dispose(); + } + + let nextToken = eotToken; + let position = promptTokens.length; + promptTokens.forEach((token, pos) => (nextToken = decode(token, pos))); + + const generated: number[] = []; + while (generated.length < maxNewTokens && nextToken !== eotToken) { + if (isCancelled.getBlocking()) break; + + generated.push(nextToken); + if (onToken) scheduleOnRN(onToken, tokenizer.decode([nextToken])); + nextToken = decode(nextToken, position); + position++; + } + + text += tokenizer.decode(generated); + offset += BUFFER_SIZE; + } + + return text.trim(); + }; + + const transcribe = wrapAsync(transcribeWorklet, runtime); + const transcribeStop = () => isCancelled.setBlocking(true); + + let isStreaming = false; + let audioBuffer = new Float32Array(0); + let signal: (() => void) | null = null; + + const streamInsert = (audioChunk: Float32Array): void => { + if (!isStreaming) return; + + const next = new Float32Array(audioBuffer.length + audioChunk.length); + next.set(audioBuffer); + next.set(audioChunk, audioBuffer.length); + audioBuffer = next; + signal?.(); + signal = null; + }; + + const streamStop = (): void => { + if (!isStreaming) return; + + isStreaming = false; + signal?.(); + signal = null; + }; + + async function* stream( + options: WhisperStreamOptions + ): AsyncGenerator<{ committed: string; nonCommitted: string }> { + if (isStreaming) { + throw new Error('Streaming is already in progress'); + } + isStreaming = true; + audioBuffer = new Float32Array(0); + + voiceDetector.resetStream(); + + let isSpeaking = false; + let currentText = ''; + let committedText = ''; + let processedLength = 0; + + const commit = () => { + if (currentText !== '') { + committedText += (committedText ? ' ' : '') + currentText; + currentText = ''; + } + audioBuffer = audioBuffer.slice(processedLength); + processedLength = 0; + }; + + try { + while (isStreaming) { + if (audioBuffer.length - processedLength < STRIDE_SIZE) { + await new Promise((resolve) => (signal = resolve)); + continue; + } + + const newSamples = audioBuffer.slice(processedLength); + processedLength = audioBuffer.length; + + const event = voiceDetector.detectVoiceOnStream(newSamples, options.vadOptions); + switch (event) { + case 'speechStart': + isSpeaking = true; + break; + case 'speechEnd': + isSpeaking = false; + currentText = await transcribe(audioBuffer.slice(0, processedLength), options); + commit(); + yield { committed: committedText, nonCommitted: '' }; + continue; + } + + if (isSpeaking) { + currentText = await transcribe(audioBuffer.slice(0, processedLength), options); + if (processedLength >= BUFFER_SIZE) commit(); + yield { committed: committedText, nonCommitted: currentText }; + } else { + const retainSamples = Math.min(audioBuffer.length, STRIDE_SIZE); + audioBuffer = audioBuffer.slice(audioBuffer.length - retainSamples); + processedLength = audioBuffer.length; + yield { committed: committedText, nonCommitted: '' }; + } + } + + if (isSpeaking && audioBuffer.length >= MIN_CHUNK_SIZE) { + currentText = await transcribe(audioBuffer, options); + commit(); + yield { committed: committedText, nonCommitted: '' }; + } + } finally { + signal = null; + isStreaming = false; + voiceDetector.resetStream(); + audioBuffer = new Float32Array(0); + } + } + + return { + dispose, + transcribe, + transcribeWorklet, + transcribeStop, + stream, + streamStop, + streamInsert, + }; +} diff --git a/packages/react-native-executorch/src/hooks/useSpeechToText.ts b/packages/react-native-executorch/src/hooks/useSpeechToText.ts new file mode 100644 index 0000000000..d75f14b8fe --- /dev/null +++ b/packages/react-native-executorch/src/hooks/useSpeechToText.ts @@ -0,0 +1,71 @@ +import { useModel } from './useModel'; +import { useResourceDownload } from './useResourceDownload'; +import { + createWhisperSpeechToText, + type WhisperSttModel, + type WhisperLanguage, +} from '../extensions/speech/tasks/whisperSpeechToText'; + +/** + * React hook to load and run a Whisper speech-to-text model. + * + * This hook manages downloading (if it's a remote URL) and loading the model, + * tokenizer, and FSMN-VAD, compiling them, tracking download progress and + * compilation errors, and cleaning up native model memory when the component + * unmounts or configuration changes. + * @category Hooks + * @param config The Whisper speech-to-text model configuration. + * @param options Hook options. + * @param options.preventLoad If true, prevents downloading and compiling the + * model. + * @returns An object containing the model's loading state, error, download + * progress, and transcription functions. + */ +export function useSpeechToText( + config: WhisperSttModel, + options?: { preventLoad?: boolean } +) { + const vadModelPath = config.vadModel.modelPath; + const vadResource = useResourceDownload(vadModelPath, options?.preventLoad); + const modelResource = useResourceDownload(config.modelPath, options?.preventLoad); + const tokenizerResource = useResourceDownload(config.tokenizerPath, options?.preventLoad); + + const localVadPath = vadResource.localPath; + const localModelPath = modelResource.localPath; + const localTokenizerPath = tokenizerResource.localPath; + + const isResourcesReady = !!(localModelPath && localTokenizerPath && localVadPath); + const whisperConfig = isResourcesReady + ? { + ...config, + modelPath: localModelPath!, + tokenizerPath: localTokenizerPath!, + vadModel: { ...config.vadModel, modelPath: localVadPath! }, + } + : null; + + const { model, error: modelError } = useModel(createWhisperSpeechToText, whisperConfig, [ + localVadPath, + localModelPath, + localTokenizerPath, + ]); + + const error = + vadResource.downloadError || + modelResource.downloadError || + tokenizerResource.downloadError || + modelError; + + return { + isReady: !!model, + error, + downloadProgress: modelResource.downloadProgress, + localPath: modelResource.localPath, + transcribe: model?.transcribe, + transcribeWorklet: model?.transcribeWorklet, + transcribeStop: model?.transcribeStop, + stream: model?.stream, + streamInsert: model?.streamInsert, + streamStop: model?.streamStop, + }; +} diff --git a/packages/react-native-executorch/src/index.ts b/packages/react-native-executorch/src/index.ts index ea9bca1a96..3aad5c8d77 100644 --- a/packages/react-native-executorch/src/index.ts +++ b/packages/react-native-executorch/src/index.ts @@ -9,6 +9,7 @@ export * from './hooks/useTokenizer'; export * from './hooks/useTextEmbedder'; export * from './hooks/useImageEmbedder'; export * from './hooks/useVoiceActivityDetector'; +export * from './hooks/useSpeechToText'; export * from './hooks/useTextToImage'; export * from './hooks/useResourceDownload'; export * from './hooks/useModel'; @@ -29,6 +30,7 @@ export * from './extensions/cv/tasks/sdxsTextToImage'; export * from './extensions/nlp/tasks/tokenization'; export * from './extensions/nlp/tasks/textEmbedding'; export * from './extensions/speech/tasks/fsmnVoiceActivityDetection'; +export * from './extensions/speech/tasks/whisperSpeechToText'; // Core primitives — for library builders and power users export { tensor } from './core/tensor'; diff --git a/packages/react-native-executorch/src/models.ts b/packages/react-native-executorch/src/models.ts index 1665f724d5..fc8e7a5e7d 100644 --- a/packages/react-native-executorch/src/models.ts +++ b/packages/react-native-executorch/src/models.ts @@ -8,6 +8,10 @@ import type { ImageEmbedderModel } from './extensions/cv/tasks/imageEmbedding'; import type { SdxsTextToImageModel } from './extensions/cv/tasks/sdxsTextToImage'; import type { TextEmbedderModel } from './extensions/nlp/tasks/textEmbedding'; import type { FsmnVadModel } from './extensions/speech/tasks/fsmnVoiceActivityDetection'; +import { + type WhisperSttModel, + WHISPER_LANGUAGES, +} from './extensions/speech/tasks/whisperSpeechToText'; import { IMAGENET_NORM, IMAGENET1K_LABELS, @@ -606,6 +610,123 @@ const FSMN_VAD_XNNPACK_FP32: FsmnVadModel = { }, }; +// ============================================================================= +// Speech-To-Text +// ============================================================================= +const WHISPER_TINY_EN_XNNPACK_FP32: WhisperSttModel<'en'> = { + modelPath: `${BASE_URL}-whisper-tiny.en/${NEXT_VERSION_TAG}/xnnpack/whisper_tiny_en_xnnpack_fp32.pte`, + tokenizerPath: `${BASE_URL}-whisper-tiny.en/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: ['en'], + vadModel: FSMN_VAD_XNNPACK_FP32, +}; +const WHISPER_TINY_EN_COREML_FP16: WhisperSttModel<'en'> = { + modelPath: `${BASE_URL}-whisper-tiny.en/${NEXT_VERSION_TAG}/coreml/whisper_tiny_en_coreml_fp16.pte`, + tokenizerPath: `${BASE_URL}-whisper-tiny.en/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: ['en'], + vadModel: FSMN_VAD_XNNPACK_FP32, +}; +const WHISPER_TINY_EN_MLX_BF16: WhisperSttModel<'en'> = { + modelPath: `${BASE_URL}-whisper-tiny.en/${NEXT_VERSION_TAG}/mlx/whisper_tiny_en_mlx_bf16.pte`, + tokenizerPath: `${BASE_URL}-whisper-tiny.en/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: ['en'], + vadModel: FSMN_VAD_XNNPACK_FP32, +}; + +const WHISPER_TINY_XNNPACK_FP32: WhisperSttModel = { + modelPath: `${BASE_URL}-whisper-tiny/${NEXT_VERSION_TAG}/xnnpack/whisper_tiny_xnnpack_fp32.pte`, + tokenizerPath: `${BASE_URL}-whisper-tiny/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: WHISPER_LANGUAGES, + vadModel: FSMN_VAD_XNNPACK_FP32, +}; +const WHISPER_TINY_COREML_FP16: WhisperSttModel = { + modelPath: `${BASE_URL}-whisper-tiny/${NEXT_VERSION_TAG}/coreml/whisper_tiny_coreml_fp16.pte`, + tokenizerPath: `${BASE_URL}-whisper-tiny/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: WHISPER_LANGUAGES, + vadModel: FSMN_VAD_XNNPACK_FP32, +}; +const WHISPER_TINY_MLX_BF16: WhisperSttModel = { + modelPath: `${BASE_URL}-whisper-tiny/${NEXT_VERSION_TAG}/mlx/whisper_tiny_mlx_bf16.pte`, + tokenizerPath: `${BASE_URL}-whisper-tiny/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: WHISPER_LANGUAGES, + vadModel: FSMN_VAD_XNNPACK_FP32, +}; + +const WHISPER_BASE_EN_XNNPACK_FP32: WhisperSttModel<'en'> = { + modelPath: `${BASE_URL}-whisper-base.en/${NEXT_VERSION_TAG}/xnnpack/whisper_base_en_xnnpack_fp32.pte`, + tokenizerPath: `${BASE_URL}-whisper-base.en/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: ['en'], + vadModel: FSMN_VAD_XNNPACK_FP32, +}; +const WHISPER_BASE_EN_COREML_FP16: WhisperSttModel<'en'> = { + modelPath: `${BASE_URL}-whisper-base.en/${NEXT_VERSION_TAG}/coreml/whisper_base_en_coreml_fp16.pte`, + tokenizerPath: `${BASE_URL}-whisper-base.en/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: ['en'], + vadModel: FSMN_VAD_XNNPACK_FP32, +}; +const WHISPER_BASE_EN_MLX_BF16: WhisperSttModel<'en'> = { + modelPath: `${BASE_URL}-whisper-base.en/${NEXT_VERSION_TAG}/mlx/whisper_base_en_mlx_bf16.pte`, + tokenizerPath: `${BASE_URL}-whisper-base.en/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: ['en'], + vadModel: FSMN_VAD_XNNPACK_FP32, +}; + +const WHISPER_BASE_XNNPACK_FP32: WhisperSttModel = { + modelPath: `${BASE_URL}-whisper-base/${NEXT_VERSION_TAG}/xnnpack/whisper_base_xnnpack_fp32.pte`, + tokenizerPath: `${BASE_URL}-whisper-base/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: WHISPER_LANGUAGES, + vadModel: FSMN_VAD_XNNPACK_FP32, +}; +const WHISPER_BASE_COREML_FP16: WhisperSttModel = { + modelPath: `${BASE_URL}-whisper-base/${NEXT_VERSION_TAG}/coreml/whisper_base_coreml_fp16.pte`, + tokenizerPath: `${BASE_URL}-whisper-base/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: WHISPER_LANGUAGES, + vadModel: FSMN_VAD_XNNPACK_FP32, +}; +const WHISPER_BASE_MLX_BF16: WhisperSttModel = { + modelPath: `${BASE_URL}-whisper-base/${NEXT_VERSION_TAG}/mlx/whisper_base_mlx_bf16.pte`, + tokenizerPath: `${BASE_URL}-whisper-base/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: WHISPER_LANGUAGES, + vadModel: FSMN_VAD_XNNPACK_FP32, +}; + +const WHISPER_SMALL_EN_XNNPACK_FP32: WhisperSttModel<'en'> = { + modelPath: `${BASE_URL}-whisper-small.en/${NEXT_VERSION_TAG}/xnnpack/whisper_small_en_xnnpack_fp32.pte`, + tokenizerPath: `${BASE_URL}-whisper-small.en/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: ['en'], + vadModel: FSMN_VAD_XNNPACK_FP32, +}; +const WHISPER_SMALL_EN_COREML_FP16: WhisperSttModel<'en'> = { + modelPath: `${BASE_URL}-whisper-small.en/${NEXT_VERSION_TAG}/coreml/whisper_small_en_coreml_fp16.pte`, + tokenizerPath: `${BASE_URL}-whisper-small.en/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: ['en'], + vadModel: FSMN_VAD_XNNPACK_FP32, +}; +const WHISPER_SMALL_EN_MLX_BF16: WhisperSttModel<'en'> = { + modelPath: `${BASE_URL}-whisper-small.en/${NEXT_VERSION_TAG}/mlx/whisper_small_en_mlx_bf16.pte`, + tokenizerPath: `${BASE_URL}-whisper-small.en/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: ['en'], + vadModel: FSMN_VAD_XNNPACK_FP32, +}; + +const WHISPER_SMALL_XNNPACK_FP32: WhisperSttModel = { + modelPath: `${BASE_URL}-whisper-small/${NEXT_VERSION_TAG}/xnnpack/whisper_small_xnnpack_fp32.pte`, + tokenizerPath: `${BASE_URL}-whisper-small/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: WHISPER_LANGUAGES, + vadModel: FSMN_VAD_XNNPACK_FP32, +}; +const WHISPER_SMALL_COREML_FP16: WhisperSttModel = { + modelPath: `${BASE_URL}-whisper-small/${NEXT_VERSION_TAG}/coreml/whisper_small_coreml_fp16.pte`, + tokenizerPath: `${BASE_URL}-whisper-small/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: WHISPER_LANGUAGES, + vadModel: FSMN_VAD_XNNPACK_FP32, +}; +const WHISPER_SMALL_MLX_BF16: WhisperSttModel = { + modelPath: `${BASE_URL}-whisper-small/${NEXT_VERSION_TAG}/mlx/whisper_small_mlx_bf16.pte`, + tokenizerPath: `${BASE_URL}-whisper-small/${NEXT_VERSION_TAG}/tokenizer.json`, + supportedLanguages: WHISPER_LANGUAGES, + vadModel: FSMN_VAD_XNNPACK_FP32, +}; + // ============================================================================= // Text to Image // ============================================================================= @@ -831,6 +952,50 @@ export const models = { XNNPACK_FP32: FSMN_VAD_XNNPACK_FP32, }, }, + speechToText: { + WHISPER: { + ...WHISPER_TINY_XNNPACK_FP32, + TINY: { + ...WHISPER_TINY_XNNPACK_FP32, + XNNPACK_FP32: WHISPER_TINY_XNNPACK_FP32, + COREML_FP16: WHISPER_TINY_COREML_FP16, + MLX_BF16: WHISPER_TINY_MLX_BF16, + }, + BASE: { + ...WHISPER_BASE_XNNPACK_FP32, + XNNPACK_FP32: WHISPER_BASE_XNNPACK_FP32, + COREML_FP16: WHISPER_BASE_COREML_FP16, + MLX_BF16: WHISPER_BASE_MLX_BF16, + }, + SMALL: { + ...WHISPER_SMALL_XNNPACK_FP32, + XNNPACK_FP32: WHISPER_SMALL_XNNPACK_FP32, + COREML_FP16: WHISPER_SMALL_COREML_FP16, + MLX_BF16: WHISPER_SMALL_MLX_BF16, + }, + EN: { + ...WHISPER_TINY_EN_XNNPACK_FP32, + TINY: { + ...WHISPER_TINY_EN_XNNPACK_FP32, + XNNPACK_FP32: WHISPER_TINY_EN_XNNPACK_FP32, + COREML_FP16: WHISPER_TINY_EN_COREML_FP16, + MLX_BF16: WHISPER_TINY_EN_MLX_BF16, + }, + BASE: { + ...WHISPER_BASE_EN_XNNPACK_FP32, + XNNPACK_FP32: WHISPER_BASE_EN_XNNPACK_FP32, + COREML_FP16: WHISPER_BASE_EN_COREML_FP16, + MLX_BF16: WHISPER_BASE_EN_MLX_BF16, + }, + SMALL: { + ...WHISPER_SMALL_EN_XNNPACK_FP32, + XNNPACK_FP32: WHISPER_SMALL_EN_XNNPACK_FP32, + COREML_FP16: WHISPER_SMALL_EN_COREML_FP16, + MLX_BF16: WHISPER_SMALL_EN_MLX_BF16, + }, + }, + }, + }, tokenizer: { ALL_MINILM_L6_V2: ALL_MINILM_L6_V2_TOKENIZER, },