react-native-executorch 0.3.3 → 0.4.2

package/src/constants/sttDefaults.ts

@@ -5,52 +5,62 @@ import {
   WHISPER_TINY_ENCODER,
   WHISPER_TINY_DECODER,
   WHISPER_TOKENIZER,
+  WHISPER_TINY_MULTILINGUAL_ENCODER,
+  WHISPER_TINY_MULTILINGUAL_DECODER,
+  WHISPER_TINY_MULTILINGUAL_TOKENIZER,
 } from './modelUrls';
+import { AvailableModels, ModelConfig } from '../types/stt';
 
 export const SAMPLE_RATE = 16_000;
 export const SECOND = SAMPLE_RATE;
 export const HAMMING_DIST_THRESHOLD = 1;
 
-export interface ModelConfig {
+const whisperTinyModelConfig = {
   sources: {
-    encoder: string;
-    decoder: string;
-  };
+    encoder: WHISPER_TINY_ENCODER,
+    decoder: WHISPER_TINY_DECODER,
+  },
   tokenizer: {
-    source: string;
-    sos: number;
-    eos: number;
-    specialChar: string;
-  };
-}
+    source: WHISPER_TOKENIZER,
+    bos: 50257, // FIXME: this is a placeholder and needs to be changed
+    eos: 50256, // FIXME: this is a placeholder and needs to be changed
+  },
+  isMultilingual: false,
+};
+
+const moonshineTinyModelConfig = {
+  sources: {
+    encoder: MOONSHINE_TINY_ENCODER,
+    decoder: MOONSHINE_TINY_DECODER,
+  },
+  tokenizer: {
+    source: MOONSHINE_TOKENIZER,
+    bos: 1, // FIXME: this is a placeholder and needs to be changed
+    eos: 2, // FIXME: this is a placeholder and needs to be changed
+  },
+  isMultilingual: false,
+};
 
-export const MODEL_CONFIGS: { [key in 'moonshine' | 'whisper']: ModelConfig } =
-  {
-    moonshine: {
-      sources: {
-        encoder: MOONSHINE_TINY_ENCODER,
-        decoder: MOONSHINE_TINY_DECODER,
-      },
-      tokenizer: {
-        source: MOONSHINE_TOKENIZER,
-        sos: 1,
-        eos: 2,
-        specialChar: '\u2581',
-      },
-    },
-    whisper: {
-      sources: {
-        encoder: WHISPER_TINY_ENCODER,
-        decoder: WHISPER_TINY_DECODER,
-      },
-      tokenizer: {
-        source: WHISPER_TOKENIZER,
-        sos: 50257,
-        eos: 50256,
-        specialChar: 'Ġ',
-      },
-    },
-  };
+const whisperTinyMultilingualModelConfig = {
+  sources: {
+    encoder: WHISPER_TINY_MULTILINGUAL_ENCODER,
+    decoder: WHISPER_TINY_MULTILINGUAL_DECODER,
+  },
+  tokenizer: {
+    source: WHISPER_TINY_MULTILINGUAL_TOKENIZER,
+    bos: 50258, // FIXME: this is a placeholder and needs to be changed
+    eos: 50257, // FIXME: this is a placeholder and needs to be changed
+  },
+  isMultilingual: true,
+};
+
+export const MODEL_CONFIGS: {
+  [key in AvailableModels]: ModelConfig;
+} = {
+  moonshine: moonshineTinyModelConfig,
+  whisper: whisperTinyModelConfig,
+  whisperMultilingual: whisperTinyMultilingualModelConfig,
+};
 
 export const MODES = {
   fast: {
@@ -66,3 +76,11 @@ export const MODES = {
     overlapSeconds: 3,
   },
 };
+
+export const NUM_TOKENS_TO_TRIM = 3;
+
+export enum STREAMING_ACTION {
+  START,
+  DATA,
+  STOP,
+}

package/src/controllers/LLMController.ts

@@ -0,0 +1,286 @@
+import { EventSubscription } from 'react-native';
+import { ResourceSource } from '../types/common';
+import { ResourceFetcher } from '../utils/ResourceFetcher';
+import { ETError, getError } from '../Error';
+import { Template } from '@huggingface/jinja';
+import { DEFAULT_CHAT_CONFIG } from '../constants/llmDefaults';
+import { readAsStringAsync } from 'expo-file-system';
+import {
+  ChatConfig,
+  LLMTool,
+  Message,
+  SPECIAL_TOKENS,
+  ToolsConfig,
+} from '../types/llm';
+import { LLMNativeModule } from '../native/RnExecutorchModules';
+import { parseToolCall } from '../utils/llm';
+
+export class LLMController {
+  private nativeModule: typeof LLMNativeModule;
+  private chatConfig: ChatConfig = DEFAULT_CHAT_CONFIG;
+  private toolsConfig: ToolsConfig | undefined;
+  private tokenizerConfig: any;
+  private onToken: EventSubscription | null = null;
+  private _response = '';
+  private _isReady = false;
+  private _isGenerating = false;
+  private _messageHistory: Message[] = [];
+
+  // User callbacks
+  private responseCallback: (response: string) => void;
+  private messageHistoryCallback: (messageHistory: Message[]) => void;
+  private isReadyCallback: (isReady: boolean) => void;
+  private isGeneratingCallback: (isGenerating: boolean) => void;
+  private onDownloadProgressCallback:
+    | ((downloadProgress: number) => void)
+    | undefined;
+
+  constructor({
+    responseCallback,
+    messageHistoryCallback,
+    isReadyCallback,
+    isGeneratingCallback,
+    onDownloadProgressCallback,
+  }: {
+    responseCallback?: (response: string) => void;
+    messageHistoryCallback?: (messageHistory: Message[]) => void;
+    isReadyCallback?: (isReady: boolean) => void;
+    isGeneratingCallback?: (isGenerating: boolean) => void;
+    onDownloadProgressCallback?: (downloadProgress: number) => void;
+  }) {
+    this.responseCallback = (response) => {
+      this._response = response;
+      responseCallback?.(response);
+    };
+    this.messageHistoryCallback = (messageHistory) => {
+      this._messageHistory = messageHistory;
+      messageHistoryCallback?.(messageHistory);
+    };
+    this.isReadyCallback = (isReady) => {
+      this._isReady = isReady;
+      isReadyCallback?.(isReady);
+    };
+    this.isGeneratingCallback = (isGenerating) => {
+      this._isGenerating = isGenerating;
+      isGeneratingCallback?.(isGenerating);
+    };
+
+    this.onDownloadProgressCallback = onDownloadProgressCallback;
+
+    this.nativeModule = LLMNativeModule;
+  }
+
+  public get response() {
+    return this._response;
+  }
+  public get isReady() {
+    return this._isReady;
+  }
+  public get isGenerating() {
+    return this._isGenerating;
+  }
+  public get messageHistory() {
+    return this._messageHistory;
+  }
+
+  public async load({
+    modelSource,
+    tokenizerSource,
+    tokenizerConfigSource,
+  }: {
+    modelSource: ResourceSource;
+    tokenizerSource: ResourceSource;
+    tokenizerConfigSource: ResourceSource;
+  }) {
+    // reset inner state when loading new model
+    this.responseCallback('');
+    this.messageHistoryCallback(this.chatConfig.initialMessageHistory);
+    this.isGeneratingCallback(false);
+    this.isReadyCallback(false);
+
+    try {
+      const tokenizerFileUri = await ResourceFetcher.fetch(tokenizerSource);
+      const tokenizerConfigFileUri = await ResourceFetcher.fetch(
+        tokenizerConfigSource
+      );
+      this.tokenizerConfig = JSON.parse(
+        await readAsStringAsync('file://' + tokenizerConfigFileUri)
+      );
+
+      const modelFileUri = await ResourceFetcher.fetch(
+        modelSource,
+        this.onDownloadProgressCallback
+      );
+
+      await this.nativeModule.loadLLM(modelFileUri, tokenizerFileUri);
+      this.isReadyCallback(true);
+      this.onToken = this.nativeModule.onToken((data: string | undefined) => {
+        if (!data) {
+          return;
+        }
+        this.responseCallback(this._response + data);
+      });
+    } catch (e) {
+      this.isReadyCallback(false);
+      throw new Error(getError(e));
+    }
+  }
+
+  public configure({
+    chatConfig,
+    toolsConfig,
+  }: {
+    chatConfig?: Partial<ChatConfig>;
+    toolsConfig?: ToolsConfig;
+  }) {
+    this.chatConfig = { ...DEFAULT_CHAT_CONFIG, ...chatConfig };
+    this.toolsConfig = toolsConfig;
+
+    // reset inner state when loading new configuration
+    this.responseCallback('');
+    this.messageHistoryCallback(this.chatConfig.initialMessageHistory);
+    this.isGeneratingCallback(false);
+  }
+
+  public delete() {
+    if (this._isGenerating) {
+      throw new Error(
+        'Model is generating! You cannot delete the model now. You need to interrupt first.'
+      );
+    }
+    this.onToken?.remove();
+    this.onToken = null;
+    this.nativeModule.releaseResources();
+    this.isReadyCallback(false);
+    this.isGeneratingCallback(false);
+  }
+
+  public async forward(input: string) {
+    if (!this._isReady) {
+      throw new Error(getError(ETError.ModuleNotLoaded));
+    }
+    try {
+      this.responseCallback('');
+      this.isGeneratingCallback(true);
+      await this.nativeModule.forward(input);
+    } catch (e) {
+      throw new Error(getError(e));
+    } finally {
+      this.isGeneratingCallback(false);
+    }
+  }
+
+  public interrupt() {
+    this.nativeModule.interrupt();
+  }
+
+  public async generate(messages: Message[], tools?: LLMTool[]) {
+    if (!this._isReady) {
+      throw new Error(getError(ETError.ModuleNotLoaded));
+    }
+    if (messages.length === 0) {
+      throw new Error(`Empty 'messages' array!`);
+    }
+    if (messages[0] && messages[0].role !== 'system') {
+      console.warn(
+        `You are not providing system prompt. You can pass it in the first message using { role: 'system', content: YOUR_PROMPT }. Otherwise prompt from your model's chat template will be used.`
+      );
+    }
+
+    const renderedChat: string = this.applyChatTemplate(
+      messages,
+      this.tokenizerConfig,
+      tools,
+      // eslint-disable-next-line camelcase
+      { tools_in_user_message: false, add_generation_prompt: true }
+    );
+
+    await this.forward(renderedChat);
+
+    if (!this._response) {
+      return;
+    }
+
+    const cleanedResponse = this._response
+      .replaceAll(this.tokenizerConfig.eos_token, '')
+      .replaceAll(this.tokenizerConfig.pad_token, '');
+    this.responseCallback(cleanedResponse);
+  }
+
+  public async sendMessage(message: string) {
+    this.messageHistoryCallback([
+      ...this._messageHistory,
+      { content: message, role: 'user' },
+    ]);
+
+    const messageHistoryWithPrompt: Message[] = [
+      { content: this.chatConfig.systemPrompt, role: 'system' },
+      ...this._messageHistory.slice(-this.chatConfig.contextWindowLength),
+    ];
+
+    await this.generate(messageHistoryWithPrompt, this.toolsConfig?.tools);
+
+    if (!this.toolsConfig || this.toolsConfig.displayToolCalls) {
+      this.responseCallback(
+        this._response.replace(this.tokenizerConfig.eos_token, '')
+      );
+      this.messageHistoryCallback([
+        ...this._messageHistory,
+        { content: this._response, role: 'assistant' },
+      ]);
+    }
+    if (!this.toolsConfig) {
+      return;
+    }
+
+    const toolCalls = parseToolCall(this._response);
+
+    for (const toolCall of toolCalls) {
+      this.toolsConfig
+        .executeToolCallback(toolCall)
+        .then((toolResponse: string | null) => {
+          if (toolResponse) {
+            this.messageHistoryCallback([
+              ...this._messageHistory,
+              { content: toolResponse, role: 'assistant' },
+            ]);
+          }
+        });
+    }
+  }
+
+  public deleteMessage(index: number) {
+    // we delete referenced message and all messages after it
+    // so the model responses that used them are deleted as well
+    const newMessageHistory = this._messageHistory.slice(0, index);
+
+    this.messageHistoryCallback(newMessageHistory);
+  }
+
+  private applyChatTemplate(
+    messages: Message[],
+    tokenizerConfig: any,
+    tools?: LLMTool[],
+    templateFlags?: Object
+  ): string {
+    if (!tokenizerConfig.chat_template) {
+      throw Error("Tokenizer config doesn't include chat_template");
+    }
+    const template = new Template(tokenizerConfig.chat_template);
+
+    const specialTokens = Object.fromEntries(
+      SPECIAL_TOKENS.filter((key) => key in tokenizerConfig).map((key) => [
+        key,
+        tokenizerConfig[key],
+      ])
+    );
+
+    const result = template.render({
+      messages,
+      tools,
+      ...templateFlags,
+      ...specialTokens,
+    });
+    return result;
+  }
+}

package/src/controllers/OCRController.ts

@@ -1,15 +1,12 @@
 import { symbols } from '../constants/ocr/symbols';
 import { ETError, getError } from '../Error';
-import { _OCRModule } from '../native/RnExecutorchModules';
+import { OCRNativeModule } 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 OCRController {
-  private nativeModule: _OCRModule;
+  private nativeModule: typeof OCRNativeModule;
   public isReady: boolean = false;
   public isGenerating: boolean = false;
   public error: string | null = null;
@@ -24,7 +21,7 @@ export class OCRController {
     isGeneratingCallback = (_isGenerating: boolean) => {},
     errorCallback = (_error: string) => {},
   }) {
-    this.nativeModule = new _OCRModule();
+    this.nativeModule = OCRNativeModule;
     this.modelDownloadProgressCallback = modelDownloadProgressCallback;
     this.isReadyCallback = isReadyCallback;
     this.isGeneratingCallback = isGeneratingCallback;
@@ -47,34 +44,23 @@ export class OCRController {
       if (!symbols[language]) {
         throw new Error(getError(ETError.LanguageNotSupported));
       }
+
       this.isReady = false;
       this.isReadyCallback(false);
 
-      const detectorPath = await fetchResource(
+      const paths = await ResourceFetcher.fetchMultipleResources(
+        this.modelDownloadProgressCallback,
         detectorSource,
-        calculateDownloadProgres(4, 0, this.modelDownloadProgressCallback)
+        recognizerSources.recognizerLarge,
+        recognizerSources.recognizerMedium,
+        recognizerSources.recognizerSmall
       );
 
-      const recognizerPaths = {
-        recognizerLarge: await fetchResource(
-          recognizerSources.recognizerLarge,
-          calculateDownloadProgres(4, 1, this.modelDownloadProgressCallback)
-        ),
-        recognizerMedium: await fetchResource(
-          recognizerSources.recognizerMedium,
-          calculateDownloadProgres(4, 2, this.modelDownloadProgressCallback)
-        ),
-        recognizerSmall: await fetchResource(
-          recognizerSources.recognizerSmall,
-          calculateDownloadProgres(4, 3, this.modelDownloadProgressCallback)
-        ),
-      };
-
       await this.nativeModule.loadModule(
-        detectorPath,
-        recognizerPaths.recognizerLarge,
-        recognizerPaths.recognizerMedium,
-        recognizerPaths.recognizerSmall,
+        paths[0]!,
+        paths[1]!,
+        paths[2]!,
+        paths[3]!,
         symbols[language]
       );