Add Text Embedder tests for text with different themes

PiperOrigin-RevId: 506023265
2023-01-31 09:23:30 -08:00
parent 2c4dece023
commit be3bddc620
3 changed files with 77 additions and 0 deletions
@@ -95,4 +95,24 @@ public class TextEmbedderTest {
            result1.embeddingResult().embeddings().get(0));
    assertThat(similarity).isWithin(DOUBLE_DIFF_TOLERANCE).of(0.999937);
  }
+
+  @Test
+  public void classify_succeedsWithBertAndDifferentThemes() throws Exception {
+    TextEmbedder textEmbedder =
+        TextEmbedder.createFromFile(ApplicationProvider.getApplicationContext(), BERT_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.3477488707202946);
+  }
 }