react-native-executorch 0.3.2 → 0.4.1

package/src/hooks/natural_language_processing/useSpeechToText.ts

@@ -1,6 +1,8 @@
-import { useEffect, useState } from 'react';
+import { useEffect, useMemo, useState } from 'react';
 import { SpeechToTextController } from '../../controllers/SpeechToTextController';
 import { ResourceSource } from '../../types/common';
+import { STREAMING_ACTION } from '../../constants/sttDefaults';
+import { AvailableModels, SpeechToTextLanguage } from '../../types/stt';
 
 interface SpeechToTextModule {
   isReady: boolean;
@@ -10,8 +12,14 @@ interface SpeechToTextModule {
   configureStreaming: SpeechToTextController['configureStreaming'];
   error: Error | undefined;
   transcribe: (
-    input: number[]
+    input: number[],
+    audioLanguage?: SpeechToTextLanguage
   ) => ReturnType<SpeechToTextController['transcribe']>;
+  streamingTranscribe: (
+    streamAction: STREAMING_ACTION,
+    input?: number[],
+    audioLanguage?: SpeechToTextLanguage
+  ) => ReturnType<SpeechToTextController['streamingTranscribe']>;
 }
 
 export const useSpeechToText = ({
@@ -22,8 +30,9 @@ export const useSpeechToText = ({
   overlapSeconds,
   windowSize,
   streamingConfig,
+  preventLoad = false,
 }: {
-  modelName: 'moonshine' | 'whisper';
+  modelName: AvailableModels;
   encoderSource?: ResourceSource;
   decoderSource?: ResourceSource;
   tokenizerSource?: ResourceSource;
@@ -36,6 +45,7 @@ export const useSpeechToText = ({
   streamingConfig?: ConstructorParameters<
     typeof SpeechToTextController
   >['0']['streamingConfig'];
+  preventLoad?: boolean;
 }): SpeechToTextModule => {
   const [sequence, setSequence] = useState<string>('');
   const [isReady, setIsReady] = useState(false);
@@ -43,20 +53,22 @@ export const useSpeechToText = ({
   const [isGenerating, setIsGenerating] = useState(false);
   const [error, setError] = useState<Error | undefined>();
 
-  const [model, _] = useState(
+  const model = useMemo(
     () =>
       new SpeechToTextController({
         transcribeCallback: setSequence,
         isReadyCallback: setIsReady,
         isGeneratingCallback: setIsGenerating,
         onErrorCallback: setError,
-        modelDownloadProgessCallback: setDownloadProgress,
-        overlapSeconds: overlapSeconds,
-        windowSize: windowSize,
-        streamingConfig: streamingConfig,
-      })
+        modelDownloadProgressCallback: setDownloadProgress,
+      }),
+    []
   );
 
+  useEffect(() => {
+    model.configureStreaming(overlapSeconds, windowSize, streamingConfig);
+  }, [model, overlapSeconds, windowSize, streamingConfig]);
+
   useEffect(() => {
     const loadModel = async () => {
       await model.loadModel(
@@ -66,8 +78,17 @@ export const useSpeechToText = ({
         tokenizerSource
       );
     };
-    loadModel();
-  }, [model, modelName, encoderSource, decoderSource, tokenizerSource]);
+    if (!preventLoad) {
+      loadModel();
+    }
+  }, [
+    model,
+    modelName,
+    encoderSource,
+    decoderSource,
+    tokenizerSource,
+    preventLoad,
+  ]);
 
   return {
     isReady,
@@ -76,6 +97,12 @@ export const useSpeechToText = ({
     configureStreaming: model.configureStreaming,
     sequence,
     error,
-    transcribe: (waveform: number[]) => model.transcribe(waveform),
+    transcribe: (waveform: number[], audioLanguage?: SpeechToTextLanguage) =>
+      model.transcribe(waveform, audioLanguage),
+    streamingTranscribe: (
+      streamAction: STREAMING_ACTION,
+      waveform?: number[],
+      audioLanguage?: SpeechToTextLanguage
+    ) => model.streamingTranscribe(streamAction, waveform, audioLanguage),
   };
 };

package/src/hooks/natural_language_processing/useTextEmbeddings.ts

@@ -0,0 +1,18 @@
+import { TextEmbeddingsModule } from '../../modules/natural_language_processing/TextEmbeddingsModule';
+import { ResourceSource } from '../../types/common';
+import { useModule } from '../useModule';
+
+export const useTextEmbeddings = ({
+  modelSource,
+  tokenizerSource,
+  preventLoad = false,
+}: {
+  modelSource: ResourceSource;
+  tokenizerSource: ResourceSource;
+  preventLoad?: boolean;
+}) =>
+  useModule({
+    module: TextEmbeddingsModule,
+    loadArgs: [modelSource, tokenizerSource],
+    preventLoad,
+  });

package/src/hooks/natural_language_processing/useTokenizer.ts

@@ -0,0 +1,61 @@
+import { useEffect, useState } from 'react';
+import { TokenizerModule } from '../../modules/natural_language_processing/TokenizerModule';
+import { ResourceSource } from '../../types/common';
+import { ETError, getError } from '../../Error';
+
+export const useTokenizer = ({
+  tokenizerSource,
+  preventLoad = false,
+}: {
+  tokenizerSource: ResourceSource;
+  preventLoad?: boolean;
+}) => {
+  const [error, setError] = useState<null | string>(null);
+  const [isReady, setIsReady] = useState(false);
+  const [isGenerating, setIsGenerating] = useState(false);
+  const [downloadProgress, setDownloadProgress] = useState(0);
+
+  useEffect(() => {
+    const loadModule = async () => {
+      try {
+        setIsReady(false);
+        TokenizerModule.onDownloadProgress(setDownloadProgress);
+        await TokenizerModule.load(tokenizerSource);
+        setIsReady(true);
+      } catch (err) {
+        setError((err as Error).message);
+      }
+    };
+    if (!preventLoad) {
+      loadModule();
+    }
+  }, [tokenizerSource, preventLoad]);
+
+  const stateWrapper = <T extends (...args: any[]) => Promise<any>>(fn: T) => {
+    const boundFn = fn.bind(TokenizerModule);
+
+    return async (...args: Parameters<T>): Promise<ReturnType<T>> => {
+      if (!isReady) throw new Error(getError(ETError.ModuleNotLoaded));
+      if (isGenerating) throw new Error(getError(ETError.ModelGenerating));
+
+      setIsGenerating(true);
+      try {
+        return await boundFn(...args);
+      } finally {
+        setIsGenerating(false);
+      }
+    };
+  };
+
+  return {
+    error,
+    isReady,
+    isGenerating,
+    downloadProgress,
+    decode: stateWrapper(TokenizerModule.decode),
+    encode: stateWrapper(TokenizerModule.encode),
+    getVocabSize: stateWrapper(TokenizerModule.getVocabSize),
+    idToToken: stateWrapper(TokenizerModule.idToToken),
+    tokenToId: stateWrapper(TokenizerModule.tokenToId),
+  };
+};

package/src/hooks/useModule.ts

@@ -1,124 +1,64 @@
 import { useEffect, useState } from 'react';
-import { fetchResource } from '../utils/fetchResource';
 import { ETError, getError } from '../Error';
-import { ETInput, Module } from '../types/common';
-import { _ETModule } from '../native/RnExecutorchModules';
-import { getTypeIdentifier } from '../types/common';
 
-interface Props {
-  modelSource: string | number;
-  module: Module;
+interface Module {
+  load: (...args: any[]) => Promise<void>;
+  forward: (...input: any[]) => Promise<any>;
+  onDownloadProgress: (cb: (progress: number) => void) => void;
 }
 
-interface _Module {
-  error: string | null;
-  isReady: boolean;
-  isGenerating: boolean;
-  downloadProgress: number;
-  forwardETInput: (
-    input: ETInput[] | ETInput,
-    shape: number[][] | number[]
-  ) => ReturnType<_ETModule['forward']>;
-  forwardImage: (input: string) => Promise<any>;
-}
-
-export const useModule = ({ modelSource, module }: Props): _Module => {
+export const useModule = <
+  M extends Module,
+  LoadArgs extends Parameters<M['load']>,
+  ForwardArgs extends Parameters<M['forward']>,
+  ForwardReturn extends Awaited<ReturnType<M['forward']>>,
+>({
+  module,
+  loadArgs,
+  preventLoad = false,
+}: {
+  module: M;
+  loadArgs: LoadArgs;
+  preventLoad?: boolean;
+}) => {
   const [error, setError] = useState<null | string>(null);
   const [isReady, setIsReady] = useState(false);
   const [isGenerating, setIsGenerating] = useState(false);
   const [downloadProgress, setDownloadProgress] = useState(0);
 
   useEffect(() => {
-    const loadModel = async () => {
-      if (!modelSource) return;
-
+    const loadModule = async () => {
       try {
         setIsReady(false);
-        const fileUri = await fetchResource(modelSource, setDownloadProgress);
-        await module.loadModule(fileUri);
+        module.onDownloadProgress(setDownloadProgress);
+        await module.load(...loadArgs);
         setIsReady(true);
-      } catch (e) {
-        setError(getError(e));
+      } catch (err) {
+        setError((err as Error).message);
       }
     };
-
-    loadModel();
-  }, [modelSource, module]);
-
-  const forwardImage = async (input: string) => {
-    if (!isReady) {
-      throw new Error(getError(ETError.ModuleNotLoaded));
-    }
-    if (isGenerating) {
-      throw new Error(getError(ETError.ModelGenerating));
+    if (!preventLoad) {
+      loadModule();
     }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [...loadArgs, preventLoad]);
 
+  const forward = async (...input: ForwardArgs): Promise<ForwardReturn> => {
+    if (!isReady) throw new Error(getError(ETError.ModuleNotLoaded));
+    if (isGenerating) throw new Error(getError(ETError.ModelGenerating));
     try {
       setIsGenerating(true);
-      const output = await module.forward(input);
-      return output;
-    } catch (e) {
-      throw new Error(getError(e));
+      return await module.forward(...input);
     } finally {
       setIsGenerating(false);
     }
   };
 
-  const forwardETInput = async (
-    input: ETInput[] | ETInput,
-    shape: number[][] | number[]
-  ) => {
-    if (!isReady) {
-      throw new Error(getError(ETError.ModuleNotLoaded));
-    }
-    if (isGenerating) {
-      throw new Error(getError(ETError.ModelGenerating));
-    }
-
-    // Since the native module expects an array of inputs and an array of shapes,
-    // if the user provides a single ETInput, we want to "unsqueeze" the array so
-    // the data is properly processed on the native side
-    if (!Array.isArray(input)) {
-      input = [input];
-    }
-
-    if (!Array.isArray(shape[0])) {
-      shape = [shape] as number[][];
-    }
-
-    let inputTypeIdentifiers: any[] = [];
-    let modelInputs: any[] = [];
-
-    for (let idx = 0; idx < input.length; idx++) {
-      let currentInputTypeIdentifier = getTypeIdentifier(input[idx] as ETInput);
-      if (currentInputTypeIdentifier === -1) {
-        throw new Error(getError(ETError.InvalidArgument));
-      }
-      inputTypeIdentifiers.push(currentInputTypeIdentifier);
-      modelInputs.push([...(input[idx] as ETInput)]);
-    }
-
-    try {
-      setIsGenerating(true);
-      const output = await module.forward(
-        modelInputs,
-        shape,
-        inputTypeIdentifiers
-      );
-      setIsGenerating(false);
-      return output;
-    } catch (e) {
-      setIsGenerating(false);
-      throw new Error(getError(e));
-    }
-  };
-
   return {
     error,
     isReady,
     isGenerating,
     downloadProgress,
-    forwardETInput,
-    forwardImage,
+    forward,
   };
 };

package/src/index.tsx

@@ -1,12 +1,17 @@
+import { SpeechToTextLanguage } from './types/stt';
+
 // hooks
 export * from './hooks/computer_vision/useClassification';
 export * from './hooks/computer_vision/useObjectDetection';
 export * from './hooks/computer_vision/useStyleTransfer';
+export * from './hooks/computer_vision/useImageSegmentation';
 export * from './hooks/computer_vision/useOCR';
 export * from './hooks/computer_vision/useVerticalOCR';
 
 export * from './hooks/natural_language_processing/useLLM';
 export * from './hooks/natural_language_processing/useSpeechToText';
+export * from './hooks/natural_language_processing/useTextEmbeddings';
+export * from './hooks/natural_language_processing/useTokenizer';
 
 export * from './hooks/general/useExecutorchModule';
 
@@ -14,20 +19,29 @@ export * from './hooks/general/useExecutorchModule';
 export * from './modules/computer_vision/ClassificationModule';
 export * from './modules/computer_vision/ObjectDetectionModule';
 export * from './modules/computer_vision/StyleTransferModule';
+export * from './modules/computer_vision/ImageSegmentationModule';
 export * from './modules/computer_vision/OCRModule';
 export * from './modules/computer_vision/VerticalOCRModule';
 
 export * from './modules/natural_language_processing/LLMModule';
 export * from './modules/natural_language_processing/SpeechToTextModule';
+export * from './modules/natural_language_processing/TextEmbeddingsModule';
+export * from './modules/natural_language_processing/TokenizerModule';
 
 export * from './modules/general/ExecutorchModule';
 
 // utils
-export * from './utils/listDownloadedResources';
+export * from './utils/ResourceFetcher';
 
 // types
-export * from './types/object_detection';
+export * from './types/objectDetection';
 export * from './types/ocr';
+export * from './types/imageSegmentation';
+export * from './types/llm';
+export { SpeechToTextLanguage };
 
 // constants
 export * from './constants/modelUrls';
+export * from './constants/ocr/models';
+export * from './constants/llmDefaults';
+export { STREAMING_ACTION, MODES } from './constants/sttDefaults';

package/src/modules/BaseModule.ts

@@ -1,38 +1,28 @@
-import {
-  _StyleTransferModule,
-  _ObjectDetectionModule,
-  _ClassificationModule,
-  _ETModule,
-} from '../native/RnExecutorchModules';
-import { fetchResource } from '../utils/fetchResource';
-import { ResourceSource } from '../types/common';
+import { ResourceFetcher } from '../utils/ResourceFetcher';
 import { getError } from '../Error';
+import { ResourceSource } from '../types/common';
 
 export class BaseModule {
-  static module:
-    | _StyleTransferModule
-    | _ObjectDetectionModule
-    | _ClassificationModule
-    | _ETModule;
-
-  static onDownloadProgressCallback = (_downloadProgress: number) => {};
-
-  static async load(modelSource: ResourceSource) {
-    if (!modelSource) return;
+  protected static nativeModule: any;
+  static onDownloadProgressCallback: (downloadProgress: number) => void =
+    () => {};
 
+  static async load(...sources: ResourceSource[]): Promise<void> {
     try {
-      const fileUri = await fetchResource(
-        modelSource,
-        this.onDownloadProgressCallback
+      const paths = await ResourceFetcher.fetchMultipleResources(
+        this.onDownloadProgressCallback,
+        ...sources
       );
-      await this.module.loadModule(fileUri);
-    } catch (e) {
-      throw new Error(getError(e));
+      await this.nativeModule.loadModule(...paths);
+    } catch (error) {
+      throw new Error(getError(error));
     }
   }
 
-  static async forward(..._: any[]): Promise<any> {
-    throw new Error('The forward method is not implemented.');
+  protected static async forward(..._args: any[]): Promise<any> {
+    throw new Error(
+      'forward method is not implemented in the BaseModule class. Please implement it in the derived class.'
+    );
   }
 
   static onDownloadProgress(callback: (downloadProgress: number) => void) {

package/src/modules/computer_vision/ClassificationModule.ts

@@ -1,12 +1,17 @@
-import { BaseCVModule } from './BaseCVModule';
-import { _ClassificationModule } from '../../native/RnExecutorchModules';
+import { ClassificationNativeModule } from '../../native/RnExecutorchModules';
+import { ResourceSource } from '../../types/common';
+import { BaseModule } from '../BaseModule';
 
-export class ClassificationModule extends BaseCVModule {
-  static module = new _ClassificationModule();
+export class ClassificationModule extends BaseModule {
+  protected static override nativeModule = ClassificationNativeModule;
 
-  static async forward(input: string) {
-    return await (super.forward(input) as ReturnType<
-      _ClassificationModule['forward']
-    >);
+  static override async load(modelSource: ResourceSource) {
+    await super.load(modelSource);
+  }
+
+  static override async forward(
+    input: string
+  ): ReturnType<typeof ClassificationNativeModule.forward> {
+    return await this.nativeModule.forward(input);
   }
 }

package/src/modules/computer_vision/ImageSegmentationModule.ts

@@ -0,0 +1,39 @@
+import { BaseModule } from '../BaseModule';
+import { getError } from '../../Error';
+import { DeeplabLabel } from '../../types/imageSegmentation';
+import { ResourceSource } from '../../types/common';
+import { ImageSegmentationNativeModule } from '../../native/RnExecutorchModules';
+
+export class ImageSegmentationModule extends BaseModule {
+  protected static override nativeModule = ImageSegmentationNativeModule;
+
+  static override async load(modelSource: ResourceSource) {
+    return await super.load(modelSource);
+  }
+
+  static override async forward(
+    input: string,
+    classesOfInterest?: DeeplabLabel[],
+    resize?: boolean
+  ) {
+    try {
+      const stringDict = await (this.nativeModule.forward(
+        input,
+        (classesOfInterest || []).map((label) => DeeplabLabel[label]),
+        resize || false
+      ) as ReturnType<(typeof this.nativeModule)['forward']>);
+
+      let enumDict: { [key in DeeplabLabel]?: number[] } = {};
+
+      for (const key in stringDict) {
+        if (key in DeeplabLabel) {
+          const enumKey = DeeplabLabel[key as keyof typeof DeeplabLabel];
+          enumDict[enumKey] = stringDict[key];
+        }
+      }
+      return enumDict;
+    } catch (e) {
+      throw new Error(getError(e));
+    }
+  }
+}

package/src/modules/computer_vision/ObjectDetectionModule.ts

@@ -1,12 +1,17 @@
-import { BaseCVModule } from './BaseCVModule';
-import { _ObjectDetectionModule } from '../../native/RnExecutorchModules';
+import { ObjectDetectionNativeModule } from '../../native/RnExecutorchModules';
+import { ResourceSource } from '../../types/common';
+import { BaseModule } from '../BaseModule';
 
-export class ObjectDetectionModule extends BaseCVModule {
-  static module = new _ObjectDetectionModule();
+export class ObjectDetectionModule extends BaseModule {
+  protected static override nativeModule = ObjectDetectionNativeModule;
 
-  static async forward(input: string) {
-    return await (super.forward(input) as ReturnType<
-      _ObjectDetectionModule['forward']
-    >);
+  static override async load(modelSource: ResourceSource) {
+    return await super.load(modelSource);
+  }
+
+  static override async forward(
+    input: string
+  ): ReturnType<typeof this.nativeModule.forward> {
+    return await this.nativeModule.forward(input);
   }
 }

package/src/modules/computer_vision/StyleTransferModule.ts

@@ -1,12 +1,17 @@
-import { BaseCVModule } from './BaseCVModule';
-import { _StyleTransferModule } from '../../native/RnExecutorchModules';
+import { StyleTransferNativeModule } from '../../native/RnExecutorchModules';
+import { ResourceSource } from '../../types/common';
+import { BaseModule } from '../BaseModule';
 
-export class StyleTransferModule extends BaseCVModule {
-  static module = new _StyleTransferModule();
+export class StyleTransferModule extends BaseModule {
+  protected static override nativeModule = StyleTransferNativeModule;
 
-  static async forward(input: string) {
-    return await (super.forward(input) as ReturnType<
-      _StyleTransferModule['forward']
-    >);
+  static override async load(modelSource: ResourceSource) {
+    return await super.load(modelSource);
+  }
+
+  static override async forward(
+    input: string
+  ): ReturnType<typeof this.nativeModule.forward> {
+    return await this.nativeModule.forward(input);
   }
 }

package/src/modules/general/ExecutorchModule.ts

@@ -1,13 +1,18 @@
-import { BaseModule } from '../BaseModule';
 import { ETError, getError } from '../../Error';
-import { _ETModule } from '../../native/RnExecutorchModules';
+import { ETModuleNativeModule } from '../../native/RnExecutorchModules';
+import { ResourceSource } from '../../types/common';
 import { ETInput } from '../../types/common';
 import { getTypeIdentifier } from '../../types/common';
+import { BaseModule } from '../BaseModule';
 
 export class ExecutorchModule extends BaseModule {
-  static module = new _ETModule();
+  protected static override nativeModule = ETModuleNativeModule;
+
+  static override async load(modelSource: ResourceSource) {
+    return await super.load(modelSource);
+  }
 
-  static async forward(input: ETInput[] | ETInput, shape: number[][]) {
+  static override async forward(input: ETInput[] | ETInput, shape: number[][]) {
     if (!Array.isArray(input)) {
       input = [input];
     }
@@ -25,7 +30,7 @@ export class ExecutorchModule extends BaseModule {
     }
 
     try {
-      return await this.module.forward(
+      return await this.nativeModule.forward(
         modelInputs,
         shape,
         inputTypeIdentifiers
@@ -37,7 +42,7 @@ export class ExecutorchModule extends BaseModule {
 
   static async loadMethod(methodName: string) {
     try {
-      await this.module.loadMethod(methodName);
+      await this.nativeModule.loadMethod(methodName);
     } catch (e) {
       throw new Error(getError(e));
     }