react-native-executorch 0.5.2 → 0.5.4

package/src/modules/natural_language_processing/SpeechToTextModule.ts

@@ -1,84 +1,154 @@
+import { Logger } from '../../common/Logger';
 import { DecodingOptions, SpeechToTextModelConfig } from '../../types/stt';
-import { ASR } from '../../utils/SpeechToTextModule/ASR';
-import { OnlineASRProcessor } from '../../utils/SpeechToTextModule/OnlineProcessor';
+import { ResourceFetcher } from '../../utils/ResourceFetcher';
 
 export class SpeechToTextModule {
-  private modelConfig!: SpeechToTextModelConfig;
-  private asr: ASR = new ASR();
+  private nativeModule: any;
 
-  private processor: OnlineASRProcessor = new OnlineASRProcessor(this.asr);
-  private isStreaming = false;
-  private readyToProcess = false;
-  private minAudioSamples: number = 1 * 16000; // 1 second
+  private modelConfig!: SpeechToTextModelConfig;
 
   public async load(
     model: SpeechToTextModelConfig,
     onDownloadProgressCallback: (progress: number) => void = () => {}
   ) {
     this.modelConfig = model;
-    return this.asr.load(model, onDownloadProgressCallback);
+
+    const tokenizerLoadPromise = ResourceFetcher.fetch(
+      undefined,
+      model.tokenizerSource
+    );
+    const encoderDecoderPromise = ResourceFetcher.fetch(
+      onDownloadProgressCallback,
+      model.encoderSource,
+      model.decoderSource
+    );
+    const [tokenizerSources, encoderDecoderResults] = await Promise.all([
+      tokenizerLoadPromise,
+      encoderDecoderPromise,
+    ]);
+    const encoderSource = encoderDecoderResults?.[0];
+    const decoderSource = encoderDecoderResults?.[1];
+    if (!encoderSource || !decoderSource || !tokenizerSources) {
+      throw new Error('Download interrupted.');
+    }
+    this.nativeModule = await global.loadSpeechToText(
+      encoderSource,
+      decoderSource,
+      tokenizerSources[0]!
+    );
   }
 
-  public async encode(waveform: Float32Array): Promise<void> {
-    return this.asr.encode(waveform);
+  public async encode(
+    waveform: Float32Array | number[]
+  ): Promise<Float32Array> {
+    if (Array.isArray(waveform)) {
+      Logger.info(
+        'Passing waveform as number[] is deprecated, use Float32Array instead'
+      );
+      waveform = new Float32Array(waveform);
+    }
+    return new Float32Array(await this.nativeModule.encode(waveform));
   }
 
-  public async decode(tokens: number[]): Promise<Float32Array> {
-    return this.asr.decode(tokens);
+  public async decode(
+    tokens: Int32Array | number[],
+    encoderOutput: Float32Array | number[]
+  ): Promise<Float32Array> {
+    if (Array.isArray(tokens)) {
+      Logger.info(
+        'Passing tokens as number[] is deprecated, use Int32Array instead'
+      );
+      tokens = new Int32Array(tokens);
+    }
+    if (Array.isArray(encoderOutput)) {
+      Logger.info(
+        'Passing encoderOutput as number[] is deprecated, use Float32Array instead'
+      );
+      encoderOutput = new Float32Array(encoderOutput);
+    }
+    return new Float32Array(
+      await this.nativeModule.decode(tokens, encoderOutput)
+    );
   }
 
   public async transcribe(
-    waveform: number[],
+    waveform: Float32Array | number[],
     options: DecodingOptions = {}
   ): Promise<string> {
     this.validateOptions(options);
 
-    const segments = await this.asr.transcribe(waveform, options);
-
-    let transcription = '';
-    for (const segment of segments) {
-      for (const word of segment.words) {
-        transcription += ` ${word.word}`;
-      }
+    if (Array.isArray(waveform)) {
+      Logger.info(
+        'Passing waveform as number[] is deprecated, use Float32Array instead'
+      );
+      waveform = new Float32Array(waveform);
     }
 
-    return transcription.trim();
+    return this.nativeModule.transcribe(waveform, options.language || '');
   }
 
-  public async *stream(options: DecodingOptions = {}) {
-    if (this.isStreaming) {
-      throw new Error('Streaming is already in progress');
-    }
+  public async *stream(
+    options: DecodingOptions = {}
+  ): AsyncGenerator<{ committed: string; nonCommitted: string }> {
     this.validateOptions(options);
-    this.resetStreamState();
-
-    this.isStreaming = true;
-    while (this.isStreaming) {
-      if (
-        !this.readyToProcess ||
-        this.processor.audioBuffer.length < this.minAudioSamples
-      ) {
-        await new Promise((resolve) => setTimeout(resolve, 100));
+
+    const queue: { committed: string; nonCommitted: string }[] = [];
+    let waiter: (() => void) | null = null;
+    let finished = false;
+    let error: unknown;
+
+    const wake = () => {
+      waiter?.();
+      waiter = null;
+    };
+
+    (async () => {
+      try {
+        await this.nativeModule.stream(
+          (committed: string, nonCommitted: string, isDone: boolean) => {
+            queue.push({ committed, nonCommitted });
+            if (isDone) {
+              finished = true;
+            }
+            wake();
+          },
+          options.language || ''
+        );
+        finished = true;
+        wake();
+      } catch (e) {
+        error = e;
+        finished = true;
+        wake();
+      }
+    })();
+
+    while (true) {
+      if (queue.length > 0) {
+        yield queue.shift()!;
+        if (finished && queue.length === 0) {
+          return;
+        }
         continue;
       }
-
-      const { committed, nonCommitted } =
-        await this.processor.processIter(options);
-      yield { committed, nonCommitted };
-      this.readyToProcess = false;
+      if (error) throw error;
+      if (finished) return;
+      await new Promise<void>((r) => (waiter = r));
     }
-
-    const { committed } = await this.processor.finish();
-    yield { committed, nonCommitted: '' };
   }
 
-  public streamStop() {
-    this.isStreaming = false;
+  public async streamInsert(waveform: Float32Array | number[]): Promise<void> {
+    if (Array.isArray(waveform)) {
+      Logger.info(
+        'Passing waveform as number[] is deprecated, use Float32Array instead'
+      );
+      waveform = new Float32Array(waveform);
+    }
+    return this.nativeModule.streamInsert(waveform);
   }
 
-  public streamInsert(waveform: number[]) {
-    this.processor.insertAudioChunk(waveform);
-    this.readyToProcess = true;
+  public async streamStop(): Promise<void> {
+    return this.nativeModule.streamStop();
   }
 
   private validateOptions(options: DecodingOptions) {
@@ -89,10 +159,4 @@ export class SpeechToTextModule {
       throw new Error('Model is multilingual, provide a language');
     }
   }
-
-  private resetStreamState() {
-    this.isStreaming = false;
-    this.readyToProcess = false;
-    this.processor = new OnlineASRProcessor(this.asr);
-  }
 }

package/src/modules/natural_language_processing/TextEmbeddingsModule.ts

@@ -1,8 +1,8 @@
 import { ResourceSource } from '../../types/common';
 import { ResourceFetcher } from '../../utils/ResourceFetcher';
-import { BaseNonStaticModule } from '../BaseNonStaticModule';
+import { BaseModule } from '../BaseModule';
 
-export class TextEmbeddingsModule extends BaseNonStaticModule {
+export class TextEmbeddingsModule extends BaseModule {
   async load(
     model: { modelSource: ResourceSource; tokenizerSource: ResourceSource },
     onDownloadProgressCallback: (progress: number) => void = () => {}

package/src/native/RnExecutorchModules.ts

@@ -1,5 +1,4 @@
 import { Platform } from 'react-native';
-import { Spec as LLMInterface } from './NativeLLM';
 import { Spec as ETInstallerInterface } from './NativeETInstaller';
 
 const LINKING_ERROR =
@@ -23,8 +22,5 @@ function returnSpecOrThrowLinkingError(spec: any) {
 
 const ETInstallerNativeModule: ETInstallerInterface =
   returnSpecOrThrowLinkingError(require('./NativeETInstaller').default);
-const LLMNativeModule: LLMInterface = returnSpecOrThrowLinkingError(
-  require('./NativeLLM').default
-);
 
-export { LLMNativeModule, ETInstallerNativeModule };
+export { ETInstallerNativeModule };

package/src/types/stt.ts

@@ -1,17 +1,5 @@
 import { ResourceSource } from './common';
 
-export type WordTuple = [number, number, string];
-
-export interface WordObject {
-  start: number;
-  end: number;
-  word: string;
-}
-
-export interface Segment {
-  words: WordObject[];
-}
-
 // Languages supported by whisper (not whisper.en)
 export type SpeechToTextLanguage =
   | 'af'