chore:refactor ts cv hooks (#64)

## Description
<!-- Provide a concise and descriptive summary of the changes
implemented in this PR. -->

### Type of change
- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [ ] Documentation update (improves or adds clarity to existing
documentation)

### Tested on
- [ ] iOS
- [x] Android

### Testing instructions
<!-- Provide step-by-step instructions on how to test your changes.
Include setup details if necessary. -->

### Screenshots
<!-- Add screenshots here, if applicable -->

### Related issues
<!-- Link related issues here using #issue-number -->

### Checklist
- [x] I have performed a self-review of my code
- [ ] I have commented my code, particularly in hard-to-understand areas
- [x] I have updated the documentation accordingly
- [x] My changes generate no new warnings

### Additional notes
<!-- Include any additional information, assumptions, or context that
reviewers might need to understand this PR. -->

---------

Co-authored-by: Mateusz Kopcinski <[email protected]>
Co-authored-by: Mateusz Kopciński <[email protected]>

Author: 3 people

android/src/main/java/com/swmansion/rnexecutorch/ObjectDetection.kt

@@ -35,7 +35,7 @@ class ObjectDetection(reactContext: ReactApplicationContext) :
       ssdLiteLarge.loadModel(modelSource)
       promise.resolve(0)
     } catch (e: Exception) {
-      promise.reject(e.message!!, ETError.InvalidModelPath.toString())
+      promise.reject(e.message!!, ETError.InvalidModelSource.toString())
     }
   }
 

docs/docs/guides/running-llms.md

@@ -23,7 +23,7 @@ const llama = useLLM({
 });
 ```
 
-The code snippet above fetches the model from the specified URL, loads it into memory, and returns an object with various methods and properties for controlling the model. You can monitor the loading progress by checking the `llama.downloadProgress` and `llama.isModelReady` property, and if anything goes wrong, the `llama.error` property will contain the error message.
+The code snippet above fetches the model from the specified URL, loads it into memory, and returns an object with various methods and properties for controlling the model. You can monitor the loading progress by checking the `llama.downloadProgress` and `llama.isReady` property, and if anything goes wrong, the `llama.error` property will contain the error message.
 
 :::danger[Danger]
 Lower-end devices might not be able to fit LLMs into memory. We recommend using quantized models to reduce the memory footprint.
@@ -50,9 +50,9 @@ Given computational constraints, our architecture is designed to support only on
 | `generate`          | `(input: string) => Promise<void>` | Function to start generating a response with the given input string.                                            |
 | `response`          | `string`                           | State of the generated response. This field is updated with each token generated by the model                   |
 | `error`             | <code>string &#124; null</code>    | Contains the error message if the model failed to load                                                          |
-| `isModelGenerating` | `boolean`                          | Indicates whether the model is currently generating a response                                                  |
+| `isGenerating` | `boolean`                          | Indicates whether the model is currently generating a response                                                  |
 | `interrupt`         | `() => void`                       | Function to interrupt the current inference                                                                     |
-| `isModelReady`      | `boolean`                          | Indicates whether the model is ready                                                                            |
+| `isReady`      | `boolean`                          | Indicates whether the model is ready                                                                            |
 | `downloadProgress`  | `number`                           | Represents the download progress as a value between 0 and 1, indicating the extent of the model file retrieval. |
 
 ### Loading models

examples/computer-vision/screens/ClassificationScreen.tsx

@@ -44,12 +44,9 @@ export const ClassificationScreen = ({
     }
   };
 
-  if (!model.isModelReady) {
+  if (!model.isReady) {
     return (
-      <Spinner
-        visible={!model.isModelReady}
-        textContent={`Loading the model...`}
-      />
+      <Spinner visible={!model.isReady} textContent={`Loading the model...`} />
     );
   }
 

examples/computer-vision/screens/ObjectDetectionScreen.tsx

@@ -52,10 +52,10 @@ export const ObjectDetectionScreen = ({
     }
   };
 
-  if (!ssdLite.isModelReady) {
+  if (!ssdLite.isReady) {
     return (
       <Spinner
-        visible={!ssdLite.isModelReady}
+        visible={!ssdLite.isReady}
         textContent={`Loading the model...`}
       />
     );

examples/computer-vision/screens/StyleTransferScreen.tsx

@@ -37,12 +37,9 @@ export const StyleTransferScreen = ({
     }
   };
 
-  if (!model.isModelReady) {
+  if (!model.isReady) {
     return (
-      <Spinner
-        visible={!model.isModelReady}
-        textContent={`Loading the model...`}
-      />
+      <Spinner visible={!model.isReady} textContent={`Loading the model...`} />
     );
   }
 

examples/computer-vision/yarn.lock

@@ -3352,7 +3352,7 @@ __metadata:
     metro-config: ^0.81.0
     react: 18.3.1
     react-native: 0.76.3
-    react-native-executorch: ../../react-native-executorch-0.1.100.tgz
+    react-native-executorch: ^0.1.3
     react-native-image-picker: ^7.2.2
     react-native-loading-spinner-overlay: ^3.0.1
     react-native-reanimated: ^3.16.3
@@ -6799,13 +6799,13 @@ __metadata:
   languageName: node
   linkType: hard
 
-"react-native-executorch@file:../../react-native-executorch-0.1.100.tgz::locator=computer-vision%40workspace%3A.":
-  version: 0.1.100
-  resolution: "react-native-executorch@file:../../react-native-executorch-0.1.100.tgz::locator=computer-vision%40workspace%3A."
+"react-native-executorch@npm:^0.1.3":
+  version: 0.1.3
+  resolution: "react-native-executorch@npm:0.1.3"
   peerDependencies:
     react: "*"
     react-native: "*"
-  checksum: f258452e2050df59e150938f6482ef8eee5fbd4ef4fc4073a920293ca87d543daddf76c560701d0c2626e6677d964b446dad8e670e978ea4f80d0a1bd17dfa03
+  checksum: b49f8ca9ae8a7de4a7f2263887537626859507c7d60af47360515b405c7778b48c4c067074e7ce7857782a6737cf47cf5dadada03ae9319a6bf577f8490f431d
   languageName: node
   linkType: hard
 

examples/llama/components/Messages.tsx

@@ -9,13 +9,13 @@ import MessageItem from './MessageItem';
 interface MessagesComponentProps {
   chatHistory: Array<MessageType>;
   llmResponse: string;
-  isModelGenerating: boolean;
+  isGenerating: boolean;
 }
 
 export default function Messages({
   chatHistory,
   llmResponse,
-  isModelGenerating,
+  isGenerating,
 }: MessagesComponentProps) {
   const scrollViewRef = useRef<ScrollView>(null);
 
@@ -29,7 +29,7 @@ export default function Messages({
           {chatHistory.map((message, index) => (
             <MessageItem key={index} message={message} />
           ))}
-          {isModelGenerating && (
+          {isGenerating && (
             <View style={styles.aiMessage}>
               <View style={styles.aiMessageIconContainer}>
                 <LlamaIcon width={24} height={24} />

examples/llama/screens/ChatScreen.tsx

@@ -31,10 +31,10 @@ export default function ChatScreen() {
   });
   const textInputRef = useRef<TextInput>(null);
   useEffect(() => {
-    if (llama.response && !llama.isModelGenerating) {
+    if (llama.response && !llama.isGenerating) {
       appendToMessageHistory(llama.response, 'ai');
     }
-  }, [llama.response, llama.isModelGenerating]);
+  }, [llama.response, llama.isGenerating]);
 
   const appendToMessageHistory = (input: string, role: SenderType) => {
     setChatHistory((prevHistory) => [
@@ -54,9 +54,9 @@ export default function ChatScreen() {
     }
   };
 
-  return !llama.isModelReady ? (
+  return !llama.isReady ? (
     <Spinner
-      visible={!llama.isModelReady}
+      visible={!llama.isReady}
       textContent={`Loading the model ${(llama.downloadProgress * 100).toFixed(0)} %`}
     />
   ) : (
@@ -76,7 +76,7 @@ export default function ChatScreen() {
               <Messages
                 chatHistory={chatHistory}
                 llmResponse={llama.response}
-                isModelGenerating={llama.isModelGenerating}
+                isGenerating={llama.isGenerating}
               />
             </View>
           ) : (
@@ -108,13 +108,13 @@ export default function ChatScreen() {
               <TouchableOpacity
                 style={styles.sendChatTouchable}
                 onPress={async () =>
-                  !llama.isModelGenerating && (await sendMessage())
+                  !llama.isGenerating && (await sendMessage())
                 }
               >
                 <SendIcon height={24} width={24} padding={4} margin={8} />
               </TouchableOpacity>
             )}
-            {llama.isModelGenerating && (
+            {llama.isGenerating && (
               <TouchableOpacity
                 style={styles.sendChatTouchable}
                 onPress={llama.interrupt}

src/ETModule.ts

@@ -1,18 +1,8 @@
-import { useEffect, useState } from 'react';
-import { Image } from 'react-native';
-import { ETModule } from './native/RnExecutorchModules';
-import { ETError, getError } from './Error';
-import { ETInput, ExecutorchModule } from './types';
-
-const getTypeIdentifier = (arr: ETInput): number => {
-  if (arr instanceof Int8Array) return 0;
-  if (arr instanceof Int32Array) return 1;
-  if (arr instanceof BigInt64Array) return 2;
-  if (arr instanceof Float32Array) return 3;
-  if (arr instanceof Float64Array) return 4;
-
-  return -1;
-};
+import { useState } from 'react';
+import { _ETModule } from './native/RnExecutorchModules';
+import { getError } from './Error';
+import { ExecutorchModule } from './types/common';
+import { useModule } from './useModule';
 
 interface Props {
   modelSource: string | number;
@@ -21,54 +11,20 @@ interface Props {
 export const useExecutorchModule = ({
   modelSource,
 }: Props): ExecutorchModule => {
-  const [error, setError] = useState<string | null>(null);
-  const [isModelLoading, setIsModelLoading] = useState(true);
-  const [isModelGenerating, setIsModelGenerating] = useState(false);
-
-  useEffect(() => {
-    const loadModel = async () => {
-      let path = modelSource;
-      if (typeof modelSource === 'number') {
-        path = Image.resolveAssetSource(modelSource).uri;
-      }
-
-      try {
-        setIsModelLoading(true);
-        await ETModule.loadModule(path);
-        setIsModelLoading(false);
-      } catch (e: unknown) {
-        setError(getError(e));
-        setIsModelLoading(false);
-      }
-    };
-    loadModel();
-  }, [modelSource]);
-
-  const forward = async (input: ETInput, shape: number[]) => {
-    if (isModelLoading) {
-      throw new Error(getError(ETError.ModuleNotLoaded));
-    }
-
-    const inputType = getTypeIdentifier(input);
-    if (inputType === -1) {
-      throw new Error(getError(ETError.InvalidArgument));
-    }
-
-    try {
-      const numberArray = [...input];
-      setIsModelGenerating(true);
-      const output = await ETModule.forward(numberArray, shape, inputType);
-      setIsModelGenerating(false);
-      return output;
-    } catch (e) {
-      setIsModelGenerating(false);
-      throw new Error(getError(e));
-    }
-  };
+  const [module] = useState(() => new _ETModule());
+  const {
+    error,
+    isReady,
+    isGenerating,
+    forwardETInput: forward,
+  } = useModule({
+    modelSource,
+    module,
+  });
 
   const loadMethod = async (methodName: string) => {
     try {
-      await ETModule.loadMethod(methodName);
+      await module.loadMethod(methodName);
     } catch (e) {
       throw new Error(getError(e));
     }
@@ -79,11 +35,11 @@ export const useExecutorchModule = ({
   };
 
   return {
-    error: error,
-    isModelLoading: isModelLoading,
-    isModelGenerating: isModelGenerating,
-    forward: forward,
-    loadMethod: loadMethod,
-    loadForward: loadForward,
+    error,
+    isReady,
+    isGenerating,
+    forward,
+    loadMethod,
+    loadForward,
   };
 };

src/LLM.ts

@@ -1,6 +1,6 @@
 import { useCallback, useEffect, useRef, useState } from 'react';
 import { EventSubscription, Image } from 'react-native';
-import { ResourceSource, Model } from './types';
+import { ResourceSource, Model } from './types/common';
 import {
   DEFAULT_CONTEXT_WINDOW_LENGTH,
   DEFAULT_SYSTEM_PROMPT,
@@ -24,8 +24,8 @@ export const useLLM = ({
   contextWindowLength?: number;
 }): Model => {
   const [error, setError] = useState<string | null>(null);
-  const [isModelReady, setIsModelReady] = useState(false);
-  const [isModelGenerating, setIsModelGenerating] = useState(false);
+  const [isReady, setIsReady] = useState(false);
+  const [isGenerating, setIsGenerating] = useState(false);
   const [response, setResponse] = useState('');
   const [downloadProgress, setDownloadProgress] = useState(0);
   const downloadProgressListener = useRef<null | EventSubscription>(null);
@@ -65,7 +65,7 @@ export const useLLM = ({
           contextWindowLength
         );
 
-        setIsModelReady(true);
+        setIsReady(true);
 
         tokenGeneratedListener.current = LLM.onToken(
           (data: string | undefined) => {
@@ -75,13 +75,13 @@ export const useLLM = ({
             if (data !== EOT_TOKEN) {
               setResponse((prevResponse) => prevResponse + data);
             } else {
-              setIsModelGenerating(false);
+              setIsGenerating(false);
             }
           }
         );
       } catch (err) {
         const message = (err as Error).message;
-        setIsModelReady(false);
+        setIsReady(false);
         setError(message);
       }
     };
@@ -99,7 +99,7 @@ export const useLLM = ({
 
   const generate = useCallback(
     async (input: string): Promise<void> => {
-      if (!isModelReady) {
+      if (!isReady) {
         throw new Error('Model is still loading');
       }
       if (error) {
@@ -108,21 +108,23 @@ export const useLLM = ({
 
       try {
         setResponse('');
-        setIsModelGenerating(true);
+        setIsGenerating(true);
         await LLM.runInference(input);
       } catch (err) {
-        setIsModelGenerating(false);
+        setIsGenerating(false);
         throw new Error((err as Error).message);
       }
     },
-    [isModelReady, error]
+    [isReady, error]
   );
 
   return {
     generate,
     error,
-    isModelReady,
-    isModelGenerating,
+    isReady,
+    isGenerating,
+    isModelReady: isReady,
+    isModelGenerating: isGenerating,
     response,
     downloadProgress,
     interrupt,