Spaces:
Sleeping
Sleeping
File size: 1,285 Bytes
3a3b216 d02622b 3a3b216 d02622b 8c66169 3a3b216 37f0321 3a3b216 37f0321 |
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 |
from abc import ABC, abstractmethod
from collections import defaultdict
from typing import Dict, List, Tuple
from graphgen.bases.base_llm_wrapper import BaseLLMWrapper
from graphgen.bases.base_storage import BaseGraphStorage
from graphgen.bases.datatypes import Chunk
class BaseKGBuilder(ABC):
def __init__(self, llm_client: BaseLLMWrapper):
self.llm_client = llm_client
self._nodes: Dict[str, List[dict]] = defaultdict(list)
self._edges: Dict[Tuple[str, str], List[dict]] = defaultdict(list)
@abstractmethod
async def extract(
self, chunk: Chunk
) -> Tuple[Dict[str, List[dict]], Dict[Tuple[str, str], List[dict]]]:
"""Extract nodes and edges from a single chunk."""
raise NotImplementedError
@abstractmethod
async def merge_nodes(
self,
node_data: tuple[str, List[dict]],
kg_instance: BaseGraphStorage,
) -> None:
"""Merge extracted nodes into the knowledge graph."""
raise NotImplementedError
@abstractmethod
async def merge_edges(
self,
edges_data: tuple[Tuple[str, str], List[dict]],
kg_instance: BaseGraphStorage,
) -> None:
"""Merge extracted edges into the knowledge graph."""
raise NotImplementedError
|