From acd8fe99c214433a9151ca07ac4533c80f7eda4a Mon Sep 17 00:00:00 2001 From: Jay Wang Date: Tue, 30 Jan 2024 16:56:35 -0500 Subject: [PATCH] Test inserting 100 items 3 layer with m=3 Signed-off-by: Jay Wang --- src/mememo.ts | 8 +- test/data/insert-100-3-layer-m=3.json | 461 ++++++++++++++++++++++++++ test/mememo.test.ts | 62 +++- 3 files changed, 523 insertions(+), 8 deletions(-) create mode 100644 test/data/insert-100-3-layer-m=3.json diff --git a/src/mememo.ts b/src/mememo.ts index 2674a82..80d436c 100644 --- a/src/mememo.ts +++ b/src/mememo.ts @@ -186,8 +186,10 @@ export class HNSW { * Insert a new element to the index. * @param key Key of the new element. * @param value The embedding of the new element to insert. + * @param maxLevel The max layer to insert this element. You don't need to set + * this value in most cases. We add this parameter for testing purpose. */ - insert(key: T, value: number[]) { + insert(key: T, value: number[], maxLevel?: number | undefined) { // If the key already exists, update the node if (this.nodes.has(key)) { // TODO: Update the node @@ -195,8 +197,8 @@ export class HNSW { } // Randomly determine the max level of this node - const level = this._getRandomLevel(); - // console.log('random level:', level); + const level = maxLevel === undefined ? this._getRandomLevel() : maxLevel; + console.log('random level:', level); // Add this node to the node index first this.nodes.set(key, new Node(key, value)); diff --git a/test/data/insert-100-3-layer-m=3.json b/test/data/insert-100-3-layer-m=3.json new file mode 100644 index 0000000..fe9de9b --- /dev/null +++ b/test/data/insert-100-3-layer-m=3.json @@ -0,0 +1,461 @@ +[ + { + "1": { "10": 0.0782, "911": 0.3178 }, + "2": { + "4": 0.094, + "8": 0.1102, + "10": 0.1143, + "5": 0.126, + "9": 0.1555, + "644": 0.2674 + }, + "3": { "4": 0.1278, "11": 0.1406 }, + "4": { + "2": 0.094, + "14": 0.1119, + "15": 0.1135, + "12": 0.127, + "3": 0.1278, + "770": 0.3129 + }, + "5": { "2": 0.126, "6": 0.1521, "375": 0.3259 }, + "6": { + "8": 0.1348, + "4": 0.1448, + "10": 0.1375, + "12": 0.1433, + "906": 0.3429 + }, + "7": { + "2": 0.1256, + "8": 0.1059, + "12": 0.1002, + "14": 0.1503, + "908": 0.3323 + }, + "8": { "7": 0.1059, "2": 0.1102, "6": 0.1348 }, + "9": { "2": 0.1555, "914": 0.2856 }, + "10": { "1": 0.0782, "2": 0.1143, "6": 0.1375, "645": 0.3273 }, + "11": { "3": 0.1406, "917": 0.2553 }, + "12": { + "7": 0.1002, + "4": 0.127, + "6": 0.1433, + "11": 0.1589, + "9": 0.165, + "150": 0.2693 + }, + "14": { "4": 0.1119, "7": 0.1503, "154": 0.37, "389": 0.3789 }, + "15": { + "4": 0.1135, + "11": 0.1436, + "1376": 0.3063, + "776": 0.3038, + "916": 0.3327 + }, + "139": { + "142": 0.0244, + "146": 0.0317, + "151": 0.062, + "145": 0.0753, + "1378": 0.2622 + }, + "141": { "153": 0.0419, "155": 0.0565, "157": 0.0341 }, + "142": { "139": 0.0244, "9": 0.3349, "143": 0.0338, "156": 0.0396 }, + "143": { "142": 0.0338, "153": 0.0347, "150": 0.0429, "155": 0.0326 }, + "144": { + "141": 0.0612, + "2": 0.2946, + "148": 0.0464, + "153": 0.0527, + "156": 0.0717 + }, + "145": { + "139": 0.0753, + "912": 0.2315, + "382": 0.2429, + "635": 0.2429, + "11": 0.2612 + }, + "146": { + "139": 0.0317, + "149": 0.0567, + "768": 0.2357, + "4": 0.3113, + "773": 0.2107 + }, + "148": { "144": 0.0464, "388": 0.2491, "2": 0.2936 }, + "149": { "153": 0.0521 }, + "150": { "143": 0.0429, "12": 0.2693, "908": 0.2783 }, + "151": { "152": 0.0531, "148": 0.1045, "12": 0.2886, "384": 0.2943 }, + "152": { "153": 0.0381, "151": 0.0531 }, + "153": { + "143": 0.0347, + "152": 0.0381, + "141": 0.0419, + "154": 0.0479, + "149": 0.0521, + "144": 0.0527 + }, + "154": { "153": 0.0479, "14": 0.37 }, + "155": { "143": 0.0326, "141": 0.0565, "373": 0.2566, "906": 0.2805 }, + "156": { + "142": 0.0396, + "144": 0.0717, + "772": 0.2231, + "647": 0.2234, + "11": 0.2811 + }, + "157": { + "141": 0.0341, + "148": 0.1041, + "12": 0.3001, + "630": 0.2463, + "642": 0.2505 + }, + "372": { "374": 0.087, "378": 0.0923, "375": 0.1097, "642": 0.2991 }, + "373": { + "381": 0.0666, + "155": 0.2566, + "386": 0.1057, + "387": 0.0724, + "634": 0.2152 + }, + "374": { + "372": 0.087, + "388": 0.0957, + "380": 0.1071, + "376": 0.11, + "389": 0.0773 + }, + "375": { "389": 0.0994, "5": 0.3259, "776": 0.3069 }, + "376": { + "374": 0.11, + "373": 0.132, + "2": 0.3283, + "378": 0.1289, + "383": 0.1093 + }, + "377": { "380": 0.0291, "378": 0.0742, "145": 0.2587, "382": 0.049 }, + "378": { + "382": 0.0514, + "385": 0.07, + "383": 0.0883, + "372": 0.0923, + "386": 0.0596 + }, + "379": { "380": 0.0379, "388": 0.0885, "145": 0.2475 }, + "380": { + "377": 0.0291, + "379": 0.0379, + "387": 0.0428, + "381": 0.0633, + "8": 0.3691 + }, + "381": { + "380": 0.0633, + "373": 0.0666, + "383": 0.0715, + "389": 0.0732, + "145": 0.263 + }, + "382": { + "377": 0.049, + "378": 0.0514, + "145": 0.2429, + "387": 0.0714, + "773": 0.2546 + }, + "383": { + "381": 0.0715, + "378": 0.0883, + "376": 0.1093, + "151": 0.2915, + "387": 0.0596 + }, + "384": { "386": 0.1003, "8": 0.363, "1542": 0.2755 }, + "385": { + "378": 0.07, + "389": 0.0938, + "1342": 0.1669, + "641": 0.2517, + "645": 0.1692 + }, + "386": { "378": 0.0596, "384": 0.1003, "373": 0.1057, "151": 0.283 }, + "387": { + "380": 0.0428, + "383": 0.0596, + "382": 0.0714, + "373": 0.0724, + "374": 0.0969, + "384": 0.1039 + }, + "388": { + "389": 0.0428, + "379": 0.0885, + "633": 0.2256, + "148": 0.2491, + "637": 0.2333 + }, + "389": { + "388": 0.0428, + "381": 0.0732, + "374": 0.0773, + "385": 0.0938, + "375": 0.0994, + "636": 0.1989 + }, + "1342": { "385": 0.1669, "644": 0.1908, "772": 0.2102 }, + "629": { "638": 0.065, "640": 0.0733, "642": 0.0757 }, + "630": { "635": 0.031, "388": 0.2285, "157": 0.2463, "11": 0.2725 }, + "631": { "636": 0.109, "629": 0.127, "9": 0.3283, "647": 0.1202 }, + "632": { "633": 0.0263, "635": 0.029 }, + "633": { + "632": 0.0263, + "640": 0.0758, + "645": 0.1183, + "1378": 0.1209, + "15": 0.2898 + }, + "634": { + "636": 0.0699, + "1377": 0.0909, + "1342": 0.2451, + "913": 0.2455, + "145": 0.2455 + }, + "635": { + "632": 0.029, + "630": 0.031, + "642": 0.0312, + "641": 0.0424, + "1542": 0.1358, + "145": 0.2429 + }, + "636": { + "634": 0.0699, + "1375": 0.1062, + "631": 0.109, + "644": 0.1415, + "389": 0.1989, + "152": 0.2781 + }, + "637": { + "629": 0.0769, + "388": 0.2333, + "156": 0.2464, + "9": 0.3336, + "638": 0.0767 + }, + "638": { + "639": 0.0549, + "629": 0.065, + "637": 0.0767, + "1377": 0.097, + "11": 0.2858 + }, + "639": { + "646": 0.039, + "1378": 0.0402, + "1375": 0.0429, + "638": 0.0549, + "647": 0.0776 + }, + "640": { "629": 0.0733, "633": 0.0758, "646": 0.112 }, + "641": { "635": 0.0424, "639": 0.123, "385": 0.2517, "11": 0.2726 }, + "642": { + "635": 0.0312, + "629": 0.0757, + "157": 0.2505, + "2": 0.2786, + "372": 0.2991 + }, + "644": { + "636": 0.1415, + "1376": 0.153, + "1342": 0.1908, + "773": 0.2105, + "2": 0.2674 + }, + "645": { + "633": 0.1183, + "644": 0.1618, + "385": 0.1692, + "775": 0.2283, + "156": 0.2415, + "10": 0.3273 + }, + "646": { "639": 0.039, "636": 0.1087, "640": 0.112 }, + "647": { "639": 0.0776, "631": 0.1202, "156": 0.2234, "387": 0.2416 }, + "1375": { "639": 0.0429, "636": 0.1062, "15": 0.3378, "911": 0.2759 }, + "1376": { + "644": 0.153, + "1542": 0.173, + "906": 0.2063, + "912": 0.2286, + "15": 0.3063 + }, + "1377": { "634": 0.0909, "638": 0.097, "2": 0.3057 }, + "1378": { + "639": 0.0402, + "636": 0.1161, + "633": 0.1209, + "139": 0.2622, + "915": 0.2596 + }, + "1542": { + "635": 0.1358, + "1376": 0.173, + "384": 0.2755, + "920": 0.2544, + "921": 0.2571 + }, + "767": { "774": 0.0867, "776": 0.1881 }, + "768": { + "773": 0.0772, + "770": 0.0783, + "767": 0.0997, + "774": 0.0539, + "775": 0.0416 + }, + "769": { "775": 0.0617, "767": 0.0923 }, + "770": { "774": 0.0636, "772": 0.0743, "775": 0.0734 }, + "771": { "774": 0.0791, "8": 0.3196, "775": 0.0887, "778": 0.074 }, + "772": { + "770": 0.0743, + "778": 0.0763, + "1342": 0.2102, + "156": 0.2231, + "1376": 0.2568 + }, + "773": { + "774": 0.0697, + "644": 0.2105, + "146": 0.2107, + "382": 0.2546, + "11": 0.2869 + }, + "774": { + "768": 0.0539, + "778": 0.0575, + "770": 0.0636, + "769": 0.069, + "773": 0.0697, + "777": 0.0734 + }, + "775": { + "768": 0.0416, + "769": 0.0617, + "770": 0.0734, + "771": 0.0887, + "645": 0.2283, + "149": 0.2691 + }, + "776": { "767": 0.1881, "15": 0.3038, "375": 0.3069, "778": 0.2063 }, + "777": { "774": 0.0734, "1376": 0.2468, "906": 0.2307 }, + "778": { + "774": 0.0575, + "771": 0.074, + "772": 0.0763, + "776": 0.2063, + "634": 0.254 + }, + "906": { + "1376": 0.2063, + "908": 0.2163, + "777": 0.2307, + "155": 0.2805, + "6": 0.3429 + }, + "908": { "913": 0.1769, "906": 0.2163, "150": 0.2783, "7": 0.3323 }, + "909": { + "912": 0.1011, + "908": 0.1823, + "11": 0.259, + "913": 0.0768, + "917": 0.0939 + }, + "910": { "917": 0.0941, "1376": 0.2509 }, + "911": { "912": 0.0777, "1375": 0.2759, "913": 0.0731, "917": 0.0609 }, + "912": { + "917": 0.0627, + "913": 0.0726, + "918": 0.0766, + "919": 0.0811, + "1376": 0.2286, + "145": 0.2315 + }, + "913": { + "912": 0.0726, + "911": 0.0731, + "915": 0.0762, + "909": 0.0768, + "916": 0.0891, + "908": 0.1769 + }, + "914": { + "913": 0.1707, + "9": 0.2856, + "915": 0.1648, + "916": 0.158, + "919": 0.1749 + }, + "915": { + "920": 0.0742, + "913": 0.0762, + "917": 0.0766, + "914": 0.1648, + "1378": 0.2596 + }, + "916": { + "913": 0.0891, + "914": 0.158, + "906": 0.2354, + "15": 0.3327, + "921": 0.097 + }, + "917": { + "911": 0.0609, + "912": 0.0627, + "915": 0.0766, + "909": 0.0939, + "910": 0.0941, + "11": 0.2553 + }, + "918": { "912": 0.0766, "906": 0.2414, "920": 0.0829 }, + "919": { "912": 0.0811, "914": 0.1749, "12": 0.2813 }, + "920": { + "915": 0.0742, + "918": 0.0829, + "906": 0.2216, + "1542": 0.2544, + "11": 0.2662 + }, + "921": { "916": 0.097, "908": 0.1775, "1542": 0.2571, "12": 0.3105 } + }, + { + "141": { "149": 0.0587, "144": 0.0612 }, + "144": { "148": 0.0464, "141": 0.0612, "156": 0.0717 }, + "145": { "146": 0.0874, "635": 0.2429, "378": 0.2676 }, + "146": { "149": 0.0567, "156": 0.0577, "145": 0.0874 }, + "148": { "144": 0.0464 }, + "149": { "146": 0.0567, "141": 0.0587 }, + "156": { "146": 0.0577, "144": 0.0717, "772": 0.2231 }, + "376": { "378": 0.1289, "1377": 0.276 }, + "378": { "376": 0.1289, "145": 0.2676 }, + "630": { "145": 0.2528, "635": 0.031 }, + "635": { "630": 0.031, "646": 0.1228, "145": 0.2429 }, + "637": { "646": 0.1048, "156": 0.2464 }, + "646": { "1377": 0.1042, "637": 0.1048, "635": 0.1228 }, + "1376": { "1377": 0.166, "145": 0.243, "910": 0.2509 }, + "1377": { "646": 0.1042, "1376": 0.166, "376": 0.276 }, + "772": { "156": 0.2231, "1376": 0.2568 }, + "910": { "1376": 0.2509, "914": 0.1729 }, + "914": { "910": 0.1729 } + }, + { + "149": { "1376": 0.2829, "914": 0.3343 }, + "1376": { "149": 0.2829, "914": 0.2723 }, + "914": { "149": 0.3343, "1376": 0.2723 } + } +] diff --git a/test/mememo.test.ts b/test/mememo.test.ts index 030be43..d1890a3 100644 --- a/test/mememo.test.ts +++ b/test/mememo.test.ts @@ -6,6 +6,7 @@ import graph10Layer1JSON from './data/insert-10-1-layer.json'; import graph10Layer2JSON from './data/insert-10-2-layer.json'; import graph30Layer3JSON from './data/insert-30-3-layer.json'; import graph100Layer6JSON from './data/insert-100-6-layer.json'; +import graph100Layer3M3JSON from './data/insert-100-3-layer-m=3.json'; interface EmbeddingData { embeddings: number[][]; @@ -19,6 +20,7 @@ const graph10Layer1 = graph10Layer1JSON as GraphLayer[]; const graph10Layer2 = graph10Layer2JSON as GraphLayer[]; const graph30Layer3 = graph30Layer3JSON as GraphLayer[]; const graph100Layer6 = graph100Layer6JSON as GraphLayer[]; +const graph100Layer3M3 = graph100Layer3M3JSON as GraphLayer[]; describe('constructor', () => { it('constructor', () => { @@ -35,7 +37,7 @@ describe('constructor', () => { //==========================================================================|| describe('insert()', () => { - it.skip('insert() 10 items, 1 layer', () => { + it('insert() 10 items, 1 layer', () => { const hnsw = new HNSW({ distanceFunction: 'cosine', seed: 20240101 @@ -72,7 +74,7 @@ describe('insert()', () => { } }); - it.skip('insert() 10 items, 2 layer', () => { + it('insert() 10 items, 2 layer', () => { const hnsw = new HNSW({ distanceFunction: 'cosine', seed: 10 @@ -113,7 +115,7 @@ describe('insert()', () => { } }); - it.skip('insert() 30 items, 3 layer', () => { + it('insert() 30 items, 3 layer', () => { const hnsw = new HNSW({ distanceFunction: 'cosine', seed: 262 @@ -200,6 +202,56 @@ describe('insert()', () => { } }); + it('insert() 100 items, 3 layer, m=3', () => { + const hnsw = new HNSW({ + distanceFunction: 'cosine', + seed: 21574, + m: 3 + }); + + // Insert 100 embeddings + const size = 100; + + // The random levels with seed 21574 (need to manually set it, because it + // would change since m and ml are different from default) + const levels = [ + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 1, 1, 2, 0, 0, + 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 2, 1, 0, 0, 0, 0, + 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0 + ]; + + const reportIDs: string[] = []; + for (let i = 0; i < size; i++) { + const curReportID = String(embeddingData.reportNumbers[i]); + reportIDs.push(curReportID); + hnsw.insert(curReportID, embeddingData.embeddings[i], levels[i]); + } + + expect(hnsw.graphLayers.length).toBe(3); + + for (const reportID of reportIDs) { + for (const [l, graphLayer] of hnsw.graphLayers.entries()) { + const curNode = graphLayer.graph.get(reportID); + + if (curNode === undefined) { + expect(graph100Layer3M3[l][reportID]).toBeUndefined(); + } else { + expect(graph100Layer3M3[l][reportID]).not.to.toBeUndefined(); + // Check the distances + const expectedNeighbors = graph100Layer3M3[l][reportID]; + for (const [neighborKey, neighborDistance] of curNode.entries()) { + expect(expectedNeighbors[neighborKey]).to.not.toBeUndefined(); + expect(neighborDistance).toBeCloseTo( + expectedNeighbors[neighborKey]!, + 1e-6 + ); + } + } + } + } + }); + it.skip('Find random seeds', () => { // Find random seed that give a nice level sequence to test const size = 100; @@ -213,9 +265,9 @@ describe('insert()', () => { curLevels.push(level); } - if (curLevels.includes(5)) { + if (Math.max(...curLevels) < 4) { const levelSum = curLevels.reduce((sum, value) => sum + value, 0); - if (levelSum > 10) { + if (levelSum > 20) { console.log('Good seed: ', i); break; }