Merge pull request #6 from alwayslove2013/hnswlib_parser

Hnswlib parser

esm/Feder.js

@@ -21,7 +21,7 @@ export default class Feder {
     }
     this.core = core;
     this.dom = dom;
-    this.setViewParams(viewParams);
+    // this.setViewParams(viewParams);
   }
   update() {
     // update core

esm/FederCore/FaissFileReader.js

@@ -1,4 +1,4 @@
-import FileReader from './FileReader.js'
+import FileReader from './FileReader.js';
 import { uint8toChars, generateArray } from '../Utils/index.js';
 
 export default class FaissFileReader extends FileReader {
@@ -14,4 +14,4 @@ export default class FaissFileReader extends FileReader {
     const dummy = this.readUint64();
     return dummy;
   }
-}
+}

esm/FederCore/FileReader.js

@@ -24,6 +24,11 @@ export default class FileReader {
     const int8 = this.readInt8();
     return Boolean(int8);
   }
+  readUint16() {
+    const uint16 = this.dataview.getUint16(this.p, true);
+    this.p += 2;
+    return uint16;
+  }
   readInt32() {
     const int32 = this.dataview.getInt32(this.p, true);
     this.p += 4;
@@ -39,10 +44,22 @@ export default class FileReader {
     const right = this.readUint32();
     const int64 = left + Math.pow(2, 32) * right;
     if (!Number.isSafeInteger(int64))
-      console.warn(int64, 'exceeds MAX_SAFE_INTEGER. Precision may be lost');
+      console.warn(int64, 'Exceeds MAX_SAFE_INTEGER. Precision may be lost');
     return int64;
   }
+  readFloat64() {
+    const float64 = this.dataview.getFloat64(this.p, true);
+    this.p += 8;
+    return float64;
+  }
+  readDouble() {
+    return this.readFloat64();
+  }
 
+  readUint32Array(n) {
+    const res = generateArray(n).map((_) => this.readUint32());
+    return res;
+  }
   readFloat32Array(n) {
     const res = new Float32Array(this.data.slice(this.p, this.p + n * 4));
     this.p += n * 4;
@@ -52,4 +69,4 @@ export default class FileReader {
     const res = generateArray(n).map((_) => this.readUint64());
     return res;
   }
-}
+}

esm/FederCore/HNSWlibFileReader.js

@@ -0,0 +1,18 @@
+import FileReader from './FileReader.js';
+
+export default class HNSWlibFileReader extends FileReader {
+  constructor(arrayBuffer) {
+    super(arrayBuffer);
+  }
+  readIsDeleted() {
+    return this.readUint8()
+  }
+  readIsReused() {
+    return this.readUint8()
+  }
+  readLevelOCount() {
+    // const res = [this.readUint8(), this.readUint8()];
+    // return res;
+    return this.readUint16();
+  }
+}

esm/FederCore/hnswlibIndexParser.js

@@ -0,0 +1,92 @@
+import HNSWlibFileReader from './HNSWlibFileReader.js';
+import { INDEX_TYPE } from '../Utils/config.js';
+import { generateArray } from '../Utils/index.js';
+
+const hnswlibIndexParser = (arrayBuffer) => {
+  const reader = new HNSWlibFileReader(arrayBuffer);
+  const index = {};
+  index.offsetLevel0_ = reader.readUint64();
+  index.max_elements_ = reader.readUint64();
+  index.cur_element_count = reader.readUint64();
+  // index.ntotal = index.cur_element_count; // consistent with Faiss
+  index.size_data_per_element_ = reader.readUint64();
+  index.label_offset_ = reader.readUint64();
+  index.offsetData_ = reader.readUint64();
+  index.dim = (index.size_data_per_element_ - index.offsetData_ - 8) / 4;
+
+  index.maxlevel_ = reader.readUint32();
+  index.enterpoint_node_ = reader.readUint32();
+  index.maxM_ = reader.readUint64();
+  index.maxM0_ = reader.readUint64();
+  index.M = reader.readUint64();
+
+  index.mult_ = reader.readFloat64();
+  index.ef_construction_ = reader.readUint64();
+
+  index.size_links_per_element_ = index.maxM_ * 4 + 4;
+  index.size_links_level0_ = index.maxM0_ * 4 + 4;
+  index.revSize_ = 1.0 / index.mult_;
+  index.ef_ = 10;
+
+  read_data_level0_memory_(reader, index);
+
+  const linkListSizes = [];
+  const linkLists_ = [];
+  for (let i = 0; i < index.cur_element_count; i++) {
+    const linkListSize = reader.readUint32();
+    linkListSizes.push(linkListSize);
+    if (linkListSize === 0) {
+      linkLists_[i] = [];
+    } else {
+      const levelCount = linkListSize / 4 / (index.maxM_ + 1);
+      linkLists_[i] = generateArray(levelCount).map((_) =>
+        reader.readUint32Array(index.maxM_ + 1)
+      );
+    }
+  }
+  index.linkListSizes = linkListSizes;
+  index.linkLists_ = linkLists_;
+
+  console.assert(
+    reader.isEmpty,
+    'HNSWlib Parser Failed. Not empty when the parser completes.'
+  );
+
+  return {
+    indexType: INDEX_TYPE.HNSW,
+    ntotal: index.cur_element_count,
+    vectors: index.vectors,
+    maxLevel: index.maxlevel_,
+    linkLists_level0_count: index.linkLists_level0_count,
+    linkLists_level_0: index.linkLists_level0,
+    linkLists_levels: index.linkLists_,
+    enterPoint: index.enterpoint_node_,
+    labels: index.externalLabel,
+    isDeleted: index.isDeleted,
+  };
+};
+
+const read_data_level0_memory_ = (reader, index) => {
+  // size_data = links_level0[M0 + 1] * 4 + vector[dim * 4] * 4 + label[1] * 8;
+  const isDeleted = [];
+  const linkLists_level0_count = [];
+  const linkLists_level0 = [];
+  const vectors = [];
+  const externalLabel = [];
+  for (let i = 0; i < index.cur_element_count; i++) {
+    linkLists_level0_count.push(reader.readLevelOCount());
+    isDeleted.push(reader.readIsDeleted());
+    reader.readIsReused(); // Unknown use.
+    linkLists_level0.push(reader.readUint32Array(index.maxM0_));
+    vectors.push(reader.readFloat32Array(index.dim));
+    externalLabel.push(reader.readUint64());
+    // console.log(isDeleted, linkLists_level0_count);
+  }
+  index.isDeleted = isDeleted;
+  index.linkLists_level0_count = linkLists_level0_count;
+  index.linkLists_level0 = linkLists_level0;
+  index.vectors = vectors;
+  index.externalLabel = externalLabel;
+};
+
+export default hnswlibIndexParser;

esm/FederCore/index.js

@@ -1,21 +1,29 @@
 import faissIndexParser from './faissIndexParser.js';
+import hnswlibIndexParser from './hnswlibIndexParser.js';
 import faissIVFFlatSearch from './faissIVFFlatSearch.js';
 import {
   ProjectMethod,
   getProjectFunc,
   getProjectParamsGuide,
 } from '../Utils/projector/index.js';
+import { SOURCE_TYPE } from '../Utils/config.js';
 
 const indexSearchHandlerMap = {
   faissIVFFlat: faissIVFFlatSearch,
   faissHNSW: null, // todo,
   hnswlibHNSW: null, // todo,
 };
+
+const indexParserMap = {
+  [SOURCE_TYPE.Faiss]: faissIndexParser,
+  [SOURCE_TYPE.HNSWlib]: hnswlibIndexParser,
+};
+
 /* Feder core class */
 export default class FederCore {
   constructor({
     data,
-    source = 'faiss',
+    source = SOURCE_TYPE.Faiss,
     projectMethod = ProjectMethod.UMAP,
     projectParams = {},
   }) {
@@ -30,25 +38,22 @@ export default class FederCore {
     this.setIndexSource(source);
     this.parseIndex();
 
+    console.log(this.index);
+
     if (this.index) {
       this.setIndexSearchHandler();
       this[`_updateId2Vec_${this.index.indexType}`]();
     }
 
     this.setProjectParams(projectMethod, projectParams);
-
-    // need parsed_index and projector.
-    // this.index && this[`_updateIndexMeta_${this.index.indexType}`]();
   }
   get indexType() {
     return this.index.indexType || '';
   }
   setIndexSource(source) {
     this.indexParser = null;
-    this.indexSource = source;
-    if (source === 'faiss') {
-      this.indexParser = faissIndexParser;
-    }
+    this.indexSource = source.toLowerCase();
+    this.indexParser = indexParserMap[source];
   }
   parseIndex() {
     if (this.indexParser) {

esm/Utils/config.js

@@ -26,3 +26,8 @@ export const ANiMATION_TYPE = {
   Enter: 'Enter',
   Exit: 'Exit',
 }
+
+export const SOURCE_TYPE = {
+  Faiss: 'faiss',
+  HNSWlib: 'hnswlib',
+}

esm/Utils/index.js

@@ -3,7 +3,7 @@ export const uint8toChars = (data) => {
 };
 
 export const generateArray = (num) => {
-  return Array.from(new Array(num).keys());
+  return Array.from(new Array(Math.floor(num)).keys());
 };
 
 export const polyPoints2path = (points) => {