TBNilles/LEANN
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial commit

Browse Source
This commit is contained in:
yichuan520030910320
2025-06-30 11:01:12 +00:00
parent 30898814ae
commit ee507bfe7a
5 changed files with 1057 additions and 295 deletions
Download Patch File Download Diff File Expand all files Collapse all files

206
demo.ipynb
View File

@@ -2,7 +2,7 @@
"cells": [
{
"cell_type": "code",
"execution_count": 3,
"execution_count": 2,
"metadata": {},
"outputs": [
{
@@ -17,7 +17,7 @@
"name": "stderr",
"output_type": "stream",
"text": [
"Batches: 100%|██████████| 1/1 [00:00<00:00, 77.61it/s]"
"Batches: 100%|██████████| 1/1 [00:00<00:00, 89.83it/s]"
]
},
{
@@ -27,11 +27,10 @@
"INFO: Building DiskANN index for 6 vectors with metric Metric.INNER_PRODUCT...\n",
"Using Inner Product search, so need to pre-process base data into temp file. Please ensure there is additional (n*(d+1)*4) bytes for storing pre-processed base vectors, apart from the interim indices created by DiskANN and the final index.\n",
"Pre-processing base file by adding extra coordinate\n",
"✅ DiskANN index built successfully at 'knowledge'\n",
"Writing bin: knowledge_disk.index_max_base_norm.bin\n",
"bin: #pts = 1, #dims = 1, size = 12B\n",
"Finished writing bin.\n",
"Time for preprocessing data for inner product: 0.000165 seconds\n",
"Time for preprocessing data for inner product: 0.000186 seconds\n",
"Reading max_norm_of_base from knowledge_disk.index_max_base_norm.bin\n",
"Reading bin file knowledge_disk.index_max_base_norm.bin ...\n",
"Opening bin file knowledge_disk.index_max_base_norm.bin... \n",
@@ -50,7 +49,8 @@
"Metadata: #pts = 256, #dims = 769...\n",
"done.\n",
"PQ pivot file exists. Not generating again\n",
"Opened: knowledge_prepped_base.bin, size: 18464, cache_size: 18464\n",
"Opened: knowledge_prep✅ DiskANN index built successfully at 'knowledge'\n",
"ped_base.bin, size: 18464, cache_size: 18464\n",
"Reading bin file knowledge_pq_pivots.bin ...\n",
"Opening bin file knowledge_pq_pivots.bin... \n",
"Metadata: #pts = 4, #dims = 1...\n",
@@ -69,17 +69,17 @@
"done.\n",
"Loaded PQ pivot information\n",
"Processing points [0, 6)...done.\n",
"Time for generating quantized data: 0.023918 seconds\n",
"Time for generating quantized data: 0.016337 seconds\n",
"Full index fits in RAM budget, should consume at most 2.03973e-05GiBs, so building in one shot\n",
"L2: Using AVX2 distance computation DistanceL2Float\n",
"Passed, empty search_params while creating index config\n",
"Using only first 6 from file.. \n",
"Starting index build with 6 points... \n",
"0% of index build completed.Starting final cleanup..done. Link time: 9e-05s\n",
"0% of index build completed.Starting final cleanup..done. Link time: 6.2e-05s\n",
"Index built with degree: max:5 avg:5 min:5 count(deg<2):0\n",
"Not saving tags as they are not enabled.\n",
"Time taken for save: 0.000178s.\n",
"Time for building merged vamana index: 0.000579 seconds\n",
"Time taken for save: 0.000156s.\n",
"Time for building merged vamana index: 0.000549 seconds\n",
"Opened: knowledge_prepped_base.bin, size: 18464, cache_size: 18464\n",
"Vamana index file size=168\n",
"Opened: knowledge_disk.index, cache_size: 67108864\n",
@@ -94,11 +94,11 @@
"Finished writing bin.\n",
"Output disk index file written to knowledge_disk.index\n",
"Finished writing 28672B\n",
"Time for generating disk layout: 0.043488 seconds\n",
"Time for generating disk layout: 0.032297 seconds\n",
"Opened: knowledge_prepped_base.bin, size: 18464, cache_size: 18464\n",
"Loading base knowledge_prepped_base.bin. #points: 6. #dim: 769.\n",
"Wrote 1 points to sample file: knowledge_sample_data.bin\n",
"Indexing time: 0.0684344\n",
"Indexing time: 0.0495994\n",
"INFO: Leann metadata saved to knowledge.leann.meta.json\n"
]
},
@@ -120,6 +120,7 @@
"Before index load\n",
"✅ DiskANN index loaded successfully.\n",
"INFO: LeannSearcher initialized with 'diskann' backend using index 'knowledge.leann'.\n",
"INFO: Terminating session server process (PID: 70114)...\n",
"Reading bin file knowledge_pq_compressed.bin ...\n",
"Opening bin file knowledge_pq_compressed.bin... \n",
"Metadata: #pts = 6, #dims = 512...\n",
@@ -147,14 +148,14 @@
"Disk-Index File Meta-data: # nodes per sector: 1, max node len (bytes): 3100, max node degree: 5\n",
"Disk-Index Meta: nodes per sector: 1, max node len: 3100, max node degree: 5\n",
"Setting up thread-specific contexts for nthreads: 8\n",
"allocating ctx: 0x78348f4de000 to thread-id:132170359560000\n",
"allocating ctx: 0x78348f4cd000 to thread-id:132158431693760\n",
"allocating ctx: 0x78348f4bc000 to thread-id:132158442179392\n",
"allocating ctx: 0x78348f4ab000 to thread-id:132158421208128\n",
"allocating ctx: 0x78348f49a000 to thread-id:132158452665024\n",
"allocating ctx: 0x78348f489000 to thread-id:132158389751232\n",
"allocating ctx: 0x78348f478000 to thread-id:132158410722496\n",
"allocating ctx: 0x78348f467000 to thread-id:132158400236864\n",
"allocating ctx: 0x73ab1e40c000 to thread-id:127182044829504\n",
"allocating ctx: 0x73ab1baef000 to thread-id:127170028506048\n",
"allocating ctx: 0x73ab1bade000 to thread-id:127170038991680\n",
"allocating ctx: 0x73ab1bacd000 to thread-id:127170049477312\n",
"allocating ctx: 0x73ab1babc000 to thread-id:127170018020416\n",
"allocating ctx: 0x73ab1baab000 to thread-id:127169986563520\n",
"allocating ctx: 0x73ab1ba9a000 to thread-id:127170007534784\n",
"allocating ctx: 0x73ab1ba89000 to thread-id:127169997049152\n",
"Loading centroid data from medoids vector data of 1 medoid(s)\n",
"Reading bin file knowledge_disk.index_max_base_norm.bin ...\n",
"Opening bin file knowledge_disk.index_max_base_norm.bin... \n",
@@ -167,6 +168,7 @@
"Graph traversal completed, hops: 3\n",
"Loading the cache list into memory....done.\n",
"After index load\n",
"INFO: Server process terminated.\n",
"Clearing scratch\n",
"INFO: Computing embeddings for 1 chunks using 'sentence-transformers/all-mpnet-base-v2'...\n"
]
@@ -175,15 +177,17 @@
"name": "stderr",
"output_type": "stream",
"text": [
"Batches: 100%|██████████| 1/1 [00:00<00:00, 92.66it/s]"
"Batches: 100%|██████████| 1/1 [00:00<00:00, 96.44it/s]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Score: -0.481 - C++ is a powerful programming language\n",
"Score: -1.049 - Java is a powerful programming language\n"
"INFO: DiskANN ZMQ mode enabled - ensuring embedding server is running\n",
"INFO: Starting session-level embedding server as a background process...\n",
"INFO: Running command from project root: /home/ubuntu/leann_release/Power-RAG\n",
"INFO: Server process started with PID: 71195\n"
]
},
{
@@ -197,9 +201,165 @@
"name": "stdout",
"output_type": "stream",
"text": [
"✅ Embedding server is up and ready for this session.\n",
"reserve ratio: 1\n",
"[EmbeddingServer LOG]: Initializing leann-backend-diskann...\n",
"[EmbeddingServer LOG]: WARNING: Could not import DiskANN backend: cannot import name '_diskannpy' from partially initialized module 'packages.leann-backend-diskann.leann_backend_diskann' (most likely due to a circular import) (/home/ubuntu/leann_release/Power-RAG/packages/leann-backend-diskann/leann_backend_diskann/__init__.py)\n",
"[EmbeddingServer LOG]: INFO: Initializing embedding server thread on port 5555\n",
"[EmbeddingServer LOG]: INFO: Using CUDA device\n",
"[EmbeddingServer LOG]: INFO: Loading model sentence-transformers/all-mpnet-base-v2\n",
"[EmbeddingServer LOG]: INFO: Using FP16 precision with model: sentence-transformers/all-mpnet-base-v2\n",
"[EmbeddingServer LOG]: INFO: Loaded 6 demo documents\n",
"[EmbeddingServer LOG]: INFO: ZMQ ROUTER server listening on port 5555\n",
"[EmbeddingServer LOG]: INFO: Embedding server ready to serve requests\n",
"Score: -0.481 - C++ is a powerful programming language\n",
"Score: -1.049 - Java is a powerful programming language\n",
"[EmbeddingServer LOG]: INFO: Received ZMQ request from client 006b8b45, size 3 bytes\n",
"[EmbeddingServer LOG]: INFO: Request for 1 node embeddings: [0]\n",
"[EmbeddingServer LOG]: DEBUG: Node ID range: 0 to 0\n",
"[EmbeddingServer LOG]: Time taken for text lookup: 0.000028 seconds\n",
"[EmbeddingServer LOG]: INFO: Total batch size: 1, max_batch_size: 128\n",
"[EmbeddingServer LOG]: INFO: Processing batch of size 1\n",
"[EmbeddingServer LOG]: Time taken for tokenization (batch): 0.019126 seconds\n",
"[EmbeddingServer LOG]: Batch size: 1, Sequence length: 256\n",
"[EmbeddingServer LOG]: Time taken for transfer to device (batch): 0.000246 seconds\n",
"[EmbeddingServer LOG]: Time taken for embedding (batch): 3.216055 seconds\n",
"[EmbeddingServer LOG]: Time taken for mean pooling (batch): 0.044091 seconds\n",
"[EmbeddingServer LOG]: INFO: Serialize time: 0.000299 seconds\n",
"[EmbeddingServer LOG]: INFO: ZMQ E2E time: 3.281315 seconds\n",
"[EmbeddingServer LOG]: INFO: Received ZMQ request from client 006b8b45, size 7 bytes\n",
"[EmbeddingServer LOG]: INFO: Request for 5 node embeddings: [1, 2, 3, 4, 5]\n",
"[EmbeddingServer LOG]: DEBUG: Node ID range: 1 to 5\n",
"[EmbeddingServer LOG]: Time taken for text lookup: 0.000048 seconds\n",
"[EmbeddingServer LOG]: INFO: Total batch size: 5, max_batch_size: 128\n",
"[EmbeddingServer LOG]: INFO: Processing batch of size 5\n",
"[EmbeddingServer LOG]: Time taken for tokenization (batch): 0.001797 seconds\n",
"[EmbeddingServer LOG]: Batch size: 5, Sequence length: 256\n",
"[EmbeddingServer LOG]: Time taken for transfer to device (batch): 0.000095 seconds\n",
"[EmbeddingServer LOG]: Time taken for embedding (batch): 3.392574 seconds\n",
"[EmbeddingServer LOG]: Time taken for mean pooling (batch): 0.000353 seconds\n",
"[EmbeddingServer LOG]: INFO: Serialize time: 0.000188 seconds\n",
"[EmbeddingServer LOG]: INFO: ZMQ E2E time: 3.395764 seconds\n",
"[EmbeddingServer LOG]: INFO: Received ZMQ request from client 006b8b45, size 7 bytes\n",
"[EmbeddingServer LOG]: INFO: Request for 5 node embeddings: [3, 4, 2, 1, 0]\n",
"[EmbeddingServer LOG]: DEBUG: Node ID range: 0 to 4\n",
"[EmbeddingServer LOG]: Time taken for text lookup: 0.000048 seconds\n",
"[EmbeddingServer LOG]: INFO: Total batch size: 5, max_batch_size: 128\n",
"[EmbeddingServer LOG]: INFO: Processing batch of size 5\n",
"[EmbeddingServer LOG]: Time taken for tokenization (batch): 0.001377 seconds\n",
"[EmbeddingServer LOG]: Batch size: 5, Sequence length: 256\n",
"[EmbeddingServer LOG]: Time taken for transfer to device (batch): 0.000085 seconds\n",
"[EmbeddingServer LOG]: Time taken for embedding (batch): 0.009257 seconds\n",
"[EmbeddingServer LOG]: Time taken for mean pooling (batch): 0.000166 seconds\n",
"[EmbeddingServer LOG]: INFO: Serialize time: 0.000074 seconds\n",
"[EmbeddingServer LOG]: INFO: ZMQ E2E time: 0.011520 seconds\n",
"[EmbeddingServer LOG]: INFO: Received ZMQ request from client 006b8b45, size 7 bytes\n",
"[EmbeddingServer LOG]: INFO: Request for 5 node embeddings: [0, 1, 2, 4, 5]\n",
"[EmbeddingServer LOG]: DEBUG: Node ID range: 0 to 5\n",
"[EmbeddingServer LOG]: Time taken for text lookup: 0.000019 seconds\n",
"[EmbeddingServer LOG]: INFO: Total batch size: 5, max_batch_size: 128\n",
"[EmbeddingServer LOG]: INFO: Processing batch of size 5\n",
"[EmbeddingServer LOG]: Time taken for tokenization (batch): 0.000792 seconds\n",
"[EmbeddingServer LOG]: Batch size: 5, Sequence length: 256\n",
"[EmbeddingServer LOG]: Time taken for transfer to device (batch): 0.000073 seconds\n",
"[EmbeddingServer LOG]: Time taken for embedding (batch): 0.008864 seconds\n",
"[EmbeddingServer LOG]: Time taken for mean pooling (batch): 0.000155 seconds\n",
"[EmbeddingServer LOG]: INFO: Serialize time: 0.000052 seconds\n",
"[EmbeddingServer LOG]: INFO: ZMQ E2E time: 0.010397 seconds\n",
"[EmbeddingServer LOG]: INFO: Received ZMQ request from client 006b8b45, size 7 bytes\n",
"[EmbeddingServer LOG]: INFO: Request for 5 node embeddings: [3, 1, 0, 2, 5]\n",
"[EmbeddingServer LOG]: DEBUG: Node ID range: 0 to 5\n",
"[EmbeddingServer LOG]: Time taken for text lookup: 0.000020 seconds\n",
"[EmbeddingServer LOG]: INFO: Total batch size: 5, max_batch_size: 128\n",
"[EmbeddingServer LOG]: INFO: Processing batch of size 5\n",
"[EmbeddingServer LOG]: Time taken for tokenization (batch): 0.000857 seconds\n",
"[EmbeddingServer LOG]: Batch size: 5, Sequence length: 256\n",
"[EmbeddingServer LOG]: Time taken for transfer to device (batch): 0.000073 seconds\n",
"[EmbeddingServer LOG]: Time taken for embedding (batch): 0.008830 seconds\n",
"[EmbeddingServer LOG]: Time taken for mean pooling (batch): 0.000143 seconds\n",
"[EmbeddingServer LOG]: INFO: Serialize time: 0.000053 seconds\n",
"[EmbeddingServer LOG]: INFO: ZMQ E2E time: 0.010439 seconds\n",
"[EmbeddingServer LOG]: INFO: Received ZMQ request from client 006b8b45, size 7 bytes\n",
"[EmbeddingServer LOG]: INFO: Request for 5 node embeddings: [0, 2, 3, 4, 5]\n",
"[EmbeddingServer LOG]: DEBUG: Node ID range: 0 to 5\n",
"[EmbeddingServer LOG]: Time taken for text lookup: 0.000020 seconds\n",
"[EmbeddingServer LOG]: INFO: Total batch size: 5, max_batch_size: 128\n",
"[EmbeddingServer LOG]: INFO: Processing batch of size 5\n",
"[EmbeddingServer LOG]: Time taken for tokenization (batch): 0.000805 seconds\n",
"[EmbeddingServer LOG]: Batch size: 5, Sequence length: 256\n",
"[EmbeddingServer LOG]: Time taken for transfer to device (batch): 0.000072 seconds\n",
"[EmbeddingServer LOG]: Time taken for embedding (batch): 0.008835 seconds\n",
"[EmbeddingServer LOG]: Time taken for mean pooling (batch): 0.000141 seconds\n",
"[EmbeddingServer LOG]: INFO: Serialize time: 0.000049 seconds\n",
"[EmbeddingServer LOG]: INFO: ZMQ E2E time: 0.010386 seconds\n",
"[EmbeddingServer LOG]: INFO: Received ZMQ request from client 006b8b45, size 7 bytes\n",
"[EmbeddingServer LOG]: INFO: Request for 5 node embeddings: [1, 0, 3, 4, 5]\n",
"[EmbeddingServer LOG]: DEBUG: Node ID range: 0 to 5\n",
"[EmbeddingServer LOG]: Time taken for text lookup: 0.000019 seconds\n",
"[EmbeddingServer LOG]: INFO: Total batch size: 5, max_batch_size: 128\n",
"[EmbeddingServer LOG]: INFO: Processing batch of size 5\n",
"[EmbeddingServer LOG]: Time taken for tokenization (batch): 0.000826 seconds\n",
"[EmbeddingServer LOG]: Batch size: 5, Sequence length: 256\n",
"[EmbeddingServer LOG]: Time taken for transfer to device (batch): 0.000092 seconds\n",
"[EmbeddingServer LOG]: Time taken for embedding (batch): 0.008809 seconds\n",
"[EmbeddingServer LOG]: Time taken for mean pooling (batch): 0.000164 seconds\n",
"[EmbeddingServer LOG]: INFO: Serialize time: 0.000048 seconds\n",
"[EmbeddingServer LOG]: INFO: ZMQ E2E time: 0.010431 seconds\n",
"Graph traversal completed, hops: 3\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n",
"[EmbeddingServer LOG]: INFO: ZMQ socket timeout, continuing to listen\n"
]
},
{
"ename": "",
"evalue": "",
"output_type": "error",
"traceback": [
"\u001b[1;31mThe Kernel crashed while executing code in the current cell or a previous cell. \n",
"\u001b[1;31mPlease review the code in the cell(s) to identify a possible cause of the failure. \n",
"\u001b[1;31mClick <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. \n",
"\u001b[1;31mView Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
]
}
],
"source": [
@@ -217,7 +377,7 @@
"\n",
"# 2. Search with real-time embeddings\n",
"searcher = LeannSearcher(\"knowledge.leann\")\n",
"results = searcher.search(\"C++ programming languages\", top_k=2)\n",
"results = searcher.search(\"C++ programming languages\", top_k=2,recompute_beighbor_embeddings=True)\n",
"\n",
"for result in results:\n",
" print(f\"Score: {result['score']:.3f} - {result['text']}\")"