react-native-executorch 0.3.2 → 0.4.1

package/src/modules/natural_language_processing/LLMModule.ts

@@ -1,62 +1,75 @@
-import { LLM } from '../../native/RnExecutorchModules';
-import { fetchResource } from '../../utils/fetchResource';
-import {
-  DEFAULT_CONTEXT_WINDOW_LENGTH,
-  DEFAULT_MESSAGE_HISTORY,
-  DEFAULT_SYSTEM_PROMPT,
-} from '../../constants/llamaDefaults';
+import { LLMController } from '../../controllers/LLMController';
 import { ResourceSource } from '../../types/common';
+import { ChatConfig, LLMTool, Message, ToolsConfig } from '../../types/llm';
 
 export class LLMModule {
-  static onDownloadProgressCallback = (_downloadProgress: number) => {};
-
-  static async load(
-    modelSource: ResourceSource,
-    tokenizerSource: ResourceSource,
-    systemPrompt = DEFAULT_SYSTEM_PROMPT,
-    messageHistory = DEFAULT_MESSAGE_HISTORY,
-    contextWindowLength = DEFAULT_CONTEXT_WINDOW_LENGTH
-  ) {
-    try {
-      const tokenizerFileUri = await fetchResource(tokenizerSource);
-      const modelFileUri = await fetchResource(
-        modelSource,
-        this.onDownloadProgressCallback
-      );
-
-      await LLM.loadLLM(
-        modelFileUri,
-        tokenizerFileUri,
-        systemPrompt,
-        messageHistory,
-        contextWindowLength
-      );
-    } catch (err) {
-      throw new Error((err as Error).message);
-    }
-  }
-
-  static async generate(input: string) {
-    try {
-      await LLM.runInference(input);
-    } catch (err) {
-      throw new Error((err as Error).message);
-    }
-  }
-
-  static onDownloadProgress(callback: (downloadProgress: number) => void) {
-    this.onDownloadProgressCallback = callback;
-  }
-
-  static onToken(callback: (data: string | undefined) => void) {
-    return LLM.onToken(callback);
+  static controller: LLMController;
+
+  static async load({
+    modelSource,
+    tokenizerSource,
+    tokenizerConfigSource,
+    onDownloadProgressCallback,
+    responseCallback,
+    messageHistoryCallback,
+  }: {
+    modelSource: ResourceSource;
+    tokenizerSource: ResourceSource;
+    tokenizerConfigSource: ResourceSource;
+    onDownloadProgressCallback?: (_downloadProgress: number) => void;
+    responseCallback?: (response: string) => void;
+    messageHistoryCallback?: (messageHistory: Message[]) => void;
+  }) {
+    this.controller = new LLMController({
+      responseCallback: responseCallback,
+      messageHistoryCallback: messageHistoryCallback,
+      onDownloadProgressCallback: onDownloadProgressCallback,
+    });
+    await this.controller.load({
+      modelSource,
+      tokenizerSource,
+      tokenizerConfigSource,
+    });
+  }
+
+  static configure({
+    chatConfig,
+    toolsConfig,
+  }: {
+    chatConfig?: Partial<ChatConfig>;
+    toolsConfig?: ToolsConfig;
+  }) {
+    this.controller.configure({ chatConfig, toolsConfig });
+  }
+
+  static async forward(input: string): Promise<string> {
+    await this.controller.forward(input);
+    return this.controller.response;
+  }
+
+  static async generate(
+    messages: Message[],
+    tools?: LLMTool[]
+  ): Promise<string> {
+    await this.controller.generate(messages, tools);
+    return this.controller.response;
+  }
+
+  static async sendMessage(message: string): Promise<Message[]> {
+    await this.controller.sendMessage(message);
+    return this.controller.messageHistory;
+  }
+
+  static async deleteMessage(index: number): Promise<Message[]> {
+    await this.controller.deleteMessage(index);
+    return this.controller.messageHistory;
   }
 
   static interrupt() {
-    LLM.interrupt();
+    this.controller.interrupt();
   }
 
   static delete() {
-    LLM.deleteModule();
+    this.controller.delete();
   }
 }

package/src/modules/natural_language_processing/SpeechToTextModule.ts

@@ -1,5 +1,7 @@
 import { ResourceSource } from '../../types/common';
 import { SpeechToTextController } from '../../controllers/SpeechToTextController';
+import { AvailableModels, SpeechToTextLanguage } from '../../types/stt';
+import { STREAMING_ACTION } from '../../constants/sttDefaults';
 
 export class SpeechToText {
   static module: SpeechToTextController;
@@ -7,9 +9,9 @@ export class SpeechToText {
   static onDownloadProgressCallback = (_downloadProgress: number) => {};
 
   static async load(
-    modelName: 'moonshine' | 'whisper',
+    modelName: AvailableModels,
     transcribeCallback: (sequence: string) => void,
-    modelDownloadProgessCallback?: (downloadProgress: number) => void,
+    modelDownloadProgressCallback?: (downloadProgress: number) => void,
     encoderSource?: ResourceSource,
     decoderSource?: ResourceSource,
     tokenizerSource?: ResourceSource,
@@ -25,7 +27,7 @@ export class SpeechToText {
   ) {
     this.module = new SpeechToTextController({
       transcribeCallback: transcribeCallback,
-      modelDownloadProgessCallback: modelDownloadProgessCallback,
+      modelDownloadProgressCallback: modelDownloadProgressCallback,
       overlapSeconds: overlapSeconds,
       windowSize: windowSize,
       streamingConfig: streamingConfig,
@@ -50,17 +52,30 @@ export class SpeechToText {
     );
   }
 
-  static async transcribe(
-    waveform: number[]
-  ): ReturnType<SpeechToTextController['transcribe']> {
-    return await this.module.transcribe(waveform);
-  }
-
   static async encode(waveform: number[]) {
     return await this.module.encode(waveform);
   }
 
-  static async decode(seq: number[], encodings?: number[]) {
+  static async decode(seq: number[], encodings: number[]) {
     return await this.module.decode(seq, encodings);
   }
+
+  static async transcribe(
+    waveform: number[],
+    audioLanguage?: SpeechToTextLanguage
+  ): ReturnType<SpeechToTextController['transcribe']> {
+    return await this.module.transcribe(waveform, audioLanguage);
+  }
+
+  static async streamingTranscribe(
+    streamAction: STREAMING_ACTION,
+    waveform?: number[],
+    audioLanguage?: SpeechToTextLanguage
+  ): ReturnType<SpeechToTextController['streamingTranscribe']> {
+    return await this.module.streamingTranscribe(
+      streamAction,
+      waveform,
+      audioLanguage
+    );
+  }
 }

package/src/modules/natural_language_processing/TextEmbeddingsModule.ts

@@ -0,0 +1,18 @@
+import { TextEmbeddingsNativeModule } from '../../native/RnExecutorchModules';
+import { ResourceSource } from '../../types/common';
+import { BaseModule } from '../BaseModule';
+
+export class TextEmbeddingsModule extends BaseModule {
+  protected static override nativeModule = TextEmbeddingsNativeModule;
+
+  static override async load(
+    modelSource: ResourceSource,
+    tokenizerSource: ResourceSource
+  ) {
+    await super.load(modelSource, tokenizerSource);
+  }
+
+  static override async forward(input: string): Promise<number[]> {
+    return this.nativeModule.forward(input);
+  }
+}

package/src/modules/natural_language_processing/TokenizerModule.ts

@@ -0,0 +1,34 @@
+import { TokenizerNativeModule } from '../../native/RnExecutorchModules';
+import { ResourceSource } from '../../types/common';
+import { BaseModule } from '../BaseModule';
+
+export class TokenizerModule extends BaseModule {
+  protected static override nativeModule = TokenizerNativeModule;
+
+  static override async load(tokenizerSource: ResourceSource) {
+    await super.load(tokenizerSource);
+  }
+
+  static async decode(
+    input: number[],
+    skipSpecialTokens = false
+  ): Promise<string> {
+    return await this.nativeModule.decode(input, skipSpecialTokens);
+  }
+
+  static async encode(input: string): Promise<number[]> {
+    return await this.nativeModule.encode(input);
+  }
+
+  static async getVocabSize(): Promise<number> {
+    return await this.nativeModule.getVocabSize();
+  }
+
+  static async idToToken(tokenId: number): Promise<string> {
+    return await this.nativeModule.idToToken(tokenId);
+  }
+
+  static async tokenToId(token: string): Promise<number> {
+    return await this.nativeModule.tokenToId(token);
+  }
+}

package/src/native/NativeClassification.ts

@@ -3,7 +3,6 @@ import { TurboModuleRegistry } from 'react-native';
 
 export interface Spec extends TurboModule {
   loadModule(modelSource: string): Promise<number>;
-
   forward(input: string): Promise<{ [category: string]: number }>;
 }
 

package/src/native/NativeImageSegmentation.ts

@@ -0,0 +1,14 @@
+import type { TurboModule } from 'react-native';
+import { TurboModuleRegistry } from 'react-native';
+
+export interface Spec extends TurboModule {
+  loadModule(modelSource: string): Promise<number>;
+
+  forward(
+    input: string,
+    classesOfInterest: string[],
+    resize: boolean
+  ): Promise<{ [category: string]: number[] }>;
+}
+
+export default TurboModuleRegistry.get<Spec>('ImageSegmentation');

package/src/native/NativeLLM.ts

@@ -1,19 +1,12 @@
 import type { TurboModule } from 'react-native';
 import { TurboModuleRegistry } from 'react-native';
 import type { EventEmitter } from 'react-native/Libraries/Types/CodegenTypes';
-import { MessageType } from '../types/common';
 
 export interface Spec extends TurboModule {
-  loadLLM(
-    modelSource: string,
-    tokenizerSource: string,
-    systemPrompt: string,
-    messageHistory: MessageType[],
-    contextWindowLength: number
-  ): Promise<string>;
-  runInference(input: string): Promise<string>;
+  loadLLM(modelSource: string, tokenizerSource: string): Promise<string>;
+  forward(input: string): Promise<string>;
   interrupt(): void;
-  deleteModule(): void;
+  releaseResources(): void;
 
   readonly onToken: EventEmitter<string>;
 }

package/src/native/NativeObjectDetection.ts

@@ -1,6 +1,6 @@
 import type { TurboModule } from 'react-native';
 import { TurboModuleRegistry } from 'react-native';
-import { Detection } from '../types/object_detection';
+import { Detection } from '../types/objectDetection';
 
 export interface Spec extends TurboModule {
   loadModule(modelSource: string): Promise<number>;

package/src/native/NativeSpeechToText.ts

@@ -8,8 +8,8 @@ export interface Spec extends TurboModule {
     modelSources: (string | number)[]
   ): Promise<number>;
   generate(waveform: number[]): Promise<number[]>;
-  encode(input: number[][]): Promise<number[]>;
-  decode(prevTokens: number[], encoderOutput: number[]): Promise<number[]>;
+  encode(input: number[]): Promise<number[]>;
+  decode(prevTokens: number[], encoderOutput: number[]): Promise<number>;
 
   readonly onToken: EventEmitter<number>;
 }

package/src/native/NativeTextEmbeddings.ts

@@ -0,0 +1,9 @@
+import type { TurboModule } from 'react-native';
+import { TurboModuleRegistry } from 'react-native';
+
+export interface Spec extends TurboModule {
+  loadModule(modelSource: string, tokenizerSource: string): Promise<number>;
+  forward(input: string): Promise<number[]>;
+}
+
+export default TurboModuleRegistry.get<Spec>('TextEmbeddings');

package/src/native/NativeTokenizer.ts

@@ -0,0 +1,13 @@
+import type { TurboModule } from 'react-native';
+import { TurboModuleRegistry } from 'react-native';
+
+export interface Spec extends TurboModule {
+  loadModule(tokenizerSource: string): Promise<number>;
+  decode(input: number[], skipSpecialTokens: boolean): Promise<string>;
+  encode(input: string): Promise<number[]>;
+  getVocabSize(): Promise<number>;
+  idToToken(tokenId: number): Promise<string>;
+  tokenToId(token: string): Promise<number>;
+}
+
+export default TurboModuleRegistry.get<Spec>('Tokenizer');

package/src/native/RnExecutorchModules.ts

@@ -1,10 +1,15 @@
 import { Platform } from 'react-native';
-import { Spec as ClassificationInterface } from './NativeClassification';
 import { Spec as ObjectDetectionInterface } from './NativeObjectDetection';
 import { Spec as StyleTransferInterface } from './NativeStyleTransfer';
+import { Spec as ImageSegmentationInterface } from './NativeImageSegmentation';
 import { Spec as ETModuleInterface } from './NativeETModule';
 import { Spec as OCRInterface } from './NativeOCR';
 import { Spec as VerticalOCRInterface } from './NativeVerticalOCR';
+import { Spec as SpeechToTextInterface } from './NativeSpeechToText';
+import { Spec as TextEmbeddingsInterface } from './NativeTextEmbeddings';
+import { Spec as LLMInterface } from './NativeLLM';
+import { Spec as ClassificationInterface } from './NativeClassification';
+import { Spec as TokenizerInterface } from './NativeTokenizer';
 
 const LINKING_ERROR =
   `The package 'react-native-executorch' doesn't seem to be linked. Make sure: \n\n` +
@@ -12,241 +17,56 @@ const LINKING_ERROR =
   '- You rebuilt the app after installing the package\n' +
   '- You are not using Expo Go\n';
 
-const LLMSpec = require('./NativeLLM').default;
-
-const LLM = LLMSpec
-  ? LLMSpec
-  : new Proxy(
-      {},
-      {
-        get() {
-          throw new Error(LINKING_ERROR);
-        },
-      }
-    );
-
-const ETModuleSpec = require('./NativeETModule').default;
-
-const ETModule = ETModuleSpec
-  ? ETModuleSpec
-  : new Proxy(
-      {},
-      {
-        get() {
-          throw new Error(LINKING_ERROR);
-        },
-      }
-    );
-
-const ClassificationSpec = require('./NativeClassification').default;
-
-const Classification = ClassificationSpec
-  ? ClassificationSpec
-  : new Proxy(
-      {},
-      {
-        get() {
-          throw new Error(LINKING_ERROR);
-        },
-      }
-    );
-
-const ObjectDetectionSpec = require('./NativeObjectDetection').default;
-
-const ObjectDetection = ObjectDetectionSpec
-  ? ObjectDetectionSpec
-  : new Proxy(
-      {},
-      {
-        get() {
-          throw new Error(LINKING_ERROR);
-        },
-      }
-    );
-
-const StyleTransferSpec = require('./NativeStyleTransfer').default;
-
-const StyleTransfer = StyleTransferSpec
-  ? StyleTransferSpec
-  : new Proxy(
-      {},
-      {
-        get() {
-          throw new Error(LINKING_ERROR);
-        },
-      }
-    );
-
-const SpeechToTextSpec = require('./NativeSpeechToText').default;
-
-const SpeechToText = SpeechToTextSpec
-  ? SpeechToTextSpec
-  : new Proxy(
-      {},
-      {
-        get() {
-          throw new Error(LINKING_ERROR);
-        },
-      }
-    );
-
-const OCRSpec = require('./NativeOCR').default;
-
-const OCR = OCRSpec
-  ? OCRSpec
-  : new Proxy(
-      {},
-      {
-        get() {
-          throw new Error(LINKING_ERROR);
-        },
-      }
-    );
-
-const VerticalOCRSpec = require('./NativeVerticalOCR').default;
-
-const VerticalOCR = VerticalOCRSpec
-  ? VerticalOCRSpec
-  : new Proxy(
-      {},
-      {
-        get() {
-          throw new Error(LINKING_ERROR);
-        },
-      }
-    );
-
-class _ObjectDetectionModule {
-  async forward(
-    input: string
-  ): ReturnType<ObjectDetectionInterface['forward']> {
-    return await ObjectDetection.forward(input);
-  }
-  async loadModule(
-    modelSource: string | number
-  ): ReturnType<ObjectDetectionInterface['loadModule']> {
-    return await ObjectDetection.loadModule(modelSource);
-  }
-}
-
-class _StyleTransferModule {
-  async forward(input: string): ReturnType<StyleTransferInterface['forward']> {
-    return await StyleTransfer.forward(input);
-  }
-  async loadModule(
-    modelSource: string | number
-  ): ReturnType<StyleTransferInterface['loadModule']> {
-    return await StyleTransfer.loadModule(modelSource);
-  }
+function returnSpecOrThrowLinkingError(spec: any) {
+  return spec
+    ? spec
+    : new Proxy(
+        {},
+        {
+          get() {
+            throw new Error(LINKING_ERROR);
+          },
+        }
+      );
 }
 
-class _SpeechToTextModule {
-  async generate(waveform: number[][]): Promise<number[]> {
-    return await SpeechToText.generate(waveform);
-  }
-
-  async loadModule(modelName: String, modelSources: (string | number)[]) {
-    return await SpeechToText.loadModule(modelName, modelSources);
-  }
-
-  async encode(input: number[]) {
-    return await SpeechToText.encode(input);
-  }
-
-  async decode(prevTokens: number[], encoderOutput?: number[]) {
-    return await SpeechToText.decode(prevTokens, encoderOutput || []);
-  }
-}
-
-class _ClassificationModule {
-  async forward(input: string): ReturnType<ClassificationInterface['forward']> {
-    return await Classification.forward(input);
-  }
-  async loadModule(
-    modelSource: string | number
-  ): ReturnType<ClassificationInterface['loadModule']> {
-    return await Classification.loadModule(modelSource);
-  }
-}
-
-class _OCRModule {
-  async forward(input: string): ReturnType<OCRInterface['forward']> {
-    return await OCR.forward(input);
-  }
-
-  async loadModule(
-    detectorSource: string,
-    recognizerSourceLarge: string,
-    recognizerSourceMedium: string,
-    recognizerSourceSmall: string,
-    symbols: string
-  ) {
-    return await OCR.loadModule(
-      detectorSource,
-      recognizerSourceLarge,
-      recognizerSourceMedium,
-      recognizerSourceSmall,
-      symbols
-    );
-  }
-}
-
-class _VerticalOCRModule {
-  async forward(input: string): ReturnType<VerticalOCRInterface['forward']> {
-    return await VerticalOCR.forward(input);
-  }
-
-  async loadModule(
-    detectorLargeSource: string,
-    detectorMediumSource: string,
-    recognizerSource: string,
-    symbols: string,
-    independentCharacters: boolean
-  ): ReturnType<VerticalOCRInterface['loadModule']> {
-    return await VerticalOCR.loadModule(
-      detectorLargeSource,
-      detectorMediumSource,
-      recognizerSource,
-      symbols,
-      independentCharacters
-    );
-  }
-}
-
-class _ETModule {
-  async forward(
-    inputs: number[][],
-    shapes: number[][],
-    inputTypes: number[]
-  ): ReturnType<ETModuleInterface['forward']> {
-    return await ETModule.forward(inputs, shapes, inputTypes);
-  }
-  async loadModule(
-    modelSource: string
-  ): ReturnType<ETModuleInterface['loadModule']> {
-    return await ETModule.loadModule(modelSource);
-  }
-  async loadMethod(
-    methodName: string
-  ): ReturnType<ETModuleInterface['loadMethod']> {
-    return await ETModule.loadMethod(methodName);
-  }
-}
+const LLMNativeModule: LLMInterface = returnSpecOrThrowLinkingError(
+  require('./NativeLLM').default
+);
+const ETModuleNativeModule: ETModuleInterface = returnSpecOrThrowLinkingError(
+  require('./NativeETModule').default
+);
+const ClassificationNativeModule: ClassificationInterface =
+  returnSpecOrThrowLinkingError(require('./NativeClassification').default);
+const ImageSegmentationNativeModule: ImageSegmentationInterface =
+  returnSpecOrThrowLinkingError(require('./NativeImageSegmentation').default);
+const ObjectDetectionNativeModule: ObjectDetectionInterface =
+  returnSpecOrThrowLinkingError(require('./NativeObjectDetection').default);
+const StyleTransferNativeModule: StyleTransferInterface =
+  returnSpecOrThrowLinkingError(require('./NativeStyleTransfer').default);
+const SpeechToTextNativeModule: SpeechToTextInterface =
+  returnSpecOrThrowLinkingError(require('./NativeSpeechToText').default);
+const OCRNativeModule: OCRInterface = returnSpecOrThrowLinkingError(
+  require('./NativeOCR').default
+);
+const VerticalOCRNativeModule: VerticalOCRInterface =
+  returnSpecOrThrowLinkingError(require('./NativeVerticalOCR').default);
+const TokenizerNativeModule: TokenizerInterface = returnSpecOrThrowLinkingError(
+  require('./NativeTokenizer').default
+);
+const TextEmbeddingsNativeModule: TextEmbeddingsInterface =
+  returnSpecOrThrowLinkingError(require('./NativeTextEmbeddings').default);
 
 export {
-  LLM,
-  ETModule,
-  Classification,
-  ObjectDetection,
-  StyleTransfer,
-  SpeechToText,
-  OCR,
-  VerticalOCR,
-  _ETModule,
-  _ClassificationModule,
-  _StyleTransferModule,
-  _ObjectDetectionModule,
-  _SpeechToTextModule,
-  _OCRModule,
-  _VerticalOCRModule,
+  LLMNativeModule,
+  ETModuleNativeModule,
+  ClassificationNativeModule,
+  ObjectDetectionNativeModule,
+  StyleTransferNativeModule,
+  ImageSegmentationNativeModule,
+  SpeechToTextNativeModule,
+  OCRNativeModule,
+  VerticalOCRNativeModule,
+  TextEmbeddingsNativeModule,
+  TokenizerNativeModule,
 };