react-native-executorch 0.4.7 → 0.5.0

package/src/modules/natural_language_processing/LLMModule.ts

@@ -3,39 +3,39 @@ import { ResourceSource } from '../../types/common';
 import { ChatConfig, LLMTool, Message, ToolsConfig } from '../../types/llm';
 
 export class LLMModule {
-  static controller: LLMController;
+  private controller: LLMController;
 
-  static async load({
-    modelSource,
-    tokenizerSource,
-    tokenizerConfigSource,
-    onDownloadProgressCallback,
+  constructor({
     tokenCallback,
     responseCallback,
     messageHistoryCallback,
   }: {
-    modelSource: ResourceSource;
-    tokenizerSource: ResourceSource;
-    tokenizerConfigSource: ResourceSource;
-    onDownloadProgressCallback?: (_downloadProgress: number) => void;
     tokenCallback?: (token: string) => void;
     responseCallback?: (response: string) => void;
     messageHistoryCallback?: (messageHistory: Message[]) => void;
-  }) {
+  } = {}) {
     this.controller = new LLMController({
-      tokenCallback: tokenCallback,
-      responseCallback: responseCallback,
-      messageHistoryCallback: messageHistoryCallback,
-      onDownloadProgressCallback: onDownloadProgressCallback,
+      tokenCallback,
+      responseCallback,
+      messageHistoryCallback,
     });
+  }
+
+  async load(
+    model: {
+      modelSource: ResourceSource;
+      tokenizerSource: ResourceSource;
+      tokenizerConfigSource: ResourceSource;
+    },
+    onDownloadProgressCallback: (progress: number) => void = () => {}
+  ) {
     await this.controller.load({
-      modelSource,
-      tokenizerSource,
-      tokenizerConfigSource,
+      ...model,
+      onDownloadProgressCallback,
     });
   }
 
-  static setTokenCallback({
+  setTokenCallback({
     tokenCallback,
   }: {
     tokenCallback: (token: string) => void;
@@ -43,7 +43,7 @@ export class LLMModule {
     this.controller.setTokenCallback(tokenCallback);
   }
 
-  static configure({
+  configure({
     chatConfig,
     toolsConfig,
   }: {
@@ -53,34 +53,31 @@ export class LLMModule {
     this.controller.configure({ chatConfig, toolsConfig });
   }
 
-  static async forward(input: string): Promise<string> {
+  async forward(input: string): Promise<string> {
     await this.controller.forward(input);
     return this.controller.response;
   }
 
-  static async generate(
-    messages: Message[],
-    tools?: LLMTool[]
-  ): Promise<string> {
+  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[]> {
+  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);
+  deleteMessage(index: number): Message[] {
+    this.controller.deleteMessage(index);
     return this.controller.messageHistory;
   }
 
-  static interrupt() {
+  interrupt() {
     this.controller.interrupt();
   }
 
-  static delete() {
+  delete() {
     this.controller.delete();
   }
 }

package/src/modules/natural_language_processing/SpeechToTextModule.ts

@@ -1,81 +1,98 @@
-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;
-
-  static onDownloadProgressCallback = (_downloadProgress: number) => {};
-
-  static async load(
-    modelName: AvailableModels,
-    transcribeCallback: (sequence: string) => void,
-    modelDownloadProgressCallback?: (downloadProgress: number) => void,
-    encoderSource?: ResourceSource,
-    decoderSource?: ResourceSource,
-    tokenizerSource?: ResourceSource,
-    overlapSeconds?: ConstructorParameters<
-      typeof SpeechToTextController
-    >['0']['overlapSeconds'],
-    windowSize?: ConstructorParameters<
-      typeof SpeechToTextController
-    >['0']['windowSize'],
-    streamingConfig?: ConstructorParameters<
-      typeof SpeechToTextController
-    >['0']['streamingConfig']
-  ) {
-    this.module = new SpeechToTextController({
-      transcribeCallback: transcribeCallback,
-      modelDownloadProgressCallback: modelDownloadProgressCallback,
-      overlapSeconds: overlapSeconds,
-      windowSize: windowSize,
-      streamingConfig: streamingConfig,
-    });
-    await this.module.loadModel(
-      (modelName = modelName),
-      (encoderSource = encoderSource),
-      (decoderSource = decoderSource),
-      (tokenizerSource = tokenizerSource)
-    );
-  }
+import { DecodingOptions, SpeechToTextModelConfig } from '../../types/stt';
+import { ASR } from '../../utils/SpeechToTextModule/ASR';
+import { OnlineASRProcessor } from '../../utils/SpeechToTextModule/OnlineProcessor';
+
+export class SpeechToTextModule {
+  private modelConfig!: SpeechToTextModelConfig;
+  private asr: ASR = new ASR();
 
-  static configureStreaming(
-    overlapSeconds: Parameters<SpeechToTextController['configureStreaming']>[0],
-    windowSize: Parameters<SpeechToTextController['configureStreaming']>[1],
-    streamingConfig: Parameters<SpeechToTextController['configureStreaming']>[2]
+  private processor: OnlineASRProcessor = new OnlineASRProcessor(this.asr);
+  private isStreaming = false;
+  private readyToProcess = false;
+  private minAudioSamples: number = 1 * 16000; // 1 second
+
+  public async load(
+    model: SpeechToTextModelConfig,
+    onDownloadProgressCallback: (progress: number) => void = () => {}
   ) {
-    this.module?.configureStreaming(
-      overlapSeconds,
-      windowSize,
-      streamingConfig
-    );
+    this.modelConfig = model;
+    return this.asr.load(model, onDownloadProgressCallback);
   }
 
-  static async encode(waveform: number[]) {
-    return await this.module.encode(waveform);
+  public async encode(waveform: Float32Array): Promise<void> {
+    return this.asr.encode(waveform);
   }
 
-  static async decode(seq: number[], encodings: number[]) {
-    return await this.module.decode(seq, encodings);
+  public async decode(tokens: number[]): Promise<Float32Array> {
+    return this.asr.decode(tokens);
   }
 
-  static async transcribe(
+  public async transcribe(
     waveform: number[],
-    audioLanguage?: SpeechToTextLanguage
-  ): ReturnType<SpeechToTextController['transcribe']> {
-    return await this.module.transcribe(waveform, audioLanguage);
+    options: DecodingOptions = {}
+  ): Promise<string> {
+    this.validateOptions(options);
+
+    const segments = await this.asr.transcribe(waveform, options);
+
+    let transcription = '';
+    for (const segment of segments) {
+      for (const word of segment.words) {
+        transcription += ` ${word.word}`;
+      }
+    }
+
+    return transcription.trim();
+  }
+
+  public async *stream(options: DecodingOptions = {}) {
+    if (this.isStreaming) {
+      throw new Error('Streaming is already in progress');
+    }
+    this.validateOptions(options);
+    this.resetStreamState();
+
+    this.isStreaming = true;
+    while (this.isStreaming) {
+      if (
+        !this.readyToProcess ||
+        this.processor.audioBuffer.length < this.minAudioSamples
+      ) {
+        await new Promise((resolve) => setTimeout(resolve, 100));
+        continue;
+      }
+
+      const { committed, nonCommitted } =
+        await this.processor.processIter(options);
+      yield { committed, nonCommitted };
+      this.readyToProcess = false;
+    }
+
+    const { committed } = await this.processor.finish();
+    yield { committed, nonCommitted: '' };
+  }
+
+  public streamStop() {
+    this.isStreaming = false;
+  }
+
+  public streamInsert(waveform: number[]) {
+    this.processor.insertAudioChunk(waveform);
+    this.readyToProcess = true;
+  }
+
+  private validateOptions(options: DecodingOptions) {
+    if (!this.modelConfig.isMultilingual && options.language) {
+      throw new Error('Model is not multilingual, cannot set language');
+    }
+    if (this.modelConfig.isMultilingual && !options.language) {
+      throw new Error('Model is multilingual, provide a language');
+    }
   }
 
-  static async streamingTranscribe(
-    streamAction: STREAMING_ACTION,
-    waveform?: number[],
-    audioLanguage?: SpeechToTextLanguage
-  ): ReturnType<SpeechToTextController['streamingTranscribe']> {
-    return await this.module.streamingTranscribe(
-      streamAction,
-      waveform,
-      audioLanguage
-    );
+  private resetStreamState() {
+    this.isStreaming = false;
+    this.readyToProcess = false;
+    this.processor = new OnlineASRProcessor(this.asr);
   }
 }

package/src/modules/natural_language_processing/TextEmbeddingsModule.ts

@@ -1,18 +1,33 @@
-import { TextEmbeddingsNativeModule } from '../../native/RnExecutorchModules';
 import { ResourceSource } from '../../types/common';
-import { BaseModule } from '../BaseModule';
+import { ResourceFetcher } from '../../utils/ResourceFetcher';
+import { BaseNonStaticModule } from '../BaseNonStaticModule';
 
-export class TextEmbeddingsModule extends BaseModule {
-  protected static override nativeModule = TextEmbeddingsNativeModule;
-
-  static override async load(
-    modelSource: ResourceSource,
-    tokenizerSource: ResourceSource
-  ) {
-    await super.load(modelSource, tokenizerSource);
+export class TextEmbeddingsModule extends BaseNonStaticModule {
+  async load(
+    model: { modelSource: ResourceSource; tokenizerSource: ResourceSource },
+    onDownloadProgressCallback: (progress: number) => void = () => {}
+  ): Promise<void> {
+    const modelPromise = ResourceFetcher.fetch(
+      onDownloadProgressCallback,
+      model.modelSource
+    );
+    const tokenizerPromise = ResourceFetcher.fetch(
+      undefined,
+      model.tokenizerSource
+    );
+    const [modelResult, tokenizerResult] = await Promise.all([
+      modelPromise,
+      tokenizerPromise,
+    ]);
+    const modelPath = modelResult?.[0];
+    const tokenizerPath = tokenizerResult?.[0];
+    if (!modelPath || !tokenizerPath) {
+      throw new Error('Download interrupted.');
+    }
+    this.nativeModule = global.loadTextEmbeddings(modelPath, tokenizerPath);
   }
 
-  static override async forward(input: string): Promise<number[]> {
-    return this.nativeModule.forward(input);
+  async forward(input: string): Promise<Float32Array> {
+    return new Float32Array(await this.nativeModule.generate(input));
   }
 }

package/src/modules/natural_language_processing/TokenizerModule.ts

@@ -1,34 +1,44 @@
-import { TokenizerNativeModule } from '../../native/RnExecutorchModules';
 import { ResourceSource } from '../../types/common';
-import { BaseModule } from '../BaseModule';
+import { ResourceFetcher } from '../../utils/ResourceFetcher';
 
-export class TokenizerModule extends BaseModule {
-  protected static override nativeModule = TokenizerNativeModule;
+export class TokenizerModule {
+  nativeModule: any;
 
-  static override async load(tokenizerSource: ResourceSource) {
-    await super.load(tokenizerSource);
+  async load(
+    tokenizer: { tokenizerSource: ResourceSource },
+    onDownloadProgressCallback: (progress: number) => void = () => {}
+  ): Promise<void> {
+    const paths = await ResourceFetcher.fetch(
+      onDownloadProgressCallback,
+      tokenizer.tokenizerSource
+    );
+    const path = paths?.[0];
+    if (!path) {
+      throw new Error('Download interrupted.');
+    }
+    this.nativeModule = global.loadTokenizerModule(path);
   }
 
-  static async decode(
-    input: number[],
-    skipSpecialTokens = false
-  ): Promise<string> {
-    return await this.nativeModule.decode(input, skipSpecialTokens);
+  async encode(s: string) {
+    return await this.nativeModule.encode(s);
   }
 
-  static async encode(input: string): Promise<number[]> {
-    return await this.nativeModule.encode(input);
+  async decode(tokens: number[], skipSpecialTokens: boolean = true) {
+    if (tokens.length === 0) {
+      return '';
+    }
+    return await this.nativeModule.decode(tokens, skipSpecialTokens);
   }
 
-  static async getVocabSize(): Promise<number> {
+  async getVocabSize(): Promise<number> {
     return await this.nativeModule.getVocabSize();
   }
 
-  static async idToToken(tokenId: number): Promise<string> {
-    return await this.nativeModule.idToToken(tokenId);
+  async idToToken(tokenId: number): Promise<string> {
+    return this.nativeModule.idToToken(tokenId);
   }
 
-  static async tokenToId(token: string): Promise<number> {
+  async tokenToId(token: string): Promise<number> {
     return await this.nativeModule.tokenToId(token);
   }
 }

package/src/native/NativeETInstaller.ts

@@ -0,0 +1,8 @@
+import type { TurboModule } from 'react-native';
+import { TurboModuleRegistry } from 'react-native';
+
+export interface Spec extends TurboModule {
+  install(): boolean;
+}
+
+export default TurboModuleRegistry.get<Spec>('ETInstaller');

package/src/native/RnExecutorchModules.ts

@@ -1,15 +1,6 @@
 import { Platform } from 'react-native';
-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';
+import { Spec as ETInstallerInterface } from './NativeETInstaller';
 
 const LINKING_ERROR =
   `The package 'react-native-executorch' doesn't seem to be linked. Make sure: \n\n` +
@@ -30,43 +21,10 @@ function returnSpecOrThrowLinkingError(spec: any) {
       );
 }
 
+const ETInstallerNativeModule: ETInstallerInterface =
+  returnSpecOrThrowLinkingError(require('./NativeETInstaller').default);
 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 {
-  LLMNativeModule,
-  ETModuleNativeModule,
-  ClassificationNativeModule,
-  ObjectDetectionNativeModule,
-  StyleTransferNativeModule,
-  ImageSegmentationNativeModule,
-  SpeechToTextNativeModule,
-  OCRNativeModule,
-  VerticalOCRNativeModule,
-  TextEmbeddingsNativeModule,
-  TokenizerNativeModule,
-};
+export { LLMNativeModule, ETInstallerNativeModule };

package/src/types/common.ts

@@ -1,17 +1,45 @@
-export const getTypeIdentifier = (input: ETInput): number => {
-  if (input instanceof Int8Array) return 1;
-  if (input instanceof Int32Array) return 3;
-  if (input instanceof BigInt64Array) return 4;
-  if (input instanceof Float32Array) return 6;
-  if (input instanceof Float64Array) return 7;
-  return -1;
-};
-
 export type ResourceSource = string | number | object;
 
-export type ETInput =
+export enum ScalarType {
+  BYTE = 0,
+  CHAR = 1,
+  SHORT = 2,
+  INT = 3,
+  LONG = 4,
+  HALF = 5,
+  FLOAT = 6,
+  DOUBLE = 7,
+  BOOL = 11,
+  QINT8 = 12,
+  QUINT8 = 13,
+  QINT32 = 14,
+  QUINT4X2 = 16,
+  QUINT2X4 = 17,
+  BITS16 = 22,
+  FLOAT8E5M2 = 23,
+  FLOAT8E4M3FN = 24,
+  FLOAT8E5M2FNUZ = 25,
+  FLOAT8E4M3FNUZ = 26,
+  UINT16 = 27,
+  UINT32 = 28,
+  UINT64 = 29,
+}
+
+export type TensorBuffer =
+  | ArrayBuffer
+  | Float32Array
+  | Float64Array
   | Int8Array
+  | Int16Array
   | Int32Array
+  | Uint8Array
+  | Uint16Array
+  | Uint32Array
   | BigInt64Array
-  | Float32Array
-  | Float64Array;
+  | BigUint64Array;
+
+export interface TensorPtr {
+  dataPtr: TensorBuffer;
+  sizes: number[];
+  scalarType: ScalarType;
+}