diff --git a/mediapipe/tasks/javatests/com/google/mediapipe/tasks/text/textembedder/TextEmbedderTest.java b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/text/textembedder/TextEmbedderTest.java index 48f21477..79628826 100644 --- a/mediapipe/tasks/javatests/com/google/mediapipe/tasks/text/textembedder/TextEmbedderTest.java +++ b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/text/textembedder/TextEmbedderTest.java @@ -28,6 +28,7 @@ import org.junit.runner.RunWith; 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 String USE_MODEL_FILE = "universal_sentence_encoder_qa_with_metadata.tflite"; private static final double DOUBLE_DIFF_TOLERANCE = 1e-4; private static final float FLOAT_DIFF_TOLERANCE = 1e-4f; @@ -70,6 +71,32 @@ public class TextEmbedderTest { assertThat(similarity).isWithin(DOUBLE_DIFF_TOLERANCE).of(0.968879); } + @Test + public void embed_succeedsWithUSE() throws Exception { + TextEmbedder textEmbedder = + TextEmbedder.createFromFile(ApplicationProvider.getApplicationContext(), USE_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(100); + assertThat(result0.embeddingResult().embeddings().get(0).floatEmbedding()[0]) + .isWithin(FLOAT_DIFF_TOLERANCE) + .of(1.422951f); + 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(100); + assertThat(result1.embeddingResult().embeddings().get(0).floatEmbedding()[0]) + .isWithin(FLOAT_DIFF_TOLERANCE) + .of(1.404664f); + + // 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.851961); + } + @Test public void embed_succeedsWithRegex() throws Exception { TextEmbedder textEmbedder = @@ -115,4 +142,24 @@ public class TextEmbedderTest { result1.embeddingResult().embeddings().get(0)); assertThat(similarity).isWithin(DOUBLE_DIFF_TOLERANCE).of(0.3477488707202946); } + + @Test + public void classify_succeedsWithUSEAndDifferentThemes() throws Exception { + TextEmbedder textEmbedder = + TextEmbedder.createFromFile(ApplicationProvider.getApplicationContext(), USE_MODEL_FILE); + + TextEmbedderResult result0 = + textEmbedder.embed( + "When you go to this restaurant, they hold the pancake upside-down before they hand " + + "it to you. It's a great gimmick."); + TextEmbedderResult result1 = + textEmbedder.embed("Let\'s make a plan to steal the declaration of independence.'"); + + // 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.7835510599396296); + } } diff --git a/mediapipe/tasks/testdata/text/BUILD b/mediapipe/tasks/testdata/text/BUILD index a0131c05..9813b654 100644 --- a/mediapipe/tasks/testdata/text/BUILD +++ b/mediapipe/tasks/testdata/text/BUILD @@ -103,5 +103,5 @@ filegroup( filegroup( name = "universal_sentence_encoder_qa", - data = ["universal_sentence_encoder_qa_with_metadata.tflite"], + srcs = ["universal_sentence_encoder_qa_with_metadata.tflite"], )