feat: Flag for mean pooling text embeddings (#384)

## Description

Some models do mean pooling, some don't. We can set flag to take that
into account - we should set this flag in all presets for our models
while completing:
#359

### Type of change

- [ ] 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

- [x] iOS
- [x] Android

### Related issues

#353 

### Checklist

- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have updated the documentation accordingly
- [ ] My changes generate no new warnings **They actually do, if users
don't set the flag**

### Additional notes

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

Author: pweglik

docs/docs/natural-language-processing/useTextEmbeddings.md

@@ -33,6 +33,7 @@ import {
 const model = useTextEmbeddings({
   modelSource: ALL_MINILM_L6_V2,
   tokenizerSource: ALL_MINILM_L6_V2_TOKENIZER,
+  meanPooling: true,
 });
 
 try {
@@ -50,6 +51,8 @@ A string that specifies the location of the model binary. For more information,
 **`tokenizerSource`**
 A string that specifies the location of the tokenizer JSON file.
 
+**`meanPooling?`** - Boolean that controls whether we perform mean pooling on the model output or not. If not set, it will default to true and display warning.
+
 **`preventLoad?`** - Boolean that can prevent automatic model loading (and downloading the data if you load it for the first time) after running the hook.
 
 ### Returns
@@ -86,6 +89,7 @@ function App() {
   const model = useTextEmbeddings({
     modelSource: ALL_MINILM_L6_V2,
     tokenizerSource: ALL_MINILM_L6_V2_TOKENIZER,
+    meanPooling: true,
   });
 
   ...

docs/docs/typescript-api/TextEmbeddingsModule.md

@@ -22,11 +22,11 @@ const embedding = await TextEmbeddingsModule.forward('Hello World!');
 
 ### Methods
 
-| Method               | Type                                                                            | Description                                                                                                                                                                                |
-| -------------------- | ------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
-| `load`               | `(modelSource: ResourceSource, tokenizerSource: ResourceSource): Promise<void>` | Loads the model, where `modelSource` is a string that specifies the location of the model binary and `tokenizerSource` is a string that specifies the location of the tokenizer JSON file. |
-| `forward`            | `(input: string): Promise<number[]>`                                            | Executes the model's forward pass, where `input` is a text that will be embedded.                                                                                                          |
-| `onDownloadProgress` | `(callback: (downloadProgress: number) => void): any`                           | Subscribe to the download progress event.                                                                                                                                                  |
+| Method               | Type                                                                                                   | Description                                                                                                                                                                                                                                                  |
+| -------------------- | ------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| `load`               | `(modelSource: ResourceSource, tokenizerSource: ResourceSource, meanPooling?: boolean): Promise<void>` | Loads the model, where `modelSource` is a string that specifies the location of the model binary, `tokenizerSource` is a string that specifies the location of the tokenizer JSON file, and `meanPooling` controls when to perform pooling on model outputs. |
+| `forward`            | `(input: string): Promise<number[]>`                                                                   | Executes the model's forward pass, where `input` is a text that will be embedded.                                                                                                                                                                            |
+| `onDownloadProgress` | `(callback: (downloadProgress: number) => void): any`                                                  | Subscribe to the download progress event.                                                                                                                                                                                                                    |
 
 <details>
 <summary>Type definitions</summary>
@@ -39,7 +39,7 @@ type ResourceSource = string | number | object;
 
 ## Loading the model
 
-To load the model, use the `load` method. It accepts the `modelSource` which is a string that specifies the location of the model binary and `tokenizerSource` which is a string that specifies the location of the tokenizer JSON file. For more information, take a look at [loading models](../fundamentals/loading-models.md) page. This method returns a promise, which can resolve to an error or void.
+To load the model, use the `load` method. It accepts the `modelSource` which is a string that specifies the location of the model binary, `tokenizerSource` which is a string that specifies the location of the tokenizer JSON file, and optional `meanPooling` flag controls when to perform pooling on model outputs. For more information, take a look at [loading models](../fundamentals/loading-models.md) page. This method returns a promise, which can resolve to an error or void.
 
 ## Running the model
 

packages/react-native-executorch/android/src/main/java/com/swmansion/rnexecutorch/TextEmbeddings.kt

@@ -34,10 +34,11 @@ class TextEmbeddings(
 
   override fun forward(
     input: String,
+    meanPooling: Boolean,
     promise: Promise,
   ) {
     try {
-      val output = textEmbeddingsModel.runModel(input)
+      val output = textEmbeddingsModel.runModel(input, meanPooling)
       val writableArray = WritableNativeArray()
       output.forEach { writableArray.pushDouble(it) }
 

packages/react-native-executorch/android/src/main/java/com/swmansion/rnexecutorch/models/TextEmbeddings/TextEmbeddingsModel.kt

@@ -24,13 +24,24 @@ class TextEmbeddingsModel(
   fun postprocess(
     modelOutput: FloatArray, // [tokens * embedding_dim]
     attentionMask: LongArray, // [tokens]
+    meanPooling: Boolean,
   ): DoubleArray {
-    val modelOutputDouble = modelOutput.map { it.toDouble() }.toDoubleArray()
-    val embeddings = TextEmbeddingsUtils.meanPooling(modelOutputDouble, attentionMask)
-    return TextEmbeddingsUtils.normalize(embeddings)
+    var embeddings = modelOutput.map { it.toDouble() }.toDoubleArray()
+    if (meanPooling) {
+      embeddings = TextEmbeddingsUtils.meanPooling(embeddings, attentionMask)
+    }
+    embeddings = TextEmbeddingsUtils.normalize(embeddings)
+    return embeddings
   }
 
   override fun runModel(input: String): DoubleArray {
+    return runModel(input, true)
+  }
+
+  fun runModel(
+    input: String,
+    meanPooling: Boolean,
+  ): DoubleArray {
     val modelInput = preprocess(input)
     val inputsIds = modelInput[0]
     val attentionMask = modelInput[1]
@@ -43,6 +54,6 @@ class TextEmbeddingsModel(
 
     val modelOutput = forward(inputIdsEValue, attentionMaskEValue)[0].toTensor().dataAsFloatArray
 
-    return postprocess(modelOutput, attentionMask)
+    return postprocess(modelOutput, attentionMask, meanPooling)
   }
 }

packages/react-native-executorch/ios/RnExecutorch/TextEmbeddings.mm

@@ -41,10 +41,11 @@ - (void)loadModule:(NSString *)modelSource
 }
 
 - (void)forward:(NSString *)input
+    meanPooling:(bool)meanPooling
         resolve:(RCTPromiseResolveBlock)resolve
          reject:(RCTPromiseRejectBlock)reject {
   @try {
-    resolve([model runModel:input]);
+    resolve([model runModel:input meanPooling:meanPooling]);
     return;
   } @catch (NSException *exception) {
     NSLog(@"An exception occurred: %@, %@", exception.name, exception.reason);

packages/react-native-executorch/ios/RnExecutorch/models/text_embeddings/TextEmbeddingsModel.h

@@ -9,6 +9,7 @@
 - (void)loadTokenizer:(NSString *)tokenizerSource;
 - (NSArray *)preprocess:(NSString *)input;
 - (NSArray *)runModel:(NSString *)input;
+- (NSArray *)runModel:(NSString *)input meanPooling:(bool)meanPooling;
 - (NSArray *)postprocess:(NSArray *)input
            attentionMask:(NSArray *)attentionMask;
 

packages/react-native-executorch/ios/RnExecutorch/models/text_embeddings/TextEmbeddingsModel.mm

@@ -14,13 +14,21 @@ - (NSArray *)preprocess:(NSString *)input {
 
 - (NSArray *)postprocess:(NSArray *)modelOutput   // [tokens * embedding_dim]
            attentionMask:(NSArray *)attentionMask // [tokens]
-{
-  NSArray *embeddings = [TextEmbeddingsUtils meanPooling:modelOutput
-                                           attentionMask:attentionMask];
-  return [TextEmbeddingsUtils normalize:embeddings];
+             meanPooling:(bool)meanPooling {
+  NSArray *embeddings = modelOutput;
+  if (meanPooling) {
+    embeddings = [TextEmbeddingsUtils meanPooling:modelOutput
+                                    attentionMask:attentionMask];
+  }
+  embeddings = [TextEmbeddingsUtils normalize:embeddings];
+  return embeddings;
 }
 
 - (NSArray *)runModel:(NSString *)input {
+  return [self runModel:input meanPooling:true];
+}
+
+- (NSArray *)runModel:(NSString *)input meanPooling:(bool)meanPooling {
   NSArray *modelInput = [self preprocess:input];
 
   NSMutableArray *inputTypes = [NSMutableArray arrayWithObjects:@4, @4, nil];
@@ -34,7 +42,9 @@ - (NSArray *)runModel:(NSString *)input {
   NSArray *modelOutput = [self forward:modelInput
                                 shapes:shapes
                             inputTypes:inputTypes];
-  return [self postprocess:modelOutput[0] attentionMask:modelInput[1]];
+  return [self postprocess:modelOutput[0]
+             attentionMask:modelInput[1]
+               meanPooling:meanPooling];
 }
 
 - (void)loadTokenizer:(NSString *)tokenizerSource {

packages/react-native-executorch/src/hooks/natural_language_processing/useTextEmbeddings.ts

@@ -5,14 +5,16 @@ import { useModule } from '../useModule';
 export const useTextEmbeddings = ({
   modelSource,
   tokenizerSource,
+  meanPooling,
   preventLoad = false,
 }: {
   modelSource: ResourceSource;
   tokenizerSource: ResourceSource;
+  meanPooling?: boolean;
   preventLoad?: boolean;
 }) =>
   useModule({
     module: TextEmbeddingsModule,
-    loadArgs: [modelSource, tokenizerSource],
+    loadArgs: [modelSource, tokenizerSource, meanPooling],
     preventLoad,
   });

packages/react-native-executorch/src/modules/natural_language_processing/TextEmbeddingsModule.ts

@@ -4,15 +4,25 @@ import { BaseModule } from '../BaseModule';
 
 export class TextEmbeddingsModule extends BaseModule {
   protected static override nativeModule = TextEmbeddingsNativeModule;
+  private static meanPooling: boolean;
 
   static override async load(
     modelSource: ResourceSource,
-    tokenizerSource: ResourceSource
+    tokenizerSource: ResourceSource,
+    meanPooling?: boolean
   ) {
+    if (meanPooling === undefined) {
+      console.warn(
+        "You haven't passed meanPooling flag. It is defaulting to true. If your model doesn't require pooling it may misbehave."
+      );
+      meanPooling = true;
+    }
+
     await super.load([modelSource, tokenizerSource]);
+    this.meanPooling = meanPooling;
   }
 
   static override async forward(input: string): Promise<number[]> {
-    return this.nativeModule.forward(input);
+    return this.nativeModule.forward(input, this.meanPooling);
   }
 }

packages/react-native-executorch/src/native/NativeTextEmbeddings.ts

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