Spaces:
Sleeping
Sleeping
| from collections.abc import Callable, Sequence | |
| from concurrent.futures import ThreadPoolExecutor | |
| from dataclasses import dataclass | |
| from techdocs_graphrag.extraction import GraphExtractor | |
| from techdocs_graphrag.ingestion import ParagraphSplitter, SourceDocument, chunk_documents | |
| from techdocs_graphrag.ports import KnowledgeGraph, LanguageModel | |
| ProgressCallback = Callable[[int, int], None] | |
| class IndexingReport: | |
| documents: int | |
| chunks: int | |
| entities: int | |
| relations: int | |
| class IndexingPipeline: | |
| def __init__( | |
| self, | |
| *, | |
| splitter: ParagraphSplitter, | |
| llm: LanguageModel, | |
| extractor: GraphExtractor, | |
| store: KnowledgeGraph, | |
| workers: int, | |
| ) -> None: | |
| self._splitter = splitter | |
| self._llm = llm | |
| self._extractor = extractor | |
| self._store = store | |
| self._workers = workers | |
| def run( | |
| self, | |
| documents: Sequence[SourceDocument], | |
| on_progress: ProgressCallback | None = None, | |
| ) -> IndexingReport: | |
| chunks = chunk_documents(documents, self._splitter) | |
| if not chunks: | |
| return IndexingReport(documents=len(documents), chunks=0, entities=0, relations=0) | |
| self._store.ensure_schema() | |
| self._store.upsert_chunks(chunks, self._llm.embed([chunk.text for chunk in chunks])) | |
| entities = relations = 0 | |
| with ThreadPoolExecutor(max_workers=self._workers) as executor: | |
| extractions = executor.map(self._extractor.extract, chunks) | |
| for completed, (chunk, extraction) in enumerate( | |
| zip(chunks, extractions, strict=True), start=1 | |
| ): | |
| self._store.upsert_extraction(chunk, extraction) | |
| entities += len(extraction.entities) | |
| relations += len(extraction.relations) | |
| if on_progress: | |
| on_progress(completed, len(chunks)) | |
| self._store.await_indexes() | |
| return IndexingReport( | |
| documents=len(documents), | |
| chunks=len(chunks), | |
| entities=entities, | |
| relations=relations, | |
| ) | |