Skip to content

Commit

Permalink
Test inserting 100 items 3 layer with m=3
Browse files Browse the repository at this point in the history
Signed-off-by: Jay Wang <[email protected]>
  • Loading branch information
xiaohk committed Jan 30, 2024
1 parent f8601f2 commit f0e61b3
Show file tree
Hide file tree
Showing 3 changed files with 63 additions and 8 deletions.
8 changes: 5 additions & 3 deletions src/mememo.ts
Original file line number Diff line number Diff line change
Expand Up @@ -186,17 +186,19 @@ export class HNSW<T = string> {
* 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
return;
}

// 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));
Expand Down
1 change: 1 addition & 0 deletions test/data/insert-100-3-layer-m=3.json
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
[{"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}}]
62 changes: 57 additions & 5 deletions test/mememo.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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[][];
Expand All @@ -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', () => {
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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;
Expand All @@ -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;
}
Expand Down

0 comments on commit f0e61b3

Please sign in to comment.