Merge pull request #373 from marklogic/feature/zero-vector-fix

Not adding invalid embeddings

Author: rjrudin

marklogic-langchain4j/src/main/java/com/marklogic/langchain4j/embedding/EmbeddingGenerator.java

@@ -97,11 +97,34 @@ private void addEmbeddingsToChunks(List<Chunk> chunks) {
         } else {
             List<Embedding> embeddings = response.content();
             for (int i = 0; i < embeddings.size(); i++) {
-                chunks.get(i).addEmbedding(embeddings.get(i));
+                addEmbeddingToChunk(chunks.get(i), embeddings.get(i));
             }
         }
     }
 
+    private void addEmbeddingToChunk(Chunk chunk, Embedding embedding) {
+        if (vectorIsAllZeroes(embedding)) {
+            if (Util.LANGCHAIN4J_LOGGER.isDebugEnabled()) {
+                Util.LANGCHAIN4J_LOGGER.debug("Not adding embedding to chunk as it only contains zeroes; source document URI: {}; text: {}",
+                    chunk.getDocumentUri(), chunk.getEmbeddingText());
+            } else {
+                Util.LANGCHAIN4J_LOGGER.warn("Not adding embedding to chunk as it only contains zeroes; source document URI: {}",
+                    chunk.getDocumentUri());
+            }
+        } else {
+            chunk.addEmbedding(embedding);
+        }
+    }
+
+    private boolean vectorIsAllZeroes(Embedding embedding) {
+        for (float f : embedding.vector()) {
+            if (f != 0.0f) {
+                return false;
+            }
+        }
+        return true;
+    }
+
     private List<TextSegment> makeTextSegments(List<Chunk> chunks) {
         return chunks.stream()
             .map(chunk -> new TextSegment(chunk.getEmbeddingText(), TEXT_SEGMENT_METADATA))

marklogic-spark-connector/src/test/java/com/marklogic/spark/writer/embedding/AddEmbeddingsToJsonTest.java

@@ -103,6 +103,42 @@ void addEmbeddingsToExistingSplits() {
         verifyEachChunkIsReturnedByAVectorQuery();
     }
 
+    @ExtendWith(RequiresMarkLogic12.class)
+    @Test
+    void vectorHasAllZeroes() {
+        readDocument("/marklogic-docs/java-client-intro.json")
+            .repartition(1)
+            .write().format(CONNECTOR_IDENTIFIER)
+            .option(Options.CLIENT_URI, makeClientUri())
+            .option(Options.WRITE_SPLITTER_JSON_POINTERS, "/text")
+            .option(Options.WRITE_SPLITTER_SIDECAR_MAX_CHUNKS, 10)
+            .option(Options.WRITE_SPLITTER_SIDECAR_COLLECTIONS, "json-vector-chunks")
+            .option(Options.WRITE_PERMISSIONS, DEFAULT_PERMISSIONS)
+            .option(Options.WRITE_URI_TEMPLATE, "/split-test.json")
+            .option(Options.WRITE_EMBEDDER_MODEL_FUNCTION_CLASS_NAME, "com.marklogic.spark.writer.embedding.TestEmbeddingModel")
+            .option(Options.WRITE_EMBEDDER_MODEL_FUNCTION_OPTION_PREFIX + "returnZeroesOnFirstCall", "true")
+            .mode(SaveMode.Append)
+            .save();
+
+        JsonNode doc = readJsonDocument("/split-test.json-chunks-1.json");
+        JsonNode firstChunk = doc.get("chunks").get(0);
+        assertFalse(firstChunk.has("embedding"), "The first chunk is given an array of all zeroes by the test " +
+            "embedding model. Flux should recognize this and not add the `embedding` field, as doing so will cause " +
+            "issues with the Optic vector library - specifically, a VEC-MAGNITUDEZERO error at least when using " +
+            "vec.cosineSimilarity and then sorting on the values. A future version of MarkLogic 12 may improve this " +
+            "by allowing for an array of zeroes to be rejected.");
+
+        JsonNode secondChunk = doc.get("chunks").get(1);
+        assertTrue(secondChunk.has("embedding"), "The test embedding model should generate a valid embedding for " +
+            "the second chunk, which means it can be queried next using Optic.");
+
+        RowManager rowManager = getDatabaseClient().newRowManager();
+        RowSet<RowRecord> rows = rowManager.resultRows(rowManager.newPlanBuilder().fromView("example", "json_chunks"));
+        assertEquals(1, rows.stream().count(), "The TDE has nullable=false for the embedding column, as a null " +
+            "vector will cause issues when querying on vectors. And since invalidValues=ignore, the first chunk " +
+            "won't be returned; only the second chunk will be.");
+    }
+
     @Test
     void passOptionsToEmbeddingModelFunction() {
         DataFrameWriter writer = readDocument("/marklogic-docs/java-client-intro.json")

marklogic-spark-connector/src/test/java/com/marklogic/spark/writer/embedding/TestEmbeddingModel.java

@@ -7,6 +7,7 @@
 import dev.langchain4j.data.embedding.Embedding;
 import dev.langchain4j.data.segment.TextSegment;
 import dev.langchain4j.model.embedding.EmbeddingModel;
+import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
 import dev.langchain4j.model.output.Response;
 
 import java.util.Arrays;
@@ -27,8 +28,13 @@ public static void reset() {
         chunkCounter = 0;
     }
 
+    private static AllMiniLmL6V2EmbeddingModel realEmbeddingModel = new AllMiniLmL6V2EmbeddingModel();
+
+    private boolean returnZeroesOnFirstCall;
+
     @Override
     public EmbeddingModel apply(Map<String, String> options) {
+        returnZeroesOnFirstCall = "true".equals(options.get("returnZeroesOnFirstCall"));
         return this;
     }
 
@@ -41,7 +47,11 @@ public int dimension() {
     public Response<List<Embedding>> embedAll(List<TextSegment> textSegments) {
         batchCounter++;
         chunkCounter += textSegments.size();
-        return Response.from(Arrays.asList(new Embedding(new float[]{1})));
+        if (returnZeroesOnFirstCall) {
+            returnZeroesOnFirstCall = false;
+            return Response.from(Arrays.asList(new Embedding(new float[384])));
+        }
+        return realEmbeddingModel.embedAll(textSegments);
     }
 
     public static class TestChunk implements Chunk {

marklogic-spark-langchain4j/src/main/java/com/marklogic/spark/langchain4j/EmbeddingAdderFactory.java

@@ -21,7 +21,7 @@ public abstract class EmbeddingAdderFactory {
     public static Optional<EmbeddingAdder> makeEmbedder(Context context, DocumentTextSplitter splitter) {
         Optional<EmbeddingModel> embeddingModel = makeEmbeddingModel(context);
         if (embeddingModel.isPresent()) {
-            EmbeddingGenerator embeddingGenerator = makeEmbeddingGenerator(context);
+            EmbeddingGenerator embeddingGenerator = makeEmbeddingGenerator(context, embeddingModel.get());
             if (splitter != null) {
                 return Optional.of(new EmbeddingAdder(splitter, embeddingGenerator));
             }
@@ -31,17 +31,12 @@ public static Optional<EmbeddingAdder> makeEmbedder(Context context, DocumentTex
         return Optional.empty();
     }
 
-    public static EmbeddingGenerator makeEmbeddingGenerator(Context context) {
-        Optional<EmbeddingModel> embeddingModel = makeEmbeddingModel(context);
-        if (embeddingModel.isPresent()) {
-            int batchSize = context.getIntOption(Options.WRITE_EMBEDDER_BATCH_SIZE, 1, 1);
-            EmbeddingModel model = embeddingModel.get();
-            if (Util.MAIN_LOGGER.isInfoEnabled()) {
-                Util.MAIN_LOGGER.info("Using embedding model with dimension: {}", model.dimension());
-            }
-            return new EmbeddingGenerator(model, batchSize);
+    private static EmbeddingGenerator makeEmbeddingGenerator(Context context, EmbeddingModel model) {
+        int batchSize = context.getIntOption(Options.WRITE_EMBEDDER_BATCH_SIZE, 1, 1);
+        if (Util.MAIN_LOGGER.isInfoEnabled()) {
+            Util.MAIN_LOGGER.info("Using embedding model with dimension: {}", model.dimension());
         }
-        return null;
+        return new EmbeddingGenerator(model, batchSize);
     }
 
     /**

test-app/src/main/ml-schemas-12/tde/json-vector-chunks.json

@@ -19,8 +19,8 @@
             "scalarType": "vector",
             "val": "vec:vector(embedding)",
             "dimension": "384",
-            "invalidValues": "reject",
-            "nullable": true
+            "invalidValues": "ignore",
+            "nullable": false
           }
         ]
       }