gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

multi batches

Browse Source
This commit is contained in:
Yannick Stephan
2025-02-19 23:53:25 +01:00
parent 4875283d9f
commit efdc8a2d26
1 changed files with 78 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
lightrag/lightrag.py
View File

@@ -792,11 +792,17 @@ class LightRAG:
] ]
logger.info(f"Number of batches to process: {len(docs_batches)}.") logger.info(f"Number of batches to process: {len(docs_batches)}.")
batches: list[Any] = []
# 3. iterate over batches # 3. iterate over batches
for batch_idx, docs_batch in enumerate(docs_batches): for batch_idx, docs_batch in enumerate(docs_batches):
logger.info(
f"Start processing batch {batch_idx + 1} of {len(docs_batches)}." async def batch(
) batch_idx: int,
docs_batch: list[tuple[str, DocProcessingStatus]],
size_batch: int,
) -> None:
logger.info(f"Start processing batch {batch_idx + 1} of {size_batch}.")
# 4. iterate over batch # 4. iterate over batch
for doc_id_processing_status in docs_batch: for doc_id_processing_status in docs_batch:
doc_id, status_doc = doc_id_processing_status doc_id, status_doc = doc_id_processing_status
@@ -834,7 +840,9 @@ class LightRAG:
tasks = [ tasks = [
self.chunks_vdb.upsert(chunks), self.chunks_vdb.upsert(chunks),
self._process_entity_relation_graph(chunks), self._process_entity_relation_graph(chunks),
self.full_docs.upsert({doc_id: {"content": status_doc.content}}), self.full_docs.upsert(
{doc_id: {"content": status_doc.content}}
),
self.text_chunks.upsert(chunks), self.text_chunks.upsert(chunks),
self.doc_status.upsert( self.doc_status.upsert(
{ {
@@ -852,7 +860,6 @@ class LightRAG:
] ]
try: try:
await asyncio.gather(*tasks) await asyncio.gather(*tasks)
await self._insert_done()
except Exception as e: except Exception as e:
logger.error(f"Failed to process document {doc_id}: {str(e)}") logger.error(f"Failed to process document {doc_id}: {str(e)}")
@@ -872,6 +879,11 @@ class LightRAG:
continue continue
logger.info(f"Completed batch {batch_idx + 1} of {len(docs_batches)}.") logger.info(f"Completed batch {batch_idx + 1} of {len(docs_batches)}.")
batches.append(batch(batch_idx, docs_batch, len(docs_batches)))
await asyncio.gather(*batches)
await self._insert_done()
async def _process_entity_relation_graph(self, chunk: dict[str, Any]) -> None: async def _process_entity_relation_graph(self, chunk: dict[str, Any]) -> None:
try: try:
new_kg = await extract_entities( new_kg = await extract_entities(