react-native-executorch 0.3.2 → 0.4.1

package/src/controllers/VerticalOCRController.ts

@@ -1,15 +1,12 @@
 import { symbols } from '../constants/ocr/symbols';
 import { ETError, getError } from '../Error';
-import { _VerticalOCRModule } from '../native/RnExecutorchModules';
+import { VerticalOCRNativeModule } from '../native/RnExecutorchModules';
 import { ResourceSource } from '../types/common';
 import { OCRLanguage } from '../types/ocr';
-import {
-  fetchResource,
-  calculateDownloadProgres,
-} from '../utils/fetchResource';
+import { ResourceFetcher } from '../utils/ResourceFetcher';
 
 export class VerticalOCRController {
-  private nativeModule: _VerticalOCRModule;
+  private ocrNativeModule: typeof VerticalOCRNativeModule;
   public isReady: boolean = false;
   public isGenerating: boolean = false;
   public error: string | null = null;
@@ -24,7 +21,7 @@ export class VerticalOCRController {
     isGeneratingCallback = (_isGenerating: boolean) => {},
     errorCallback = (_error: string) => {},
   }) {
-    this.nativeModule = new _VerticalOCRModule();
+    this.ocrNativeModule = VerticalOCRNativeModule;
     this.modelDownloadProgressCallback = modelDownloadProgressCallback;
     this.isReadyCallback = isReadyCallback;
     this.isGeneratingCallback = isGeneratingCallback;
@@ -57,31 +54,19 @@ export class VerticalOCRController {
       this.isReady = false;
       this.isReadyCallback(this.isReady);
 
-      const recognizerPath = independentCharacters
-        ? await fetchResource(
-            recognizerSources.recognizerSmall,
-            calculateDownloadProgres(3, 0, this.modelDownloadProgressCallback)
-          )
-        : await fetchResource(
-            recognizerSources.recognizerLarge,
-            calculateDownloadProgres(3, 0, this.modelDownloadProgressCallback)
-          );
-
-      const detectorPaths = {
-        detectorLarge: await fetchResource(
-          detectorSources.detectorLarge,
-          calculateDownloadProgres(3, 1, this.modelDownloadProgressCallback)
-        ),
-        detectorNarrow: await fetchResource(
-          detectorSources.detectorNarrow,
-          calculateDownloadProgres(3, 2, this.modelDownloadProgressCallback)
-        ),
-      };
+      const paths = await ResourceFetcher.fetchMultipleResources(
+        this.modelDownloadProgressCallback,
+        detectorSources.detectorLarge,
+        detectorSources.detectorNarrow,
+        independentCharacters
+          ? recognizerSources.recognizerSmall
+          : recognizerSources.recognizerLarge
+      );
 
-      await this.nativeModule.loadModule(
-        detectorPaths.detectorLarge,
-        detectorPaths.detectorNarrow,
-        recognizerPath,
+      await this.ocrNativeModule.loadModule(
+        paths[0]!,
+        paths[1]!,
+        paths[2]!,
         symbols[language],
         independentCharacters
       );
@@ -108,7 +93,7 @@ export class VerticalOCRController {
     try {
       this.isGenerating = true;
       this.isGeneratingCallback(this.isGenerating);
-      return await this.nativeModule.forward(input);
+      return await this.ocrNativeModule.forward(input);
     } catch (e) {
       throw new Error(getError(e));
     } finally {

package/src/hooks/computer_vision/useClassification.ts

@@ -1,31 +1,16 @@
-import { useState } from 'react';
-import { _ClassificationModule } from '../../native/RnExecutorchModules';
+import { ClassificationModule } from '../../modules/computer_vision/ClassificationModule';
+import { ResourceSource } from '../../types/common';
 import { useModule } from '../useModule';
 
-interface Props {
-  modelSource: string | number;
-}
-
 export const useClassification = ({
   modelSource,
-}: Props): {
-  error: string | null;
-  isReady: boolean;
-  isGenerating: boolean;
-  downloadProgress: number;
-  forward: (input: string) => Promise<{ [category: string]: number }>;
-} => {
-  const [module, _] = useState(() => new _ClassificationModule());
-  const {
-    error,
-    isReady,
-    isGenerating,
-    downloadProgress,
-    forwardImage: forward,
-  } = useModule({
-    modelSource,
-    module,
+  preventLoad = false,
+}: {
+  modelSource: ResourceSource;
+  preventLoad?: boolean;
+}) =>
+  useModule({
+    module: ClassificationModule,
+    loadArgs: [modelSource],
+    preventLoad,
   });
-
-  return { error, isReady, isGenerating, downloadProgress, forward };
-};

package/src/hooks/computer_vision/useImageSegmentation.ts

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

package/src/hooks/computer_vision/useOCR.ts

@@ -1,4 +1,4 @@
-import { useEffect, useState } from 'react';
+import { useEffect, useMemo, useState } from 'react';
 import { ResourceSource } from '../../types/common';
 import { OCRDetection, OCRLanguage } from '../../types/ocr';
 import { OCRController } from '../../controllers/OCRController';
@@ -15,6 +15,7 @@ export const useOCR = ({
   detectorSource,
   recognizerSources,
   language = 'en',
+  preventLoad = false,
 }: {
   detectorSource: ResourceSource;
   recognizerSources: {
@@ -23,20 +24,22 @@ export const useOCR = ({
     recognizerSmall: ResourceSource;
   };
   language?: OCRLanguage;
+  preventLoad?: boolean;
 }): OCRModule => {
   const [error, setError] = useState<string | null>(null);
   const [isReady, setIsReady] = useState(false);
   const [isGenerating, setIsGenerating] = useState(false);
   const [downloadProgress, setDownloadProgress] = useState(0);
 
-  const [model, _] = useState(
+  const model = useMemo(
     () =>
       new OCRController({
         modelDownloadProgressCallback: setDownloadProgress,
         isReadyCallback: setIsReady,
         isGeneratingCallback: setIsGenerating,
         errorCallback: setError,
-      })
+      }),
+    []
   );
 
   useEffect(() => {
@@ -44,9 +47,18 @@ export const useOCR = ({
       await model.loadModel(detectorSource, recognizerSources, language);
     };
 
-    loadModel();
+    if (!preventLoad) {
+      loadModel();
+    }
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [detectorSource, language, JSON.stringify(recognizerSources)]);
+  }, [
+    model,
+    detectorSource,
+    language,
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+    JSON.stringify(recognizerSources),
+    preventLoad,
+  ]);
 
   return {
     error,

package/src/hooks/computer_vision/useObjectDetection.ts

@@ -1,32 +1,18 @@
-import { useState } from 'react';
-import { _ObjectDetectionModule } from '../../native/RnExecutorchModules';
+import { ResourceSource } from '../../types/common';
 import { useModule } from '../useModule';
-import { Detection } from '../../types/object_detection';
+import { ObjectDetectionModule } from '../../modules/computer_vision/ObjectDetectionModule';
 
 interface Props {
-  modelSource: string | number;
+  modelSource: ResourceSource;
+  preventLoad?: boolean;
 }
 
 export const useObjectDetection = ({
   modelSource,
-}: Props): {
-  error: string | null;
-  isReady: boolean;
-  isGenerating: boolean;
-  downloadProgress: number;
-  forward: (input: string) => Promise<Detection[]>;
-} => {
-  const [module, _] = useState(() => new _ObjectDetectionModule());
-  const {
-    error,
-    isReady,
-    isGenerating,
-    downloadProgress,
-    forwardImage: forward,
-  } = useModule({
-    modelSource,
-    module,
+  preventLoad = false,
+}: Props) =>
+  useModule({
+    module: ObjectDetectionModule,
+    loadArgs: [modelSource],
+    preventLoad,
   });
-
-  return { error, isReady, isGenerating, downloadProgress, forward };
-};

package/src/hooks/computer_vision/useStyleTransfer.ts

@@ -1,31 +1,15 @@
-import { useState } from 'react';
-import { _StyleTransferModule } from '../../native/RnExecutorchModules';
+import { ResourceSource } from '../../types/common';
 import { useModule } from '../useModule';
+import { StyleTransferModule } from '../../modules/computer_vision/StyleTransferModule';
 
 interface Props {
-  modelSource: string | number;
+  modelSource: ResourceSource;
+  preventLoad?: boolean;
 }
 
-export const useStyleTransfer = ({
-  modelSource,
-}: Props): {
-  error: string | null;
-  isReady: boolean;
-  isGenerating: boolean;
-  downloadProgress: number;
-  forward: (input: string) => Promise<string>;
-} => {
-  const [module, _] = useState(() => new _StyleTransferModule());
-  const {
-    error,
-    isReady,
-    isGenerating,
-    downloadProgress,
-    forwardImage: forward,
-  } = useModule({
-    modelSource,
-    module,
+export const useStyleTransfer = ({ modelSource, preventLoad = false }: Props) =>
+  useModule({
+    module: StyleTransferModule,
+    loadArgs: [modelSource],
+    preventLoad,
   });
-
-  return { error, isReady, isGenerating, downloadProgress, forward };
-};

package/src/hooks/computer_vision/useVerticalOCR.ts

@@ -1,4 +1,4 @@
-import { useEffect, useState } from 'react';
+import { useEffect, useMemo, useState } from 'react';
 import { ResourceSource } from '../../types/common';
 import { OCRDetection, OCRLanguage } from '../../types/ocr';
 import { VerticalOCRController } from '../../controllers/VerticalOCRController';
@@ -16,6 +16,7 @@ export const useVerticalOCR = ({
   recognizerSources,
   language = 'en',
   independentCharacters = false,
+  preventLoad = false,
 }: {
   detectorSources: {
     detectorLarge: ResourceSource;
@@ -27,20 +28,22 @@ export const useVerticalOCR = ({
   };
   language?: OCRLanguage;
   independentCharacters?: boolean;
+  preventLoad?: boolean;
 }): OCRModule => {
   const [error, setError] = useState<string | null>(null);
   const [isReady, setIsReady] = useState(false);
   const [isGenerating, setIsGenerating] = useState(false);
   const [downloadProgress, setDownloadProgress] = useState(0);
 
-  const [model, _] = useState(
+  const model = useMemo(
     () =>
       new VerticalOCRController({
         modelDownloadProgressCallback: setDownloadProgress,
         isReadyCallback: setIsReady,
         isGeneratingCallback: setIsGenerating,
         errorCallback: setError,
-      })
+      }),
+    []
   );
 
   useEffect(() => {
@@ -53,15 +56,19 @@ export const useVerticalOCR = ({
       );
     };
 
-    loadModel();
+    if (!preventLoad) {
+      loadModel();
+    }
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [
+    model,
     // eslint-disable-next-line react-hooks/exhaustive-deps
     JSON.stringify(detectorSources),
     language,
     independentCharacters,
     // eslint-disable-next-line react-hooks/exhaustive-deps
     JSON.stringify(recognizerSources),
+    preventLoad,
   ]);
 
   return {

package/src/hooks/general/useExecutorchModule.ts

@@ -1,58 +1,18 @@
-import { useState } from 'react';
-import { _ETModule } from '../../native/RnExecutorchModules';
+import { ExecutorchModule } from '../../modules/general/ExecutorchModule';
+import { ResourceSource } from '../../types/common';
 import { useModule } from '../useModule';
-import { ETInput } from '../../types/common';
-import { getError } from '../../Error';
 
 interface Props {
-  modelSource: string | number;
+  modelSource: ResourceSource;
+  preventLoad?: boolean;
 }
 
 export const useExecutorchModule = ({
   modelSource,
-}: Props): {
-  error: string | null;
-  isReady: boolean;
-  isGenerating: boolean;
-  downloadProgress: number;
-  forward: (
-    input: ETInput | ETInput[],
-    shape: number[] | number[][]
-  ) => Promise<number[][]>;
-  loadMethod: (methodName: string) => Promise<void>;
-  loadForward: () => Promise<void>;
-} => {
-  const [module] = useState(() => new _ETModule());
-  const {
-    error,
-    isReady,
-    isGenerating,
-    downloadProgress,
-    forwardETInput: forward,
-  } = useModule({
-    modelSource,
-    module,
+  preventLoad = false,
+}: Props) =>
+  useModule({
+    module: ExecutorchModule,
+    loadArgs: [modelSource],
+    preventLoad,
   });
-
-  const loadMethod = async (methodName: string) => {
-    try {
-      await module.loadMethod(methodName);
-    } catch (e) {
-      throw new Error(getError(e));
-    }
-  };
-
-  const loadForward = async () => {
-    await loadMethod('forward');
-  };
-
-  return {
-    error,
-    isReady,
-    isGenerating,
-    downloadProgress,
-    forward,
-    loadMethod,
-    loadForward,
-  };
-};

package/src/hooks/natural_language_processing/useLLM.ts

@@ -1,126 +1,109 @@
-import { useCallback, useEffect, useRef, useState } from 'react';
-import { EventSubscription } from 'react-native';
-import { LLM } from '../../native/RnExecutorchModules';
-import { fetchResource } from '../../utils/fetchResource';
-import { ResourceSource, Model, MessageType } from '../../types/common';
+import { useCallback, useEffect, useMemo, useState } from 'react';
+import { ResourceSource } from '../../types/common';
 import {
-  DEFAULT_CONTEXT_WINDOW_LENGTH,
-  DEFAULT_MESSAGE_HISTORY,
-  DEFAULT_SYSTEM_PROMPT,
-  EOT_TOKEN,
-} from '../../constants/llamaDefaults';
-
-const interrupt = () => {
-  LLM.interrupt();
-};
+  ChatConfig,
+  LLMTool,
+  LLMType,
+  Message,
+  ToolsConfig,
+} from '../../types/llm';
+import { LLMController } from '../../controllers/LLMController';
 
+/*
+Hook version of LLMController
+*/
 export const useLLM = ({
   modelSource,
   tokenizerSource,
-  systemPrompt = DEFAULT_SYSTEM_PROMPT,
-  messageHistory = DEFAULT_MESSAGE_HISTORY,
-  contextWindowLength = DEFAULT_CONTEXT_WINDOW_LENGTH,
+  tokenizerConfigSource,
+  preventLoad = false,
 }: {
   modelSource: ResourceSource;
   tokenizerSource: ResourceSource;
-  systemPrompt?: string;
-  messageHistory?: MessageType[];
-  contextWindowLength?: number;
-}): Model => {
-  const [error, setError] = useState<string | null>(null);
+  tokenizerConfigSource: ResourceSource;
+  preventLoad?: boolean;
+}): LLMType => {
+  const [response, setResponse] = useState('');
+  const [messageHistory, setMessageHistory] = useState<Message[]>([]);
   const [isReady, setIsReady] = useState(false);
   const [isGenerating, setIsGenerating] = useState(false);
-  const [response, setResponse] = useState('');
   const [downloadProgress, setDownloadProgress] = useState(0);
-  const tokenGeneratedListener = useRef<null | EventSubscription>(null);
-
-  useEffect(() => {
-    const loadModel = async () => {
-      try {
-        setIsReady(false);
+  const [error, setError] = useState<any>(null);
 
-        const tokenizerFileUri = await fetchResource(tokenizerSource);
-        const modelFileUri = await fetchResource(
-          modelSource,
-          setDownloadProgress
-        );
-
-        await LLM.loadLLM(
-          modelFileUri,
-          tokenizerFileUri,
-          systemPrompt,
-          messageHistory,
-          contextWindowLength
-        );
+  const model = useMemo(
+    () =>
+      new LLMController({
+        responseCallback: setResponse,
+        messageHistoryCallback: setMessageHistory,
+        isReadyCallback: setIsReady,
+        isGeneratingCallback: setIsGenerating,
+        onDownloadProgressCallback: setDownloadProgress,
+      }),
+    []
+  );
 
-        setIsReady(true);
+  useEffect(() => {
+    setDownloadProgress(0);
+    setError(null);
 
-        tokenGeneratedListener.current = LLM.onToken(
-          (data: string | undefined) => {
-            if (!data) {
-              return;
-            }
-            if (data !== EOT_TOKEN) {
-              setResponse((prevResponse) => prevResponse + data);
-            } else {
-              setIsGenerating(false);
-            }
-          }
-        );
-      } catch (err) {
-        const message = (err as Error).message;
-        setIsReady(false);
-        setError(message);
-      } finally {
-        setDownloadProgress(0);
+    if (!preventLoad) {
+      try {
+        (async () => {
+          await model.load({
+            modelSource,
+            tokenizerSource,
+            tokenizerConfigSource,
+          });
+        })();
+      } catch (e) {
+        setError(e);
       }
-    };
-
-    loadModel();
+    }
 
     return () => {
-      tokenGeneratedListener.current?.remove();
-      tokenGeneratedListener.current = null;
-      LLM.deleteModule();
+      model.delete();
     };
-  }, [
-    modelSource,
-    tokenizerSource,
-    systemPrompt,
-    messageHistory,
-    contextWindowLength,
-  ]);
+  }, [modelSource, tokenizerSource, tokenizerConfigSource, preventLoad, model]);
+
+  // memoization of returned functions
+  const configure = useCallback(
+    ({
+      chatConfig,
+      toolsConfig,
+    }: {
+      chatConfig?: Partial<ChatConfig>;
+      toolsConfig?: ToolsConfig;
+    }) => model.configure({ chatConfig, toolsConfig }),
+    [model]
+  );
 
   const generate = useCallback(
-    async (input: string): Promise<void> => {
-      if (!isReady) {
-        throw new Error('Model is still loading');
-      }
-      if (error) {
-        throw new Error(error);
-      }
+    (messages: Message[], tools?: LLMTool[]) => model.generate(messages, tools),
+    [model]
+  );
 
-      try {
-        setResponse('');
-        setIsGenerating(true);
-        await LLM.runInference(input);
-      } catch (err) {
-        setIsGenerating(false);
-        throw new Error((err as Error).message);
-      }
-    },
-    [isReady, error]
+  const sendMessage = useCallback(
+    (message: string) => model.sendMessage(message),
+    [model]
   );
 
+  const deleteMessage = useCallback(
+    (index: number) => model.deleteMessage(index),
+    [model]
+  );
+  const interrupt = useCallback(() => model.interrupt(), [model]);
+
   return {
-    generate,
-    error,
+    messageHistory,
+    response,
     isReady,
     isGenerating,
-    isModelReady: isReady,
-    isModelGenerating: isGenerating,
-    response,
     downloadProgress,
-    interrupt,
+    error,
+    configure: configure,
+    generate: generate,
+    sendMessage: sendMessage,
+    deleteMessage: deleteMessage,
+    interrupt: interrupt,
   };
 };