Add Java TextEmbedder API.
PiperOrigin-RevId: 488427327
This commit is contained in:
committed by
Copybara-Service
parent
b40b2ade14
commit
34daba4747
+1
-3
@@ -67,9 +67,7 @@ public class TextClassifierTest {
|
||||
ApplicationProvider.getApplicationContext(), options));
|
||||
// TODO: Make MediaPipe InferenceCalculator report the detailed.
|
||||
// interpreter errors (e.g., "Encountered unresolved custom op").
|
||||
assertThat(exception)
|
||||
.hasMessageThat()
|
||||
.contains("interpreter_builder(&interpreter) == kTfLiteOk");
|
||||
assertThat(exception).hasMessageThat().contains("== kTfLiteOk");
|
||||
}
|
||||
|
||||
@Test
|
||||
|
||||
+24
@@ -0,0 +1,24 @@
|
||||
<?xml version="1.0" encoding="utf-8"?>
|
||||
<manifest xmlns:android="http://schemas.android.com/apk/res/android"
|
||||
package="com.google.mediapipe.tasks.text.textembeddertest"
|
||||
android:versionCode="1"
|
||||
android:versionName="1.0" >
|
||||
|
||||
<uses-permission android:name="android.permission.READ_EXTERNAL_STORAGE"/>
|
||||
<uses-permission android:name="android.permission.WRITE_EXTERNAL_STORAGE"/>
|
||||
|
||||
<uses-sdk android:minSdkVersion="24"
|
||||
android:targetSdkVersion="30" />
|
||||
|
||||
<application
|
||||
android:label="textembeddertest"
|
||||
android:name="android.support.multidex.MultiDexApplication"
|
||||
android:taskAffinity="">
|
||||
<uses-library android:name="android.test.runner" />
|
||||
</application>
|
||||
|
||||
<instrumentation
|
||||
android:name="com.google.android.apps.common.testing.testrunner.GoogleInstrumentationTestRunner"
|
||||
android:targetPackage="com.google.mediapipe.tasks.text.textembeddertest" />
|
||||
|
||||
</manifest>
|
||||
@@ -0,0 +1,19 @@
|
||||
# Copyright 2022 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"])
|
||||
|
||||
# TODO: Enable this in OSS
|
||||
+98
@@ -0,0 +1,98 @@
|
||||
// Copyright 2022 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 com.google.mediapipe.tasks.text.textembedder;
|
||||
|
||||
import static com.google.common.truth.Truth.assertThat;
|
||||
import static org.junit.Assert.assertThrows;
|
||||
|
||||
import androidx.test.core.app.ApplicationProvider;
|
||||
import androidx.test.ext.junit.runners.AndroidJUnit4;
|
||||
import com.google.mediapipe.framework.MediaPipeException;
|
||||
import org.junit.Test;
|
||||
import org.junit.runner.RunWith;
|
||||
|
||||
/** Test for {@link TextEmbedder}/ */
|
||||
@RunWith(AndroidJUnit4.class)
|
||||
public class TextEmbedderTest {
|
||||
private static final String BERT_MODEL_FILE = "mobilebert_embedding_with_metadata.tflite";
|
||||
private static final String REGEX_MODEL_FILE = "regex_one_embedding_with_metadata.tflite";
|
||||
|
||||
private static final double DOUBLE_DIFF_TOLERANCE = 1e-4;
|
||||
private static final float FLOAT_DIFF_TOLERANCE = 1e-4f;
|
||||
|
||||
@Test
|
||||
public void create_failsWithMissingModel() throws Exception {
|
||||
String nonExistentFile = "/path/to/non/existent/file";
|
||||
MediaPipeException exception =
|
||||
assertThrows(
|
||||
MediaPipeException.class,
|
||||
() ->
|
||||
TextEmbedder.createFromFile(
|
||||
ApplicationProvider.getApplicationContext(), nonExistentFile));
|
||||
assertThat(exception).hasMessageThat().contains(nonExistentFile);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void embed_succeedsWithBert() throws Exception {
|
||||
TextEmbedder textEmbedder =
|
||||
TextEmbedder.createFromFile(ApplicationProvider.getApplicationContext(), BERT_MODEL_FILE);
|
||||
|
||||
TextEmbedderResult result0 = textEmbedder.embed("it's a charming and often affecting journey");
|
||||
assertThat(result0.embeddingResult().embeddings().size()).isEqualTo(1);
|
||||
assertThat(result0.embeddingResult().embeddings().get(0).floatEmbedding()).hasLength(512);
|
||||
assertThat(result0.embeddingResult().embeddings().get(0).floatEmbedding()[0])
|
||||
.isWithin(FLOAT_DIFF_TOLERANCE)
|
||||
.of(20.59746f);
|
||||
TextEmbedderResult result1 = textEmbedder.embed("what a great and fantastic trip");
|
||||
assertThat(result1.embeddingResult().embeddings().size()).isEqualTo(1);
|
||||
assertThat(result1.embeddingResult().embeddings().get(0).floatEmbedding()).hasLength(512);
|
||||
assertThat(result1.embeddingResult().embeddings().get(0).floatEmbedding()[0])
|
||||
.isWithin(FLOAT_DIFF_TOLERANCE)
|
||||
.of(21.774776f);
|
||||
|
||||
// Check cosine similarity.
|
||||
double similarity =
|
||||
TextEmbedder.cosineSimilarity(
|
||||
result0.embeddingResult().embeddings().get(0),
|
||||
result1.embeddingResult().embeddings().get(0));
|
||||
assertThat(similarity).isWithin(DOUBLE_DIFF_TOLERANCE).of(0.968879);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void embed_succeedsWithRegex() throws Exception {
|
||||
TextEmbedder textEmbedder =
|
||||
TextEmbedder.createFromFile(ApplicationProvider.getApplicationContext(), REGEX_MODEL_FILE);
|
||||
|
||||
TextEmbedderResult result0 = textEmbedder.embed("it's a charming and often affecting journey");
|
||||
assertThat(result0.embeddingResult().embeddings().size()).isEqualTo(1);
|
||||
assertThat(result0.embeddingResult().embeddings().get(0).floatEmbedding()).hasLength(16);
|
||||
assertThat(result0.embeddingResult().embeddings().get(0).floatEmbedding()[0])
|
||||
.isWithin(FLOAT_DIFF_TOLERANCE)
|
||||
.of(0.030935612f);
|
||||
TextEmbedderResult result1 = textEmbedder.embed("what a great and fantastic trip");
|
||||
assertThat(result1.embeddingResult().embeddings().size()).isEqualTo(1);
|
||||
assertThat(result1.embeddingResult().embeddings().get(0).floatEmbedding()).hasLength(16);
|
||||
assertThat(result1.embeddingResult().embeddings().get(0).floatEmbedding()[0])
|
||||
.isWithin(FLOAT_DIFF_TOLERANCE)
|
||||
.of(0.0312863f);
|
||||
|
||||
// Check cosine similarity.
|
||||
double similarity =
|
||||
TextEmbedder.cosineSimilarity(
|
||||
result0.embeddingResult().embeddings().get(0),
|
||||
result1.embeddingResult().embeddings().get(0));
|
||||
assertThat(similarity).isWithin(DOUBLE_DIFF_TOLERANCE).of(0.999937);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user