Spaces:
Sleeping
Sleeping
File size: 2,164 Bytes
da21448 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | 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]
@dataclass(frozen=True, slots=True)
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,
)
|