Spaces:
Sleeping
Sleeping
File size: 1,563 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 | from collections import Counter
from collections.abc import Sequence
from techdocs_graphrag.models import RetrievalContext, RetrievalMode, RetrievedChunk
from techdocs_graphrag.ports import KnowledgeGraph, LanguageModel
MAX_LEXICAL_SEEDS = 5
MAX_MENTION_SEEDS = 8
class HybridRetriever:
def __init__(
self,
*,
llm: LanguageModel,
store: KnowledgeGraph,
top_k: int,
max_hops: int,
max_facts: int,
) -> None:
self._llm = llm
self._store = store
self._top_k = top_k
self._max_hops = max_hops
self._max_facts = max_facts
def retrieve(self, question: str, mode: RetrievalMode = RetrievalMode.GRAPH) -> RetrievalContext:
embedding = self._llm.embed([question])[0]
chunks = self._store.vector_search(embedding, self._top_k)
if mode is RetrievalMode.VECTOR:
return RetrievalContext(chunks=chunks)
seeds = self._seed_entities(question, chunks)
facts = self._store.expand(seeds, self._max_hops, self._max_facts)
return RetrievalContext(chunks=chunks, facts=facts, seed_entities=seeds)
def _seed_entities(self, question: str, chunks: Sequence[RetrievedChunk]) -> list[str]:
lexical_matches = self._store.search_entities(question, MAX_LEXICAL_SEEDS)
mentions = Counter(key for chunk in chunks for key in chunk.entity_keys)
most_mentioned = [key for key, _ in mentions.most_common(MAX_MENTION_SEEDS)]
return list(dict.fromkeys([*lexical_matches, *most_mentioned]))
|