You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@hugegraph.apache.org by ji...@apache.org on 2026/10/03 07:38:03 UTC
(hugegraph-ai) branch main updated: fix(llm): preserve semantic embedding batch order
This is an automated email from the ASF dual-hosted git repository.
imbajin pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/hugegraph-ai.git
The following commit(s) were added to refs/heads/main by this push:
new 99225684 fix(llm): preserve semantic embedding batch order
99225684 is described below
commit 992256845370a9812d29ca7d43c282f1783af920
Author: JesusMan0529 <16...@qq.com>
AuthorDate: Sat Oct 3 03:54:12 2026 +0800
fix(llm): preserve semantic embedding batch order
---
.../operators/index_op/build_semantic_index.py | 15 ++++---
.../index_op/test_build_semantic_index.py | 46 ++++++++++++++++++++++
2 files changed, 53 insertions(+), 8 deletions(-)
diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_semantic_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_semantic_index.py
index 1c75ea4b..be8282d0 100644
--- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_semantic_index.py
+++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_semantic_index.py
@@ -40,20 +40,19 @@ class BuildSemanticIndex:
sem = asyncio.Semaphore(10)
batch_size = 1000
- async def get_embeddings_with_semaphore(vid_list: list[str]) -> Any:
+ async def get_embeddings_with_semaphore(vid_list: list[str], pbar: tqdm) -> Any:
async with sem:
loop = asyncio.get_running_loop()
- return await loop.run_in_executor(None, self.embedding.get_texts_embeddings, vid_list)
+ result = await loop.run_in_executor(None, self.embedding.get_texts_embeddings, vid_list)
+ pbar.update(1)
+ return result
vid_batches = [vids[i : i + batch_size] for i in range(0, len(vids), batch_size)]
- tasks = [get_embeddings_with_semaphore(batch) for batch in vid_batches]
-
embeddings = []
- with tqdm(total=len(tasks)) as pbar:
- for future in asyncio.as_completed(tasks):
- batch_embeddings = await future
+ with tqdm(total=len(vid_batches)) as pbar:
+ tasks = [get_embeddings_with_semaphore(batch, pbar) for batch in vid_batches]
+ for batch_embeddings in await asyncio.gather(*tasks):
embeddings.extend(batch_embeddings)
- pbar.update(1)
return embeddings
def run(self, context: Dict[str, Any]) -> Dict[str, Any]:
diff --git a/hugegraph-llm/src/tests/operators/index_op/test_build_semantic_index.py b/hugegraph-llm/src/tests/operators/index_op/test_build_semantic_index.py
index 2befaa38..4ec66f5f 100644
--- a/hugegraph-llm/src/tests/operators/index_op/test_build_semantic_index.py
+++ b/hugegraph-llm/src/tests/operators/index_op/test_build_semantic_index.py
@@ -21,10 +21,12 @@ import asyncio
import os
import shutil
import tempfile
+import threading
import unittest
from unittest.mock import MagicMock, patch
from hugegraph_llm.indices.vector_index.base import VectorStoreBase
+from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex
from hugegraph_llm.models.embeddings.base import BaseEmbedding
from hugegraph_llm.operators.index_op.build_semantic_index import BuildSemanticIndex
@@ -114,6 +116,50 @@ class TestBuildSemanticIndex(unittest.TestCase):
finally:
loop.close()
+ def test_get_embeddings_parallel_empty(self):
+ builder = BuildSemanticIndex(self.mock_embedding, self.mock_vector_store_class)
+
+ self.assertEqual(asyncio.run(builder._get_embeddings_parallel([])), [])
+ self.mock_embedding.get_texts_embeddings.assert_not_called()
+
+ def test_get_embeddings_parallel_propagates_errors(self):
+ builder = BuildSemanticIndex(self.mock_embedding, self.mock_vector_store_class)
+ self.mock_embedding.get_texts_embeddings.side_effect = RuntimeError("embedding failed")
+
+ with self.assertRaisesRegex(RuntimeError, "embedding failed"):
+ asyncio.run(builder._get_embeddings_parallel(["vid1"]))
+
+ def test_run_preserves_embedding_alignment_across_batches(self):
+ vertices = [f"person:vid{i}" for i in range(1001)]
+ self.mock_embedding.get_embedding_dim.return_value = 1
+
+ for strategy in ("PRIMARY_KEY", "CUSTOMIZE"):
+ with self.subTest(strategy=strategy):
+ self.mock_settings.graph_name = os.path.join(self.temp_dir, strategy)
+ self.mock_schema_manager.schema.getSchema.return_value = {"vertexlabels": [{"id_strategy": strategy}]}
+ second_batch_completed = threading.Event()
+
+ def embed_batch(batch):
+ if len(batch) == 1000:
+ if not second_batch_completed.wait(timeout=5):
+ raise RuntimeError("The second batch did not complete")
+ return [[float(vid.rsplit("vid", 1)[1])] for vid in batch]
+
+ self.mock_embedding.get_texts_embeddings.side_effect = embed_batch
+ builder = BuildSemanticIndex(self.mock_embedding, FaissVectorIndex)
+
+ with patch("hugegraph_llm.operators.index_op.build_semantic_index.tqdm") as mock_tqdm:
+ progress = mock_tqdm.return_value.__enter__.return_value
+ # Release the first batch only after the second batch reports completion.
+ progress.update.side_effect = lambda _: second_batch_completed.set()
+ result = builder.run({"vertices": vertices})
+
+ self.assertEqual(result["added_vid_vector_num"], len(vertices))
+ self.assertEqual(progress.update.call_count, 2)
+ loaded_index = FaissVectorIndex.from_name(1, self.mock_settings.graph_name, "graph_vids")
+ for i in (0, 500, 1000):
+ self.assertEqual(loaded_index.search([float(i)], top_k=1), [f"person:vid{i}"])
+
def test_run_with_primary_key_strategy(self):
# Create a builder
builder = BuildSemanticIndex(self.mock_embedding, self.mock_vector_store_class)