react-native-executorch 0.4.7 → 0.5.0

package/src/hooks/computer_vision/useStyleTransfer.ts

@@ -1,15 +1,15 @@
 import { ResourceSource } from '../../types/common';
-import { useModule } from '../useModule';
+import { useNonStaticModule } from '../useNonStaticModule';
 import { StyleTransferModule } from '../../modules/computer_vision/StyleTransferModule';
 
 interface Props {
-  modelSource: ResourceSource;
+  model: { modelSource: ResourceSource };
   preventLoad?: boolean;
 }
 
-export const useStyleTransfer = ({ modelSource, preventLoad = false }: Props) =>
-  useModule({
+export const useStyleTransfer = ({ model, preventLoad = false }: Props) =>
+  useNonStaticModule({
     module: StyleTransferModule,
-    loadArgs: [modelSource],
-    preventLoad,
+    model,
+    preventLoad: preventLoad,
   });

package/src/hooks/computer_vision/useVerticalOCR.ts

@@ -12,21 +12,17 @@ interface OCRModule {
 }
 
 export const useVerticalOCR = ({
-  detectorSources,
-  recognizerSources,
-  language = 'en',
+  model,
   independentCharacters = false,
   preventLoad = false,
 }: {
-  detectorSources: {
+  model: {
     detectorLarge: ResourceSource;
     detectorNarrow: ResourceSource;
-  };
-  recognizerSources: {
     recognizerLarge: ResourceSource;
     recognizerSmall: ResourceSource;
+    language: OCRLanguage;
   };
-  language?: OCRLanguage;
   independentCharacters?: boolean;
   preventLoad?: boolean;
 }): OCRModule => {
@@ -35,10 +31,9 @@ export const useVerticalOCR = ({
   const [isGenerating, setIsGenerating] = useState(false);
   const [downloadProgress, setDownloadProgress] = useState(0);
 
-  const model = useMemo(
+  const controllerInstance = useMemo(
     () =>
       new VerticalOCRController({
-        modelDownloadProgressCallback: setDownloadProgress,
         isReadyCallback: setIsReady,
         isGeneratingCallback: setIsGenerating,
         errorCallback: setError,
@@ -47,27 +42,35 @@ export const useVerticalOCR = ({
   );
 
   useEffect(() => {
-    const loadModel = async () => {
-      await model.loadModel(
-        detectorSources,
-        recognizerSources,
-        language,
-        independentCharacters
+    if (preventLoad) return;
+
+    (async () => {
+      await controllerInstance.load(
+        {
+          detectorLarge: model.detectorLarge,
+          detectorNarrow: model.detectorNarrow,
+        },
+        {
+          recognizerLarge: model.recognizerLarge,
+          recognizerSmall: model.recognizerSmall,
+        },
+        model.language,
+        independentCharacters,
+        setDownloadProgress
       );
-    };
+    })();
 
-    if (!preventLoad) {
-      loadModel();
-    }
-    // eslint-disable-next-line react-hooks/exhaustive-deps
+    return () => {
+      controllerInstance.delete();
+    };
   }, [
-    model,
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-    JSON.stringify(detectorSources),
-    language,
+    controllerInstance,
+    model.detectorLarge,
+    model.detectorNarrow,
+    model.recognizerLarge,
+    model.recognizerSmall,
+    model.language,
     independentCharacters,
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-    JSON.stringify(recognizerSources),
     preventLoad,
   ]);
 
@@ -75,7 +78,7 @@ export const useVerticalOCR = ({
     error,
     isReady,
     isGenerating,
-    forward: model.forward,
+    forward: controllerInstance.forward,
     downloadProgress,
   };
 };

package/src/hooks/general/useExecutorchModule.ts

@@ -1,6 +1,6 @@
 import { ExecutorchModule } from '../../modules/general/ExecutorchModule';
 import { ResourceSource } from '../../types/common';
-import { useModule } from '../useModule';
+import { useNonStaticModule } from '../useNonStaticModule';
 
 interface Props {
   modelSource: ResourceSource;
@@ -11,8 +11,8 @@ export const useExecutorchModule = ({
   modelSource,
   preventLoad = false,
 }: Props) =>
-  useModule({
+  useNonStaticModule({
     module: ExecutorchModule,
-    loadArgs: [modelSource],
+    model: { modelSource },
     preventLoad,
   });

package/src/hooks/natural_language_processing/useLLM.ts

@@ -13,14 +13,14 @@ import { LLMController } from '../../controllers/LLMController';
 Hook version of LLMModule
 */
 export const useLLM = ({
-  modelSource,
-  tokenizerSource,
-  tokenizerConfigSource,
+  model,
   preventLoad = false,
 }: {
-  modelSource: ResourceSource;
-  tokenizerSource: ResourceSource;
-  tokenizerConfigSource: ResourceSource;
+  model: {
+    modelSource: ResourceSource;
+    tokenizerSource: ResourceSource;
+    tokenizerConfigSource: ResourceSource;
+  };
   preventLoad?: boolean;
 }): LLMType => {
   const [token, setToken] = useState<string>('');
@@ -36,14 +36,13 @@ export const useLLM = ({
     setResponse((prevResponse) => prevResponse + newToken);
   }, []);
 
-  const model = useMemo(
+  const controllerInstance = useMemo(
     () =>
       new LLMController({
         tokenCallback: tokenCallback,
         messageHistoryCallback: setMessageHistory,
         isReadyCallback: setIsReady,
         isGeneratingCallback: setIsGenerating,
-        onDownloadProgressCallback: setDownloadProgress,
       }),
     [tokenCallback]
   );
@@ -52,24 +51,31 @@ export const useLLM = ({
     setDownloadProgress(0);
     setError(null);
 
-    if (!preventLoad) {
+    if (preventLoad) return;
+
+    (async () => {
       try {
-        (async () => {
-          await model.load({
-            modelSource,
-            tokenizerSource,
-            tokenizerConfigSource,
-          });
-        })();
+        await controllerInstance.load({
+          modelSource: model.modelSource,
+          tokenizerSource: model.tokenizerSource,
+          tokenizerConfigSource: model.tokenizerConfigSource,
+          onDownloadProgressCallback: setDownloadProgress,
+        });
       } catch (e) {
         setError(e);
       }
-    }
+    })();
 
     return () => {
-      model.delete();
+      controllerInstance.delete();
     };
-  }, [modelSource, tokenizerSource, tokenizerConfigSource, preventLoad, model]);
+  }, [
+    controllerInstance,
+    model.modelSource,
+    model.tokenizerSource,
+    model.tokenizerConfigSource,
+    preventLoad,
+  ]);
 
   // memoization of returned functions
   const configure = useCallback(
@@ -79,31 +85,35 @@ export const useLLM = ({
     }: {
       chatConfig?: Partial<ChatConfig>;
       toolsConfig?: ToolsConfig;
-    }) => model.configure({ chatConfig, toolsConfig }),
-    [model]
+    }) => controllerInstance.configure({ chatConfig, toolsConfig }),
+    [controllerInstance]
   );
 
   const generate = useCallback(
     (messages: Message[], tools?: LLMTool[]) => {
       setResponse('');
-      return model.generate(messages, tools);
+      return controllerInstance.generate(messages, tools);
     },
-    [model]
+    [controllerInstance]
   );
 
   const sendMessage = useCallback(
     (message: string) => {
       setResponse('');
-      return model.sendMessage(message);
+      return controllerInstance.sendMessage(message);
     },
-    [model]
+    [controllerInstance]
   );
 
   const deleteMessage = useCallback(
-    (index: number) => model.deleteMessage(index),
-    [model]
+    (index: number) => controllerInstance.deleteMessage(index),
+    [controllerInstance]
+  );
+
+  const interrupt = useCallback(
+    () => controllerInstance.interrupt(),
+    [controllerInstance]
   );
-  const interrupt = useCallback(() => model.interrupt(), [model]);
 
   return {
     messageHistory,

package/src/hooks/natural_language_processing/useSpeechToText.ts

@@ -1,108 +1,111 @@
-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;
-  isGenerating: boolean;
-  sequence: string;
-  downloadProgress: number;
-  configureStreaming: SpeechToTextController['configureStreaming'];
-  error: Error | undefined;
-  transcribe: (
-    input: number[],
-    audioLanguage?: SpeechToTextLanguage
-  ) => ReturnType<SpeechToTextController['transcribe']>;
-  streamingTranscribe: (
-    streamAction: STREAMING_ACTION,
-    input?: number[],
-    audioLanguage?: SpeechToTextLanguage
-  ) => ReturnType<SpeechToTextController['streamingTranscribe']>;
-}
+import { useEffect, useCallback, useState } from 'react';
+import { ETError, getError } from '../../Error';
+import { SpeechToTextModule } from '../../modules/natural_language_processing/SpeechToTextModule';
+import { SpeechToTextModelConfig } from '../../types/stt';
 
 export const useSpeechToText = ({
-  modelName,
-  encoderSource,
-  decoderSource,
-  tokenizerSource,
-  overlapSeconds,
-  windowSize,
-  streamingConfig,
+  model,
   preventLoad = false,
 }: {
-  modelName: AvailableModels;
-  encoderSource?: ResourceSource;
-  decoderSource?: ResourceSource;
-  tokenizerSource?: ResourceSource;
-  overlapSeconds?: ConstructorParameters<
-    typeof SpeechToTextController
-  >['0']['overlapSeconds'];
-  windowSize?: ConstructorParameters<
-    typeof SpeechToTextController
-  >['0']['windowSize'];
-  streamingConfig?: ConstructorParameters<
-    typeof SpeechToTextController
-  >['0']['streamingConfig'];
+  model: SpeechToTextModelConfig;
   preventLoad?: boolean;
-}): SpeechToTextModule => {
-  const [sequence, setSequence] = useState<string>('');
+}) => {
+  const [error, setError] = useState<null | string>(null);
   const [isReady, setIsReady] = useState(false);
-  const [downloadProgress, setDownloadProgress] = useState(0);
   const [isGenerating, setIsGenerating] = useState(false);
-  const [error, setError] = useState<Error | undefined>();
-
-  const model = useMemo(
-    () =>
-      new SpeechToTextController({
-        transcribeCallback: setSequence,
-        isReadyCallback: setIsReady,
-        isGeneratingCallback: setIsGenerating,
-        onErrorCallback: setError,
-        modelDownloadProgressCallback: setDownloadProgress,
-      }),
-    []
-  );
+  const [downloadProgress, setDownloadProgress] = useState(0);
 
-  useEffect(() => {
-    model.configureStreaming(overlapSeconds, windowSize, streamingConfig);
-  }, [model, overlapSeconds, windowSize, streamingConfig]);
+  const [modelInstance] = useState(() => new SpeechToTextModule());
+  const [committedTranscription, setCommittedTranscription] = useState('');
+  const [nonCommittedTranscription, setNonCommittedTranscription] =
+    useState('');
 
   useEffect(() => {
-    const loadModel = async () => {
-      await model.loadModel(
-        modelName,
-        encoderSource,
-        decoderSource,
-        tokenizerSource
-      );
-    };
-    if (!preventLoad) {
-      loadModel();
-    }
+    if (preventLoad) return;
+    (async () => {
+      setDownloadProgress(0);
+      setError(null);
+      try {
+        setIsReady(false);
+        await modelInstance.load(
+          {
+            isMultilingual: model.isMultilingual,
+            encoderSource: model.encoderSource,
+            decoderSource: model.decoderSource,
+            tokenizerSource: model.tokenizerSource,
+          },
+          setDownloadProgress
+        );
+        setIsReady(true);
+      } catch (err) {
+        setError((err as Error).message);
+      }
+    })();
   }, [
-    model,
-    modelName,
-    encoderSource,
-    decoderSource,
-    tokenizerSource,
+    modelInstance,
+    model.isMultilingual,
+    model.encoderSource,
+    model.decoderSource,
+    model.tokenizerSource,
     preventLoad,
   ]);
 
+  const stateWrapper = useCallback(
+    <T extends (...args: any[]) => Promise<any>>(fn: T) =>
+      async (...args: Parameters<T>): Promise<Awaited<ReturnType<T>>> => {
+        if (!isReady) throw new Error(getError(ETError.ModuleNotLoaded));
+        if (isGenerating) throw new Error(getError(ETError.ModelGenerating));
+        setIsGenerating(true);
+        try {
+          return await fn.apply(modelInstance, args);
+        } finally {
+          setIsGenerating(false);
+        }
+      },
+    [isReady, isGenerating, modelInstance]
+  );
+
+  const stream = useCallback(async () => {
+    if (!isReady) throw new Error(getError(ETError.ModuleNotLoaded));
+    if (isGenerating) throw new Error(getError(ETError.ModelGenerating));
+    setIsGenerating(true);
+    setCommittedTranscription('');
+    setNonCommittedTranscription('');
+    let transcription = '';
+    try {
+      for await (const { committed, nonCommitted } of modelInstance.stream()) {
+        setCommittedTranscription((prev) => prev + committed);
+        setNonCommittedTranscription(nonCommitted);
+        transcription += committed;
+      }
+    } finally {
+      setIsGenerating(false);
+    }
+    return transcription;
+  }, [isReady, isGenerating, modelInstance]);
+
+  const wrapper = useCallback(
+    <T extends (...args: any[]) => any>(fn: T) => {
+      return (...args: Parameters<T>): ReturnType<T> => {
+        if (!isReady) throw new Error(getError(ETError.ModuleNotLoaded));
+        return fn.apply(modelInstance, args);
+      };
+    },
+    [isReady, modelInstance]
+  );
+
   return {
+    error,
     isReady,
     isGenerating,
     downloadProgress,
-    configureStreaming: model.configureStreaming,
-    sequence,
-    error,
-    transcribe: (waveform: number[], audioLanguage?: SpeechToTextLanguage) =>
-      model.transcribe(waveform, audioLanguage),
-    streamingTranscribe: (
-      streamAction: STREAMING_ACTION,
-      waveform?: number[],
-      audioLanguage?: SpeechToTextLanguage
-    ) => model.streamingTranscribe(streamAction, waveform, audioLanguage),
+    committedTranscription,
+    nonCommittedTranscription,
+    encode: stateWrapper(SpeechToTextModule.prototype.encode),
+    decode: stateWrapper(SpeechToTextModule.prototype.decode),
+    transcribe: stateWrapper(SpeechToTextModule.prototype.transcribe),
+    stream,
+    streamStop: wrapper(SpeechToTextModule.prototype.streamStop),
+    streamInsert: wrapper(SpeechToTextModule.prototype.streamInsert),
   };
 };

package/src/hooks/natural_language_processing/useTextEmbeddings.ts

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

package/src/hooks/natural_language_processing/useTokenizer.ts

@@ -1,46 +1,46 @@
-import { useEffect, useState } from 'react';
+import { useEffect, useMemo, 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,
+  tokenizer,
   preventLoad = false,
 }: {
-  tokenizerSource: ResourceSource;
+  tokenizer: { 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);
+  const _tokenizer = useMemo(() => new TokenizerModule(), []);
 
   useEffect(() => {
-    const loadModule = async () => {
+    if (preventLoad) return;
+    (async () => {
+      setDownloadProgress(0);
+      setError(null);
       try {
         setIsReady(false);
-        TokenizerModule.onDownloadProgress(setDownloadProgress);
-        await TokenizerModule.load(tokenizerSource);
+        await _tokenizer.load(
+          { tokenizerSource: tokenizer.tokenizerSource },
+          setDownloadProgress
+        );
         setIsReady(true);
       } catch (err) {
         setError((err as Error).message);
       }
-    };
-    if (!preventLoad) {
-      loadModule();
-    }
-  }, [tokenizerSource, preventLoad]);
+    })();
+  }, [_tokenizer, tokenizer.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>> => {
+    return (...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);
+        setIsGenerating(true);
+        return fn.apply(_tokenizer, args);
       } finally {
         setIsGenerating(false);
       }
@@ -52,10 +52,10 @@ export const useTokenizer = ({
     isReady,
     isGenerating,
     downloadProgress,
-    decode: stateWrapper(TokenizerModule.decode),
-    encode: stateWrapper(TokenizerModule.encode),
-    getVocabSize: stateWrapper(TokenizerModule.getVocabSize),
-    idToToken: stateWrapper(TokenizerModule.idToToken),
-    tokenToId: stateWrapper(TokenizerModule.tokenToId),
+    decode: stateWrapper(TokenizerModule.prototype.decode),
+    encode: stateWrapper(TokenizerModule.prototype.encode),
+    getVocabSize: stateWrapper(TokenizerModule.prototype.getVocabSize),
+    idToToken: stateWrapper(TokenizerModule.prototype.idToToken),
+    tokenToId: stateWrapper(TokenizerModule.prototype.tokenToId),
   };
 };

package/src/hooks/useNonStaticModule.ts

@@ -0,0 +1,74 @@
+import { useEffect, useState, useMemo } from 'react';
+import { ETError, getError } from '../Error';
+
+interface Module {
+  load: (...args: any[]) => Promise<void>;
+  forward: (...args: any[]) => Promise<any>;
+  delete: () => void;
+}
+
+interface ModuleConstructor<M extends Module> {
+  new (): M;
+}
+
+export const useNonStaticModule = <
+  M extends Module,
+  LoadArgs extends Parameters<M['load']>,
+  ForwardArgs extends Parameters<M['forward']>,
+  ForwardReturn extends Awaited<ReturnType<M['forward']>>,
+>({
+  module,
+  model,
+  preventLoad = false,
+}: {
+  module: ModuleConstructor<M>;
+  model: LoadArgs[0];
+  preventLoad?: boolean;
+}) => {
+  const [error, setError] = useState<null | string>(null);
+  const [isReady, setIsReady] = useState(false);
+  const [isGenerating, setIsGenerating] = useState(false);
+  const [downloadProgress, setDownloadProgress] = useState(0);
+  const moduleInstance = useMemo(() => new module(), [module]);
+
+  useEffect(() => {
+    if (preventLoad) return;
+
+    (async () => {
+      setDownloadProgress(0);
+      setError(null);
+      try {
+        setIsReady(false);
+        await moduleInstance.load(model, setDownloadProgress);
+        setIsReady(true);
+      } catch (err) {
+        setError((err as Error).message);
+      }
+    })();
+
+    return () => {
+      moduleInstance.delete();
+    };
+
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [moduleInstance, ...Object.values(model), 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);
+      return await moduleInstance.forward(...input);
+    } finally {
+      setIsGenerating(false);
+    }
+  };
+
+  return {
+    error,
+    isReady,
+    isGenerating,
+    downloadProgress,
+    forward,
+  };
+};