refactor: ImageEmbedding and TextEmbedding, remove normalization, update URLs (#445)

## Description

Changes:
- Refactored CLIP URLs
- Updated textEmbeddings models URLs to v0.5.0
- Refactored textEmbeddings demo
- Removed normalization from ImageEmbeddings
- Refactored ImageEmbeddings c++ code
- Removed attentionMask from TextEmbeddings postProcessing

### Type of change

- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [x] 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

### Checklist

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

Author: jakmro

apps/text-embeddings/app/clip-embeddings/index.tsx

@@ -14,11 +14,12 @@ import { Ionicons } from '@expo/vector-icons';
 import {
   useTextEmbeddings,
   useImageEmbeddings,
-  CLIP_VIT_BASE_PATCH_32_TEXT_ENCODER,
-  CLIP_VIT_BASE_PATCH_32_IMAGE_ENCODER_MODEL,
+  CLIP_VIT_BASE_PATCH32_TEXT,
+  CLIP_VIT_BASE_PATCH32_IMAGE,
 } from 'react-native-executorch';
 import { launchImageLibrary } from 'react-native-image-picker';
 import { useIsFocused } from '@react-navigation/native';
+import { dotProduct } from '../../utils/math';
 
 export default function ClipEmbeddingsScreenWrapper() {
   const isFocused = useIsFocused();
@@ -27,11 +28,8 @@ export default function ClipEmbeddingsScreenWrapper() {
 }
 
 function ClipEmbeddingsScreen() {
-  const model = useTextEmbeddings({ ...CLIP_VIT_BASE_PATCH_32_TEXT_ENCODER });
-
-  const imageModel = useImageEmbeddings({
-    modelSource: CLIP_VIT_BASE_PATCH_32_IMAGE_ENCODER_MODEL,
-  });
+  const textModel = useTextEmbeddings(CLIP_VIT_BASE_PATCH32_TEXT);
+  const imageModel = useImageEmbeddings(CLIP_VIT_BASE_PATCH32_IMAGE);
 
   const [inputSentence, setInputSentence] = useState('');
   const [sentencesWithEmbeddings, setSentencesWithEmbeddings] = useState<
@@ -41,18 +39,10 @@ function ClipEmbeddingsScreen() {
     { sentence: string; similarity: number }[]
   >([]);
 
-  const dotProduct = (a: Float32Array, b: Float32Array) => {
-    let sum = 0;
-    for (let i = 0; i < a.length; i++) {
-      sum += a[i] * b[i];
-    }
-    return sum;
-  };
-
   useEffect(
     () => {
       const computeEmbeddings = async () => {
-        if (!model.isReady) return;
+        if (!textModel.isReady) return;
 
         const sentences = [
           'The weather is lovely today.',
@@ -64,7 +54,7 @@ function ClipEmbeddingsScreen() {
         try {
           const embeddings = [];
           for (const sentence of sentences) {
-            const embedding = await model.forward(sentence);
+            const embedding = await textModel.forward(sentence);
             embeddings.push({ sentence, embedding });
           }
           setSentencesWithEmbeddings(embeddings);
@@ -76,14 +66,14 @@ function ClipEmbeddingsScreen() {
       computeEmbeddings();
     },
     // eslint-disable-next-line react-hooks/exhaustive-deps
-    [model.isReady]
+    [textModel.isReady]
   );
 
   const checkSimilarities = async () => {
-    if (!model.isReady || !inputSentence.trim()) return;
+    if (!textModel.isReady || !inputSentence.trim()) return;
 
     try {
-      const inputEmbedding = await model.forward(inputSentence);
+      const inputEmbedding = await textModel.forward(inputSentence);
       const matches = sentencesWithEmbeddings.map(
         ({ sentence, embedding }) => ({
           sentence,
@@ -98,10 +88,10 @@ function ClipEmbeddingsScreen() {
   };
 
   const addToSentences = async () => {
-    if (!model.isReady || !inputSentence.trim()) return;
+    if (!textModel.isReady || !inputSentence.trim()) return;
 
     try {
-      const embedding = await model.forward(inputSentence);
+      const embedding = await textModel.forward(inputSentence);
       setSentencesWithEmbeddings((prev) => [
         ...prev,
         { sentence: inputSentence, embedding },
@@ -115,7 +105,7 @@ function ClipEmbeddingsScreen() {
   };
 
   const clearList = async () => {
-    if (!model.isReady) return;
+    if (!textModel.isReady) return;
     try {
       setSentencesWithEmbeddings([]);
     } catch (error) {
@@ -149,16 +139,14 @@ function ClipEmbeddingsScreen() {
     }
   };
 
-  const getModelStatusText = () => {
-    if (model.error || imageModel.error) {
-      return `Oops! Error: ${model.error || imageModel.error}`;
+  const getModelStatusText = (model: typeof textModel | typeof imageModel) => {
+    if (model.error) {
+      return `Oops! Error: ${model.error}`;
     }
-    if (!model.isReady || !imageModel.isReady) {
-      return `Loading model ${(((model.downloadProgress + imageModel.downloadProgress) / 2) * 100).toFixed(2)}%`;
+    if (!model.isReady) {
+      return `Loading model ${(model.downloadProgress * 100).toFixed(2)}%`;
     }
-    return model.isGenerating || imageModel.isGenerating
-      ? 'Generating...'
-      : 'Model is ready';
+    return model.isGenerating ? 'Generating...' : 'Model is ready';
   };
 
   return (
@@ -169,8 +157,12 @@ function ClipEmbeddingsScreen() {
       >
         <ScrollView contentContainerStyle={styles.scrollContainer}>
           <Text style={styles.heading}>Text Embeddings Playground</Text>
-          <Text style={styles.sectionTitle}>{getModelStatusText()}</Text>
-
+          <Text style={styles.sectionTitle}>
+            Text Model: {getModelStatusText(textModel)}
+          </Text>
+          <Text style={styles.sectionTitle}>
+            Image Model: {getModelStatusText(imageModel)}
+          </Text>
           <View style={styles.card}>
             <Text style={styles.sectionTitle}>List of Existing Sentences</Text>
             {sentencesWithEmbeddings.map((item, index) => (

apps/text-embeddings/app/text-embeddings/index.tsx

@@ -17,6 +17,7 @@ import {
   ALL_MINILM_L6_V2_TOKENIZER,
 } from 'react-native-executorch';
 import { useIsFocused } from '@react-navigation/native';
+import { dotProduct } from '../../utils/math';
 
 export default function TextEmbeddingsScreenWrapper() {
   const isFocused = useIsFocused();
@@ -38,14 +39,6 @@ function TextEmbeddingsScreen() {
     { sentence: string; similarity: number }[]
   >([]);
 
-  const dotProduct = (a: Float32Array, b: Float32Array) => {
-    let sum = 0;
-    for (let i = 0; i < a.length; i++) {
-      sum += a[i] * b[i];
-    }
-    return sum;
-  };
-
   useEffect(
     () => {
       const computeEmbeddings = async () => {

apps/text-embeddings/utils/math.ts

@@ -0,0 +1,11 @@
+export const dotProduct = (a: Float32Array, b: Float32Array) => {
+  if (a.length !== b.length) {
+    throw new Error('Vectors must be of the same length');
+  }
+
+  let sum = 0;
+  for (let i = 0; i < a.length; i++) {
+    sum += a[i] * b[i];
+  }
+  return sum;
+};

packages/react-native-executorch/common/rnexecutorch/RnExecutorchInstaller.cpp

@@ -3,13 +3,13 @@
 #include <rnexecutorch/TokenizerModule.h>
 #include <rnexecutorch/host_objects/JsiConversions.h>
 #include <rnexecutorch/models/classification/Classification.h>
-#include <rnexecutorch/models/image_embeddings/ImageEmbeddings.h>
+#include <rnexecutorch/models/embeddings/image/ImageEmbeddings.h>
+#include <rnexecutorch/models/embeddings/text/TextEmbeddings.h>
 #include <rnexecutorch/models/image_segmentation/ImageSegmentation.h>
 #include <rnexecutorch/models/llm/LLM.h>
 #include <rnexecutorch/models/object_detection/ObjectDetection.h>
 #include <rnexecutorch/models/speech_to_text/SpeechToText.h>
 #include <rnexecutorch/models/style_transfer/StyleTransfer.h>
-#include <rnexecutorch/models/text_embeddings/TextEmbeddings.h>
 
 namespace rnexecutorch {
 

packages/react-native-executorch/common/rnexecutorch/models/embeddings/BaseEmbeddings.cpp

@@ -0,0 +1,27 @@
+#include "BaseEmbeddings.h"
+
+#include <span>
+
+namespace rnexecutorch {
+
+BaseEmbeddings::BaseEmbeddings(const std::string &modelSource,
+                               std::shared_ptr<react::CallInvoker> callInvoker)
+    : BaseModel(modelSource, callInvoker) {}
+
+std::shared_ptr<OwningArrayBuffer>
+BaseEmbeddings::postprocess(const Result<std::vector<EValue>> &forwardResult) {
+  auto forwardResultTensor = forwardResult->at(0).toTensor();
+  auto dataPtr = forwardResultTensor.mutable_data_ptr();
+  auto outputNumel = forwardResultTensor.numel();
+
+  std::span<float> modelOutput(static_cast<float *>(dataPtr), outputNumel);
+
+  auto createBuffer = [](const auto &data, size_t size) {
+    auto buffer = std::make_shared<OwningArrayBuffer>(size);
+    std::memcpy(buffer->data(), data, size);
+    return buffer;
+  };
+  return createBuffer(modelOutput.data(), modelOutput.size_bytes());
+}
+
+} // namespace rnexecutorch

packages/react-native-executorch/common/rnexecutorch/models/embeddings/BaseEmbeddings.h

@@ -0,0 +1,17 @@
+#pragma once
+
+#include <rnexecutorch/models/BaseModel.h>
+
+namespace rnexecutorch {
+
+class BaseEmbeddings : public BaseModel {
+public:
+  BaseEmbeddings(const std::string &modelSource,
+                 std::shared_ptr<react::CallInvoker> callInvoker);
+
+protected:
+  std::shared_ptr<OwningArrayBuffer>
+  postprocess(const Result<std::vector<EValue>> &forwardResult);
+};
+
+}; // namespace rnexecutorch

…els/image_embeddings/ImageEmbeddings.cpp …els/embeddings/image/ImageEmbeddings.cpppackages/react-native-executorch/common/rnexecutorch/models/image_embeddings/ImageEmbeddings.cpp renamed to packages/react-native-executorch/common/rnexecutorch/models/embeddings/image/ImageEmbeddings.cpp packages/react-native-executorch/common/rnexecutorch/models/image_embeddings/ImageEmbeddings.cpp renamed to packages/react-native-executorch/common/rnexecutorch/models/embeddings/image/ImageEmbeddings.cpp

@@ -3,12 +3,13 @@
 #include <executorch/extension/tensor/tensor.h>
 #include <rnexecutorch/data_processing/ImageProcessing.h>
 #include <rnexecutorch/data_processing/Numerical.h>
+
 namespace rnexecutorch {
 
 ImageEmbeddings::ImageEmbeddings(
     const std::string &modelSource,
     std::shared_ptr<react::CallInvoker> callInvoker)
-    : BaseModel(modelSource, callInvoker) {
+    : BaseEmbeddings(modelSource, callInvoker) {
   auto inputTensors = getAllInputShapes();
   if (inputTensors.size() == 0) {
     throw std::runtime_error("Model seems to not take any input tensors.");
@@ -31,30 +32,14 @@ ImageEmbeddings::generate(std::string imageSource) {
   auto [inputTensor, originalSize] =
       imageprocessing::readImageToTensor(imageSource, getAllInputShapes()[0]);
 
-  auto result = BaseModel::forward(inputTensor);
-  if (!result.ok()) {
-    throw std::runtime_error("Forward pass failed: Error " +
-                             std::to_string(static_cast<int>(result.error())));
+  auto forwardResult = BaseModel::forward(inputTensor);
+  if (!forwardResult.ok()) {
+    throw std::runtime_error(
+        "Function forward in ImageEmbeddings failed with error code: " +
+        std::to_string(static_cast<uint32_t>(forwardResult.error())));
   }
 
-  auto &outputs = result.get();
-
-  if (outputs.size() > 1) {
-    throw std::runtime_error("It returned multiple outputs!");
-  }
-
-  auto &outputTensor = outputs.at(0).toTensor();
-  std::span<float> outputTensorSpan(
-      static_cast<float *>(outputTensor.mutable_data_ptr()),
-      outputTensor.numel());
-
-  numerical::normalize(outputTensorSpan);
-
-  size_t bufferSize = outputTensorSpan.size_bytes();
-  auto buffer = std::make_shared<OwningArrayBuffer>(bufferSize);
-
-  std::memcpy(buffer->data(), outputTensorSpan.data(), bufferSize);
-
-  return buffer;
+  return BaseEmbeddings::postprocess(forwardResult);
 }
+
 } // namespace rnexecutorch

…odels/image_embeddings/ImageEmbeddings.h …odels/embeddings/image/ImageEmbeddings.hpackages/react-native-executorch/common/rnexecutorch/models/image_embeddings/ImageEmbeddings.h renamed to packages/react-native-executorch/common/rnexecutorch/models/embeddings/image/ImageEmbeddings.h packages/react-native-executorch/common/rnexecutorch/models/image_embeddings/ImageEmbeddings.h renamed to packages/react-native-executorch/common/rnexecutorch/models/embeddings/image/ImageEmbeddings.h

@@ -4,13 +4,13 @@
 #include <executorch/runtime/core/evalue.h>
 #include <opencv2/opencv.hpp>
 
-#include <rnexecutorch/models/BaseModel.h>
+#include <rnexecutorch/models/embeddings/BaseEmbeddings.h>
 
 namespace rnexecutorch {
 using executorch::extension::TensorPtr;
 using executorch::runtime::EValue;
 
-class ImageEmbeddings : public BaseModel {
+class ImageEmbeddings final : public BaseEmbeddings {
 public:
   ImageEmbeddings(const std::string &modelSource,
                   std::shared_ptr<react::CallInvoker> callInvoker);
@@ -20,4 +20,4 @@ class ImageEmbeddings : public BaseModel {
   cv::Size modelImageSize{0, 0};
 };
 
-} // namespace rnexecutorch
+} // namespace rnexecutorch

…odels/text_embeddings/TextEmbeddings.cpp …odels/embeddings/text/TextEmbeddings.cpppackages/react-native-executorch/common/rnexecutorch/models/text_embeddings/TextEmbeddings.cpp renamed to packages/react-native-executorch/common/rnexecutorch/models/embeddings/text/TextEmbeddings.cpp packages/react-native-executorch/common/rnexecutorch/models/text_embeddings/TextEmbeddings.cpp renamed to packages/react-native-executorch/common/rnexecutorch/models/embeddings/text/TextEmbeddings.cpp

@@ -9,7 +9,7 @@ using namespace executorch::extension;
 TextEmbeddings::TextEmbeddings(const std::string &modelSource,
                                const std::string &tokenizerSource,
                                std::shared_ptr<react::CallInvoker> callInvoker)
-    : BaseModel(modelSource, callInvoker),
+    : BaseEmbeddings(modelSource, callInvoker),
       tokenizer(
           std::make_unique<TokenizerModule>(tokenizerSource, callInvoker)) {}
 
@@ -36,42 +36,26 @@ TokenIdsWithAttentionMask TextEmbeddings::preprocess(const std::string &input) {
 std::shared_ptr<OwningArrayBuffer>
 TextEmbeddings::generate(const std::string input) {
   auto preprocessed = preprocess(input);
+
   std::vector<int32_t> tokenIdsShape = {
       1, static_cast<int32_t>(preprocessed.inputIds.size())};
   std::vector<int32_t> attnMaskShape = {
       1, static_cast<int32_t>(preprocessed.attentionMask.size())};
+
   auto tokenIds = make_tensor_ptr(tokenIdsShape, preprocessed.inputIds.data(),
                                   ScalarType::Long);
   auto attnMask = make_tensor_ptr(
       attnMaskShape, preprocessed.attentionMask.data(), ScalarType::Long);
+
   auto forwardResult = BaseModel::forward({tokenIds, attnMask});
+
   if (!forwardResult.ok()) {
     throw std::runtime_error(
-        "Failed to forward, error: " +
+        "Function forward in TextEmbeddings failed with error code: " +
         std::to_string(static_cast<uint32_t>(forwardResult.error())));
   }
 
-  auto forwardResultTensor = forwardResult->at(0).toTensor();
-  auto dataPtr = forwardResultTensor.mutable_data_ptr();
-  auto outputNumel = forwardResultTensor.numel();
-
-  std::span<float> modelOutputSpan(static_cast<float *>(dataPtr), outputNumel);
-  std::span<const int64_t> attnMaskSpan(preprocessed.attentionMask.data(),
-                                        preprocessed.attentionMask.size());
-
-  return postprocess(modelOutputSpan, attnMaskSpan);
-}
-
-std::shared_ptr<OwningArrayBuffer>
-TextEmbeddings::postprocess(std::span<float> modelOutput,
-                            std::span<const int64_t> attnMask) {
-  auto createBuffer = [](const auto &data, size_t size) {
-    auto buffer = std::make_shared<OwningArrayBuffer>(size);
-    std::memcpy(buffer->data(), data, size);
-    return buffer;
-  };
-
-  return createBuffer(modelOutput.data(), modelOutput.size_bytes());
+  return BaseEmbeddings::postprocess(forwardResult);
 }
 
-} // namespace rnexecutorch
+} // namespace rnexecutorch

…/models/text_embeddings/TextEmbeddings.h …/models/embeddings/text/TextEmbeddings.hpackages/react-native-executorch/common/rnexecutorch/models/text_embeddings/TextEmbeddings.h renamed to packages/react-native-executorch/common/rnexecutorch/models/embeddings/text/TextEmbeddings.h packages/react-native-executorch/common/rnexecutorch/models/text_embeddings/TextEmbeddings.h renamed to packages/react-native-executorch/common/rnexecutorch/models/embeddings/text/TextEmbeddings.h

@@ -1,8 +1,7 @@
 #pragma once
 
 #include <rnexecutorch/TokenizerModule.h>
-#include <rnexecutorch/models/BaseModel.h>
-#include <span>
+#include <rnexecutorch/models/embeddings/BaseEmbeddings.h>
 
 namespace rnexecutorch {
 
@@ -11,7 +10,7 @@ struct TokenIdsWithAttentionMask {
   std::vector<int64_t> attentionMask;
 };
 
-class TextEmbeddings : public BaseModel {
+class TextEmbeddings final : public BaseEmbeddings {
 public:
   TextEmbeddings(const std::string &modelSource,
                  const std::string &tokenizerSource,
@@ -21,9 +20,7 @@ class TextEmbeddings : public BaseModel {
 private:
   std::vector<std::vector<int32_t>> inputShapes;
   TokenIdsWithAttentionMask preprocess(const std::string &input);
-  std::shared_ptr<OwningArrayBuffer>
-  postprocess(std::span<float> modelOutput,
-              std::span<const int64_t> attentionMask);
   std::unique_ptr<TokenizerModule> tokenizer;
 };
+
 }; // namespace rnexecutorch