forked from opensearch-project/k-NN
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Integration With Qunatization Config
Signed-off-by: VIKASH TIWARI <viktari@amazon.com>
- Loading branch information
Showing
20 changed files
with
326 additions
and
135 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
40 changes: 40 additions & 0 deletions
40
src/main/java/org/opensearch/knn/index/codec/nativeindex/IndexBuildSetup.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,40 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.knn.index.codec.nativeindex; | ||
|
|
||
| import lombok.AllArgsConstructor; | ||
| import lombok.Getter; | ||
| import org.opensearch.knn.quantization.models.quantizationOutput.QuantizationOutput; | ||
| import org.opensearch.knn.quantization.models.quantizationState.QuantizationState; | ||
|
|
||
| /** | ||
| * IndexBuildSetup encapsulates the configuration and parameters required for building an index. | ||
| * This includes the size of each vector, the dimensions of the vectors, and any quantization-related | ||
| * settings such as the output and state of quantization. | ||
| */ | ||
| @Getter | ||
| @AllArgsConstructor | ||
| final class IndexBuildSetup { | ||
| /** | ||
| * The number of bytes per vector. | ||
| */ | ||
| private final int bytesPerVector; | ||
|
|
||
| /** | ||
| * Dimension of Vector for Indexing | ||
| */ | ||
| private final int dimensions; | ||
|
|
||
| /** | ||
| * The quantization output that will hold the quantized vector. | ||
| */ | ||
| private final QuantizationOutput quantizationOutput; | ||
|
|
||
| /** | ||
| * The state of quantization, which may include parameters and trained models. | ||
| */ | ||
| private final QuantizationState quantizationState; | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
68 changes: 68 additions & 0 deletions
68
src/main/java/org/opensearch/knn/index/codec/nativeindex/QuantizationIndexUtils.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,68 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.knn.index.codec.nativeindex; | ||
|
|
||
| import lombok.experimental.UtilityClass; | ||
| import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams; | ||
| import org.opensearch.knn.index.quantizationService.QuantizationService; | ||
| import org.opensearch.knn.index.vectorvalues.KNNVectorValues; | ||
| import org.opensearch.knn.quantization.models.quantizationOutput.QuantizationOutput; | ||
| import org.opensearch.knn.quantization.models.quantizationState.QuantizationState; | ||
|
|
||
| import java.io.IOException; | ||
|
|
||
| @UtilityClass | ||
| class QuantizationIndexUtils { | ||
|
|
||
| /** | ||
| * Processes and returns the vector based on whether quantization is applied or not. | ||
| * | ||
| * @param knnVectorValues the KNN vector values to be processed. | ||
| * @param indexBuildSetup the setup containing quantization state and output, along with other parameters. | ||
| * @return the processed vector, either quantized or original. | ||
| * @throws IOException if an I/O error occurs during processing. | ||
| */ | ||
| static Object processAndReturnVector(KNNVectorValues<?> knnVectorValues, IndexBuildSetup indexBuildSetup) throws IOException { | ||
| QuantizationService quantizationService = QuantizationService.getInstance(); | ||
| if (indexBuildSetup.getQuantizationState() != null && indexBuildSetup.getQuantizationOutput() != null) { | ||
| quantizationService.quantize( | ||
| indexBuildSetup.getQuantizationState(), | ||
| knnVectorValues.getVector(), | ||
| indexBuildSetup.getQuantizationOutput() | ||
| ); | ||
| return indexBuildSetup.getQuantizationOutput().getQuantizedVector(); | ||
| } else { | ||
| return knnVectorValues.conditionalCloneVector(); | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Prepares the quantization setup including bytes per vector and dimensions. | ||
| * | ||
| * @param knnVectorValues the KNN vector values. | ||
| * @param indexInfo the index build parameters. | ||
| * @return an instance of QuantizationSetup containing relevant information. | ||
| */ | ||
| static IndexBuildSetup prepareIndexBuild(KNNVectorValues<?> knnVectorValues, BuildIndexParams indexInfo) { | ||
| QuantizationState quantizationState = indexInfo.getQuantizationState(); | ||
| QuantizationOutput quantizationOutput = null; | ||
| QuantizationService quantizationService = QuantizationService.getInstance(); | ||
|
|
||
| int bytesPerVector; | ||
| int dimensions; | ||
|
|
||
| if (quantizationState != null) { | ||
| bytesPerVector = quantizationState.getBytesPerVector(); | ||
| dimensions = quantizationState.getDimensions(); | ||
| quantizationOutput = quantizationService.createQuantizationOutput(quantizationState.getQuantizationParams()); | ||
| } else { | ||
| bytesPerVector = knnVectorValues.bytesPerVector(); | ||
| dimensions = knnVectorValues.dimension(); | ||
| } | ||
|
|
||
| return new IndexBuildSetup(bytesPerVector, dimensions, quantizationOutput, quantizationState); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.