Merge pull request #375 from marklogic/feature/allow-zero-vectors

Allowing for zero vectors

Author: rjrudin

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

@@ -97,34 +97,11 @@ private void addEmbeddingsToChunks(List<Chunk> chunks) {
         } else {
             List<Embedding> embeddings = response.content();
             for (int i = 0; i < embeddings.size(); i++) {
-                addEmbeddingToChunk(chunks.get(i), embeddings.get(i));
+                chunks.get(i).addEmbedding(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,42 +103,6 @@ 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

@@ -10,7 +10,6 @@
 import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
 import dev.langchain4j.model.output.Response;
 
-import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Function;
@@ -30,11 +29,8 @@ public static void reset() {
 
     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;
     }
 
@@ -47,10 +43,6 @@ public int dimension() {
     public Response<List<Embedding>> embedAll(List<TextSegment> textSegments) {
         batchCounter++;
         chunkCounter += textSegments.size();
-        if (returnZeroesOnFirstCall) {
-            returnZeroesOnFirstCall = false;
-            return Response.from(Arrays.asList(new Embedding(new float[384])));
-        }
         return realEmbeddingModel.embedAll(textSegments);
     }
 

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": "ignore",
-            "nullable": false
+            "invalidValues": "reject",
+            "nullable": true
           }
         ]
       }