Added files for the TextEmbedder C API and tests
This commit is contained in:
Kinar
@@ -0,0 +1,85 @@
|
||||
# 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.
|
||||
|
||||
package(default_visibility = ["//mediapipe/tasks:internal"])
|
||||
|
||||
licenses(["notice"])
|
||||
|
||||
cc_library(
|
||||
name = "text_embedder_lib",
|
||||
srcs = ["text_embedder.cc"],
|
||||
hdrs = ["text_embedder.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/tasks/c/components/containers:embedding_result",
|
||||
"//mediapipe/tasks/c/components/containers:embedding_result_converter",
|
||||
"//mediapipe/tasks/c/components/processors:embedder_options",
|
||||
"//mediapipe/tasks/c/components/processors:embedder_options_converter",
|
||||
"//mediapipe/tasks/c/core:base_options",
|
||||
"//mediapipe/tasks/c/core:base_options_converter",
|
||||
"//mediapipe/tasks/cc/text/text_embedder",
|
||||
"@com_google_absl//absl/log:absl_log",
|
||||
"@com_google_absl//absl/status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
# bazel build -c opt --linkopt -s --strip always --define MEDIAPIPE_DISABLE_GPU=1 \
|
||||
# //mediapipe/tasks/c/text/text_embedder:libtext_embedder.so
|
||||
cc_binary(
|
||||
name = "libtext_embedder.so",
|
||||
linkopts = [
|
||||
"-Wl,-soname=libtext_embedder.so",
|
||||
"-fvisibility=hidden",
|
||||
],
|
||||
linkshared = True,
|
||||
tags = [
|
||||
"manual",
|
||||
"nobuilder",
|
||||
"notap",
|
||||
],
|
||||
deps = [":text_embedder_lib"],
|
||||
)
|
||||
|
||||
# bazel build --config darwin_arm64 -c opt --strip always --define MEDIAPIPE_DISABLE_GPU=1 \
|
||||
# //mediapipe/tasks/c/text/text_embedder:libtext_embedder.dylib
|
||||
cc_binary(
|
||||
name = "libtext_embedder.dylib",
|
||||
linkopts = [
|
||||
"-Wl,-install_name,libtext_embedder.dylib",
|
||||
"-fvisibility=hidden",
|
||||
],
|
||||
linkshared = True,
|
||||
tags = [
|
||||
"manual",
|
||||
"nobuilder",
|
||||
"notap",
|
||||
],
|
||||
deps = [":text_embedder_lib"],
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "text_embedder_test",
|
||||
srcs = ["text_embedder_test.cc"],
|
||||
data = ["//mediapipe/tasks/testdata/text:mobilebert_embedding_model"],
|
||||
linkstatic = 1,
|
||||
deps = [
|
||||
":text_embedder_lib",
|
||||
"//mediapipe/framework/deps:file_path",
|
||||
"//mediapipe/framework/port:gtest",
|
||||
"@com_google_absl//absl/flags:flag",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@com_google_googletest//:gtest_main",
|
||||
],
|
||||
)
|
||||
@@ -0,0 +1,124 @@
|
||||
/* 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.
|
||||
==============================================================================*/
|
||||
|
||||
#include "mediapipe/tasks/c/text/text_embedder/text_embedder.h"
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include "absl/log/absl_log.h"
|
||||
#include "absl/status/status.h"
|
||||
#include "mediapipe/tasks/c/components/containers/embedding_result_converter.h"
|
||||
#include "mediapipe/tasks/c/components/processors/embedder_options.h"
|
||||
#include "mediapipe/tasks/c/components/processors/embedder_options_converter.h"
|
||||
#include "mediapipe/tasks/c/core/base_options.h"
|
||||
#include "mediapipe/tasks/c/core/base_options_converter.h"
|
||||
#include "mediapipe/tasks/cc/text/text_embedder/text_embedder.h"
|
||||
|
||||
namespace mediapipe::tasks::c::text::text_embedder {
|
||||
|
||||
namespace {
|
||||
|
||||
using ::mediapipe::tasks::c::components::containers::
|
||||
CppCloseEmbeddingResult;
|
||||
using ::mediapipe::tasks::c::components::containers::
|
||||
CppConvertToEmbeddingResult;
|
||||
using ::mediapipe::tasks::c::components::processors::
|
||||
CppConvertToEmbedderOptions;
|
||||
using ::mediapipe::tasks::c::core::CppConvertToBaseOptions;
|
||||
using ::mediapipe::tasks::text::text_embedder::TextEmbedder;
|
||||
|
||||
int CppProcessError(absl::Status status, char** error_msg) {
|
||||
if (error_msg) {
|
||||
*error_msg = strdup(status.ToString().c_str());
|
||||
}
|
||||
return status.raw_code();
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
TextEmbedder* CppTextEmbedderCreate(const TextEmbedderOptions& options,
|
||||
char** error_msg) {
|
||||
auto cpp_options = std::make_unique<
|
||||
::mediapipe::tasks::text::text_embedder::TextEmbedderOptions>();
|
||||
|
||||
CppConvertToBaseOptions(options.base_options, &cpp_options->base_options);
|
||||
CppConvertToEmbedderOptions(options.embedder_options,
|
||||
&cpp_options->embedder_options);
|
||||
|
||||
auto embedder = TextEmbedder::Create(std::move(cpp_options));
|
||||
if (!embedder.ok()) {
|
||||
ABSL_LOG(ERROR) << "Failed to create TextEmbedder: "
|
||||
<< embedder.status();
|
||||
CppProcessError(embedder.status(), error_msg);
|
||||
return nullptr;
|
||||
}
|
||||
return embedder->release();
|
||||
}
|
||||
|
||||
int CppTextEmbedderEmbed(void* embedder, const char* utf8_str,
|
||||
TextEmbedderResult* result, char** error_msg) {
|
||||
auto cpp_embedder = static_cast<TextEmbedder*>(embedder);
|
||||
auto cpp_result = cpp_embedder->Embed(utf8_str);
|
||||
if (!cpp_result.ok()) {
|
||||
ABSL_LOG(ERROR) << "Embedding extraction failed: " << cpp_result.status();
|
||||
return CppProcessError(cpp_result.status(), error_msg);
|
||||
}
|
||||
CppConvertToEmbeddingResult(*cpp_result, result);
|
||||
return 0;
|
||||
}
|
||||
|
||||
void CppTextEmbedderCloseResult(TextEmbedderResult* result) {
|
||||
CppCloseEmbeddingResult(result);
|
||||
}
|
||||
|
||||
int CppTextEmbedderClose(void* embedder, char** error_msg) {
|
||||
auto cpp_embedder = static_cast<TextEmbedder*>(embedder);
|
||||
auto result = cpp_embedder->Close();
|
||||
if (!result.ok()) {
|
||||
ABSL_LOG(ERROR) << "Failed to close TextEmbedder: " << result;
|
||||
return CppProcessError(result, error_msg);
|
||||
}
|
||||
delete cpp_embedder;
|
||||
return 0;
|
||||
}
|
||||
|
||||
} // namespace mediapipe::tasks::c::text::text_embedder
|
||||
|
||||
extern "C" {
|
||||
|
||||
void* text_embedder_create(struct TextEmbedderOptions* options,
|
||||
char** error_msg) {
|
||||
return mediapipe::tasks::c::text::text_embedder::CppTextEmbedderCreate(
|
||||
*options, error_msg);
|
||||
}
|
||||
|
||||
int text_embedder_embed(void* embedder, const char* utf8_str,
|
||||
TextEmbedderResult* result, char** error_msg) {
|
||||
return mediapipe::tasks::c::text::text_embedder::CppTextEmbedderEmbed(
|
||||
embedder, utf8_str, result, error_msg);
|
||||
}
|
||||
|
||||
void text_embedder_close_result(TextEmbedderResult* result) {
|
||||
mediapipe::tasks::c::text::text_embedder::CppTextEmbedderCloseResult(
|
||||
result);
|
||||
}
|
||||
|
||||
int text_embedder_close(void* embedder, char** error_ms) {
|
||||
return mediapipe::tasks::c::text::text_embedder::CppTextEmbedderClose(
|
||||
embedder, error_ms);
|
||||
}
|
||||
|
||||
} // extern "C"
|
||||
@@ -0,0 +1,75 @@
|
||||
/* 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.
|
||||
==============================================================================*/
|
||||
|
||||
#ifndef MEDIAPIPE_TASKS_C_TEXT_TEXT_CLASSIFIER_TEXT_EMBEDDER_H_
|
||||
#define MEDIAPIPE_TASKS_C_TEXT_TEXT_CLASSIFIER_TEXT_EMBEDDER_H_
|
||||
|
||||
#include "mediapipe/tasks/c/components/containers/embedding_result.h"
|
||||
#include "mediapipe/tasks/c/components/processors/embedder_options.h"
|
||||
#include "mediapipe/tasks/c/core/base_options.h"
|
||||
|
||||
#ifndef MP_EXPORT
|
||||
#define MP_EXPORT __attribute__((visibility("default")))
|
||||
#endif // MP_EXPORT
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
typedef struct EmbeddingResult TextEmbedderResult;
|
||||
|
||||
// The options for configuring a MediaPipe text embedder task.
|
||||
struct TextEmbedderOptions {
|
||||
// Base options for configuring MediaPipe Tasks, such as specifying the model
|
||||
// file with metadata, accelerator options, op resolver, etc.
|
||||
struct BaseOptions base_options;
|
||||
|
||||
// Options for configuring the embedder behavior, such as score threshold,
|
||||
// number of results, etc.
|
||||
struct EmbedderOptions embedder_options;
|
||||
};
|
||||
|
||||
// Creates a TextEmbedder from the provided `options`.
|
||||
// Returns a pointer to the text embedder on success.
|
||||
// If an error occurs, returns `nullptr` and sets the error parameter to an
|
||||
// an error message (if `error_msg` is not nullptr). You must free the memory
|
||||
// allocated for the error message.
|
||||
MP_EXPORT void* text_embedder_create(struct TextEmbedderOptions* options,
|
||||
char** error_msg = nullptr);
|
||||
|
||||
// Performs embedding extraction on the input `text`. Returns `0` on success.
|
||||
// If an error occurs, returns an error code and sets the error parameter to an
|
||||
// an error message (if `error_msg` is not nullptr). You must free the memory
|
||||
// allocated for the error message.
|
||||
MP_EXPORT int text_embedder_embed(void* embedder, const char* utf8_str,
|
||||
TextEmbedderResult* result,
|
||||
char** error_msg = nullptr);
|
||||
|
||||
// Frees the memory allocated inside a TextEmbedderResult result. Does not
|
||||
// free the result pointer itself.
|
||||
MP_EXPORT void text_embedder_close_result(TextEmbedderResult* result);
|
||||
|
||||
// Shuts down the TextEmbedder when all the work is done. Frees all memory.
|
||||
// If an error occurs, returns an error code and sets the error parameter to an
|
||||
// an error message (if `error_msg` is not nullptr). You must free the memory
|
||||
// allocated for the error message.
|
||||
MP_EXPORT int text_embedder_close(void* embedder,
|
||||
char** error_msg = nullptr);
|
||||
|
||||
#ifdef __cplusplus
|
||||
} // extern C
|
||||
#endif
|
||||
|
||||
#endif // MEDIAPIPE_TASKS_C_TEXT_TEXT_CLASSIFIER_TEXT_EMBEDDER_H_
|
||||
@@ -0,0 +1,80 @@
|
||||
/* 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.
|
||||
==============================================================================*/
|
||||
|
||||
#include "mediapipe/tasks/c/text/text_embedder/text_embedder.h"
|
||||
|
||||
#include <cstdlib>
|
||||
#include <string>
|
||||
|
||||
#include "absl/flags/flag.h"
|
||||
#include "absl/strings/string_view.h"
|
||||
#include "mediapipe/framework/deps/file_path.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
|
||||
namespace {
|
||||
|
||||
using ::mediapipe::file::JoinPath;
|
||||
using testing::HasSubstr;
|
||||
|
||||
constexpr char kTestDataDirectory[] = "/mediapipe/tasks/testdata/text/";
|
||||
constexpr char kTestBertModelPath[] = "mobilebert_embedding_with_metadata.tflite";
|
||||
constexpr char kTestString[] = "It's beautiful outside.";
|
||||
constexpr float kPrecision = 1e-6;
|
||||
|
||||
std::string GetFullPath(absl::string_view file_name) {
|
||||
return JoinPath("./", kTestDataDirectory, file_name);
|
||||
}
|
||||
|
||||
TEST(TextEmbedderTest, SmokeTest) {
|
||||
std::string model_path = GetFullPath(kTestBertModelPath);
|
||||
TextEmbedderOptions options = {
|
||||
/* base_options= */ {/* model_asset_buffer= */ nullptr,
|
||||
/* model_asset_path= */ model_path.c_str()},
|
||||
/* embedder_options= */
|
||||
{/* l2_normalize= */ false,
|
||||
/* quantize= */ true},
|
||||
};
|
||||
|
||||
void* embedder = text_embedder_create(&options);
|
||||
EXPECT_NE(embedder, nullptr);
|
||||
|
||||
TextEmbedderResult result;
|
||||
text_embedder_embed(embedder, kTestString, &result);
|
||||
EXPECT_EQ(result.embeddings_count, 1);
|
||||
EXPECT_EQ(result.embeddings[0].values_count, 512);
|
||||
|
||||
text_embedder_close_result(&result);
|
||||
text_embedder_close(embedder);
|
||||
}
|
||||
|
||||
TEST(TextEmbedderTest, ErrorHandling) {
|
||||
// It is an error to set neither the asset buffer nor the path.
|
||||
TextEmbedderOptions options = {
|
||||
/* base_options= */ {/* model_asset_buffer= */ nullptr,
|
||||
/* model_asset_path= */ nullptr},
|
||||
/* embedder_options= */ {},
|
||||
};
|
||||
|
||||
char* error_msg;
|
||||
void* embedder = text_embedder_create(&options, &error_msg);
|
||||
EXPECT_EQ(embedder, nullptr);
|
||||
|
||||
EXPECT_THAT(error_msg, HasSubstr("INVALID_ARGUMENT"));
|
||||
|
||||
free(error_msg);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
Reference in New Issue
Block a user