From 7d624027687894705a68835a2cee239101f1fb94 Mon Sep 17 00:00:00 2001
From: Prianka Liz Kariat <prianka.kariat@ymedialabs.com>
Date: Wed, 25 Jan 2023 20:16:54 +0530
Subject: [PATCH] Added MPPEmbedding

---
 .../tasks/ios/components/containers/BUILD     |  7 ++
 .../containers/sources/MPPEmbedding.h         | 70 +++++++++++++++++++
 .../containers/sources/MPPEmbedding.m         | 34 +++++++++
 3 files changed, 111 insertions(+)
 create mode 100644 mediapipe/tasks/ios/components/containers/sources/MPPEmbedding.h
 create mode 100644 mediapipe/tasks/ios/components/containers/sources/MPPEmbedding.m

diff --git a/mediapipe/tasks/ios/components/containers/BUILD b/mediapipe/tasks/ios/components/containers/BUILD
index fb23160b8..22bbe7731 100644
--- a/mediapipe/tasks/ios/components/containers/BUILD
+++ b/mediapipe/tasks/ios/components/containers/BUILD
@@ -28,3 +28,10 @@ objc_library(
     hdrs = ["sources/MPPClassificationResult.h"],
     deps = [":MPPCategory"],
 )
+
+objc_library(
+    name = "MPPEmbedding",
+    srcs = ["sources/MPPEmbedding.m"],
+    hdrs = ["sources/MPPEmbedding.h"],
+)
+
diff --git a/mediapipe/tasks/ios/components/containers/sources/MPPEmbedding.h b/mediapipe/tasks/ios/components/containers/sources/MPPEmbedding.h
new file mode 100644
index 000000000..a9db8e579
--- /dev/null
+++ b/mediapipe/tasks/ios/components/containers/sources/MPPEmbedding.h
@@ -0,0 +1,70 @@
+// 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 Floating-point embedding.
+ * Empty if the embedder was configured to perform scalar quantization.
+ */
+@property(nonatomic, readonly, nullable) float *floatEmbedding;
+
+/**
+ * @brief The Quantized embedding.
+ * Empty if the embedder was not configured to perform scalar quantization.
+ */
+@property(nonatomic, readonly, nullable) char *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 float *)floatEmbedding
+                    quantizedEmbedding:(nullable char *)quantizedEmbedding
+                             headIndex:(NSInteger)headIndex
+                              headName:(nullable NSString *)headName NS_DESIGNATED_INITIALIZER;
+
+- (instancetype)init NS_UNAVAILABLE;
+
++ (instancetype)new NS_UNAVAILABLE;
+
+@end
+
+NS_ASSUME_NONNULL_END
diff --git a/mediapipe/tasks/ios/components/containers/sources/MPPEmbedding.m b/mediapipe/tasks/ios/components/containers/sources/MPPEmbedding.m
new file mode 100644
index 000000000..642853ef1
--- /dev/null
+++ b/mediapipe/tasks/ios/components/containers/sources/MPPEmbedding.m
@@ -0,0 +1,34 @@
+// 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 float *)floatEmbedding
+                    quantizedEmbedding:(nullable char *)quantizedEmbedding
+                             headIndex:(NSInteger)headIndex
+                              headName:(nullable NSString *)headName {
+  // TODO: Should null check for embeddings be done here ?
+  self = [super init];
+  if (self) {
+    _headIndex = headIndex;
+    _headName = headName;
+    _floatEmbedding = floatEmbedding;
+    _quantizedEmbedding = quantizedEmbedding;
+  }
+  return self;
+}
+
+@end