Merge pull request #4030 from priankakariatyml:ios-text-embedder

PiperOrigin-RevId: 506048365

mediapipe/tasks/ios/components/containers/BUILD

@@ -28,3 +28,16 @@ objc_library(
     hdrs = ["sources/MPPClassificationResult.h"],
     deps = [":MPPCategory"],
 )
+
+objc_library(
+    name = "MPPEmbedding",
+    srcs = ["sources/MPPEmbedding.m"],
+    hdrs = ["sources/MPPEmbedding.h"],
+)
+
+objc_library(
+    name = "MPPEmbeddingResult",
+    srcs = ["sources/MPPEmbeddingResult.m"],
+    hdrs = ["sources/MPPEmbeddingResult.h"],
+    deps = [":MPPEmbedding"],
+)

mediapipe/tasks/ios/components/containers/sources/MPPEmbedding.h

@@ -0,0 +1,69 @@
+// Copyright 2023 The MediaPipe Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#import <Foundation/Foundation.h>
+
+NS_ASSUME_NONNULL_BEGIN
+
+/**
+ * Represents the embedding for a given embedder head. Typically used in embedding tasks.
+ *
+ * One and only one of the two 'floatEmbedding' and 'quantizedEmbedding' will contain data, based on
+ * whether or not the embedder was configured to perform scala quantization.
+ */
+NS_SWIFT_NAME(Embedding)
+@interface MPPEmbedding : NSObject
+
+/**
+ * @brief The embedding represented as an `NSArray` of `Float` values.
+ * Empty if the embedder was configured to perform scalar quantization.
+ */
+@property(nonatomic, readonly, nullable) NSArray<NSNumber *> *floatEmbedding;
+
+/**
+ * @brief The embedding represented as an `NSArray` of `UInt8` values.
+ * Empty if the embedder was not configured to perform scalar quantization.
+ */
+@property(nonatomic, readonly, nullable) NSArray<NSNumber *> *quantizedEmbedding;
+
+/** The index of the embedder head these entries refer to. This is useful for multi-head models. */
+@property(nonatomic, readonly) NSInteger headIndex;
+
+/** The optional name of the embedder head, which is the corresponding tensor metadata name. */
+@property(nonatomic, readonly, nullable) NSString *headName;
+
+/**
+ * Initializes a new `MPPEmbedding` with the given float embedding, quantized embedding, head index
+ * and head name.
+ *
+ * @param floatEmbedding The optional Floating-point embedding.
+ * @param quantizedEmbedding The optional Quantized embedding.
+ * @param headIndex The index of the embedder head.
+ * @param headName The optional name of the embedder head.
+ *
+ * @return An instance of `MPPEmbedding` initialized with the given float embedding, quantized
+ * embedding, head index and head name.
+ */
+- (instancetype)initWithFloatEmbedding:(nullable NSArray<NSNumber *> *)floatEmbedding
+                    quantizedEmbedding:(nullable NSArray<NSNumber *> *)quantizedEmbedding
+                             headIndex:(NSInteger)headIndex
+                              headName:(nullable NSString *)headName NS_DESIGNATED_INITIALIZER;
+
+- (instancetype)init NS_UNAVAILABLE;
+
++ (instancetype)new NS_UNAVAILABLE;
+
+@end
+
+NS_ASSUME_NONNULL_END

mediapipe/tasks/ios/components/containers/sources/MPPEmbedding.m

@@ -0,0 +1,33 @@
+// Copyright 2023 The MediaPipe Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#import "mediapipe/tasks/ios/components/containers/sources/MPPEmbedding.h"
+
+@implementation MPPEmbedding
+
+- (instancetype)initWithFloatEmbedding:(nullable NSArray<NSNumber *> *)floatEmbedding
+                    quantizedEmbedding:(nullable NSArray<NSNumber *> *)quantizedEmbedding
+                             headIndex:(NSInteger)headIndex
+                              headName:(nullable NSString *)headName {
+  self = [super init];
+  if (self) {
+    _headIndex = headIndex;
+    _headName = headName;
+    _floatEmbedding = floatEmbedding;
+    _quantizedEmbedding = quantizedEmbedding;
+  }
+  return self;
+}
+
+@end

mediapipe/tasks/ios/components/containers/sources/MPPEmbeddingResult.h

@@ -0,0 +1,59 @@
+// Copyright 2023 The MediaPipe Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#import <Foundation/Foundation.h>
+#import "mediapipe/tasks/ios/components/containers/sources/MPPEmbedding.h"
+
+NS_ASSUME_NONNULL_BEGIN
+
+/** Represents the embedding results of a model. Typically used as a result for embedding tasks. */
+NS_SWIFT_NAME(EmbeddingResult)
+@interface MPPEmbeddingResult : NSObject
+
+/**
+ * An Array of `MPPEmbedding` objects containing the embedding results for each head of the model.
+ */
+@property(nonatomic, readonly) NSArray<MPPEmbedding *> *embeddings;
+
+/**
+ * @brief The optional timestamp (in milliseconds) of the start of the chunk of data corresponding
+ * to these results.
+ * This is only used for embedding extraction on time series (e.g. audio embedder). In these use
+ * cases, the amount of data to process might exceed the maximum size that the model can process. To
+ * solve this, the input data is split into multiple chunks starting at different timestamps.
+ */
+@property(nonatomic, readonly) NSInteger timestampMs;
+
+/**
+ * Initializes a new `MPPEmbedding` with the given array of embeddings and timestamp (in
+ * milliseconds).
+ *
+ * @param embeddings An Array of `MPPEmbedding` objects containing the embedding results for each
+ * head of the model.
+ * @param timestampMs The optional timestamp (in milliseconds) of the start of the chunk of data
+ * corresponding to these results. Pass `0` if timestamp is absent.
+ *
+ * @return An instance of `MPPEmbeddingResult` initialized with the given array of embeddings and
+ * timestampMs.
+ */
+- (instancetype)initWithEmbeddings:(NSArray<MPPEmbedding *> *)embeddings
+                       timestampMs:(NSInteger)timestampMs NS_DESIGNATED_INITIALIZER;
+
+- (instancetype)init NS_UNAVAILABLE;
+
++ (instancetype)new NS_UNAVAILABLE;
+
+@end
+
+NS_ASSUME_NONNULL_END

mediapipe/tasks/ios/components/containers/sources/MPPEmbeddingResult.m

@@ -0,0 +1,30 @@
+// Copyright 2023 The MediaPipe Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#import "mediapipe/tasks/ios/components/containers/sources/MPPEmbeddingResult.h"
+
+@implementation MPPEmbeddingResult
+
+- (instancetype)initWithEmbeddings:(NSArray<MPPEmbedding *> *)embeddings
+                       timestampMs:(NSInteger)timestampMs {
+  self = [super init];
+  if (self) {
+    _embeddings = embeddings;
+    _timestampMs = timestampMs;
+  }
+
+  return self;
+}
+
+@end

mediapipe/tasks/ios/text/text_classifier/sources/MPPTextClassifierOptions.h

@@ -19,7 +19,7 @@
 NS_ASSUME_NONNULL_BEGIN
 
 /**
- * Options for setting up a `MPPTextClassifierOptions`.
+ * Options for setting up a `MPPTextClassifier`.
  */
 NS_SWIFT_NAME(TextClassifierOptions)
 @interface MPPTextClassifierOptions : MPPTaskOptions <NSCopying>

mediapipe/tasks/ios/text/text_embedder/BUILD

@@ -0,0 +1,34 @@
+# Copyright 2023 The MediaPipe Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+package(default_visibility = ["//mediapipe/tasks:internal"])
+
+licenses(["notice"])
+
+objc_library(
+    name = "MPPTextEmbedderOptions",
+    srcs = ["sources/MPPTextEmbedderOptions.m"],
+    hdrs = ["sources/MPPTextEmbedderOptions.h"],
+    deps = ["//mediapipe/tasks/ios/core:MPPTaskOptions"],
+)
+
+objc_library(
+    name = "MPPTextEmbedderResult",
+    srcs = ["sources/MPPTextEmbedderResult.m"],
+    hdrs = ["sources/MPPTextEmbedderResult.h"],
+    deps = [
+        "//mediapipe/tasks/ios/components/containers:MPPEmbeddingResult",
+        "//mediapipe/tasks/ios/core:MPPTaskResult",
+    ],
+)

mediapipe/tasks/ios/text/text_embedder/sources/MPPTextEmbedderOptions.h

@@ -0,0 +1,47 @@
+// Copyright 2023 The MediaPipe Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#import <Foundation/Foundation.h>
+
+#import "mediapipe/tasks/ios/core/sources/MPPTaskOptions.h"
+
+NS_ASSUME_NONNULL_BEGIN
+
+/**
+ * Options for setting up a `MPPTextEmbedder`.
+ */
+NS_SWIFT_NAME(TextEmbedderOptions)
+@interface MPPTextEmbedderOptions : MPPTaskOptions <NSCopying>
+
+/**
+ * @brief Sets whether L2 normalization should be performed on the returned embeddings.
+ * Use this option only if the model does not already contain a native L2_NORMALIZATION TF Lite Op.
+ * In most cases, this is already the case and L2 norm is thus achieved through TF Lite inference.
+ *
+ * `NO` by default.
+ */
+@property(nonatomic) BOOL l2Normalize;
+
+/**
+ * @brief Sets whether the returned embedding should be quantized to bytes via scalar quantization.
+ * Embeddings are implicitly assumed to be unit-norm and therefore any dimensions is guaranteed to
+ * have value in [-1.0, 1.0]. Use the `l2Normalize` property if this is not the case.
+ *
+ * `NO` by default.
+ */
+@property(nonatomic) BOOL quantize;
+
+@end
+
+NS_ASSUME_NONNULL_END

mediapipe/tasks/ios/text/text_embedder/sources/MPPTextEmbedderOptions.m

@@ -0,0 +1,28 @@
+// Copyright 2023 The MediaPipe Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#import "mediapipe/tasks/ios/text/text_embedder/sources/MPPTextEmbedderOptions.h"
+
+@implementation MPPTextEmbedderOptions
+
+- (id)copyWithZone:(NSZone *)zone {
+  MPPTextEmbedderOptions *textEmbedderOptions = [super copyWithZone:zone];
+
+  textEmbedderOptions.l2Normalize = self.l2Normalize;
+  textEmbedderOptions.quantize = self.quantize;
+
+  return textEmbedderOptions;
+}
+
+@end

mediapipe/tasks/ios/text/text_embedder/sources/MPPTextEmbedderResult.h

@@ -0,0 +1,48 @@
+// Copyright 2023 The MediaPipe Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#import <Foundation/Foundation.h>
+#import "mediapipe/tasks/ios/components/containers/sources/MPPEmbeddingResult.h"
+#import "mediapipe/tasks/ios/core/sources/MPPTaskResult.h"
+
+NS_ASSUME_NONNULL_BEGIN
+
+/** Represents the embedding results generated by `MPPTextEmbedder`. **/
+NS_SWIFT_NAME(TextEmbedderResult)
+@interface MPPTextEmbedderResult : MPPTaskResult
+
+/** The `MPPEmbedderResult` instance containing one embedding per embedder head. **/
+@property(nonatomic, readonly) MPPEmbeddingResult *embeddingResult;
+
+/**
+ * Initializes a new `MPPTextEmbedderResult` with the given `MPPEmbeddingResult` and
+ * timestamp (in milliseconds).
+ *
+ * @param embeddingResult The `MPPEmbeddingResult` instance containing one set of results
+ * per classifier head.
+ * @param timestampMs The timestamp for this result.
+ *
+ * @return An instance of `MPPTextEmbedderResult` initialized with the given
+ * `MPPEmbeddingResult` and timestamp (in milliseconds).
+ */
+- (instancetype)initWithEmbeddingResult:(MPPEmbeddingResult *)embeddingResult
+                            timestampMs:(NSInteger)timestampMs;
+
+- (instancetype)init NS_UNAVAILABLE;
+
++ (instancetype)new NS_UNAVAILABLE;
+
+@end
+
+NS_ASSUME_NONNULL_END

mediapipe/tasks/ios/text/text_embedder/sources/MPPTextEmbedderResult.m

@@ -0,0 +1,28 @@
+// Copyright 2023 The MediaPipe Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#import "mediapipe/tasks/ios/text/text_embedder/sources/MPPTextEmbedderResult.h"
+
+@implementation MPPTextEmbedderResult
+
+- (instancetype)initWithEmbeddingResult:(MPPEmbeddingResult *)embeddingResult
+                            timestampMs:(NSInteger)timestampMs {
+  self = [super initWithTimestampMs:timestampMs];
+  if (self) {
+    _embeddingResult = embeddingResult;
+  }
+  return self;
+}
+
+@end