# Adapt from https://github.com/HKUDS/LightRAG import asyncio import os import time from dataclasses import dataclass, field from typing import List, Union, cast import gradio as gr from tqdm.asyncio import tqdm as tqdm_async from .models import ( Chunk, JsonKVStorage, NetworkXStorage, OpenAIModel, Tokenizer, TraverseStrategy, WikiSearch, ) from .models.storage.base_storage import StorageNameSpace from .operators import ( extract_kg, judge_statement, quiz, search_wikipedia, skip_judge_statement, traverse_graph_atomically, traverse_graph_by_edge, traverse_graph_for_multi_hop, ) from .utils import compute_content_hash, create_event_loop, logger sys_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) @dataclass class GraphGen: unique_id: int = int(time.time()) working_dir: str = os.path.join(sys_path, "cache") # text chunking chunk_size: int = 1024 chunk_overlap_size: int = 100 # llm synthesizer_llm_client: OpenAIModel = None trainee_llm_client: OpenAIModel = None tokenizer_instance: Tokenizer = None # web search if_web_search: bool = False wiki_client: WikiSearch = field(default_factory=WikiSearch) # traverse strategy traverse_strategy: TraverseStrategy = field(default_factory=TraverseStrategy) # webui progress_bar: gr.Progress = None def __post_init__(self): self.full_docs_storage: JsonKVStorage = JsonKVStorage( self.working_dir, namespace="full_docs" ) self.text_chunks_storage: JsonKVStorage = JsonKVStorage( self.working_dir, namespace="text_chunks" ) self.wiki_storage: JsonKVStorage = JsonKVStorage( self.working_dir, namespace="wiki" ) self.graph_storage: NetworkXStorage = NetworkXStorage( self.working_dir, namespace="graph" ) self.rephrase_storage: JsonKVStorage = JsonKVStorage( self.working_dir, namespace="rephrase" ) self.qa_storage: JsonKVStorage = JsonKVStorage( os.path.join(self.working_dir, "data", "graphgen", str(self.unique_id)), namespace=f"qa-{self.unique_id}" ) async def async_split_chunks(self, data: Union[List[list], List[dict]], data_type: str) -> dict: # TODO： 是否进行指代消解 if len(data) == 0: return {} new_docs = {} inserting_chunks = {} if data_type == "raw": assert isinstance(data, list) and isinstance(data[0], dict) # compute hash for each document new_docs = { compute_content_hash(doc['content'], prefix="doc-"): {'content': doc['content']} for doc in data } _add_doc_keys = await self.full_docs_storage.filter_keys(list(new_docs.keys())) new_docs = {k: v for k, v in new_docs.items() if k in _add_doc_keys} if len(new_docs) == 0: logger.warning("All docs are already in the storage") return {} logger.info("[New Docs] inserting %d docs", len(new_docs)) cur_index = 1 doc_number = len(new_docs) async for doc_key, doc in tqdm_async( new_docs.items(), desc="[1/4]Chunking documents", unit="doc" ): chunks = { compute_content_hash(dp["content"], prefix="chunk-"): { **dp, 'full_doc_id': doc_key } for dp in self.tokenizer_instance.chunk_by_token_size(doc["content"], self.chunk_overlap_size, self.chunk_size) } inserting_chunks.update(chunks) if self.progress_bar is not None: self.progress_bar( cur_index / doc_number, f"Chunking {doc_key}" ) cur_index += 1 _add_chunk_keys = await self.text_chunks_storage.filter_keys(list(inserting_chunks.keys())) inserting_chunks = {k: v for k, v in inserting_chunks.items() if k in _add_chunk_keys} elif data_type == "chunked": assert isinstance(data, list) and isinstance(data[0], list) new_docs = { compute_content_hash("".join(chunk['content']), prefix="doc-"): {'content': "".join(chunk['content'])} for doc in data for chunk in doc } _add_doc_keys = await self.full_docs_storage.filter_keys(list(new_docs.keys())) new_docs = {k: v for k, v in new_docs.items() if k in _add_doc_keys} if len(new_docs) == 0: logger.warning("All docs are already in the storage") return {} logger.info("[New Docs] inserting %d docs", len(new_docs)) async for doc in tqdm_async(data, desc="[1/4]Chunking documents", unit="doc"): doc_str = "".join([chunk['content'] for chunk in doc]) for chunk in doc: chunk_key = compute_content_hash(chunk['content'], prefix="chunk-") inserting_chunks[chunk_key] = { **chunk, 'full_doc_id': compute_content_hash(doc_str, prefix="doc-") } _add_chunk_keys = await self.text_chunks_storage.filter_keys(list(inserting_chunks.keys())) inserting_chunks = {k: v for k, v in inserting_chunks.items() if k in _add_chunk_keys} await self.full_docs_storage.upsert(new_docs) await self.text_chunks_storage.upsert(inserting_chunks) return inserting_chunks def insert(self, data: Union[List[list], List[dict]], data_type: str): loop = create_event_loop() loop.run_until_complete(self.async_insert(data, data_type)) async def async_insert(self, data: Union[List[list], List[dict]], data_type: str): """ insert chunks into the graph """ inserting_chunks = await self.async_split_chunks(data, data_type) if len(inserting_chunks) == 0: logger.warning("All chunks are already in the storage") return logger.info("[New Chunks] inserting %d chunks", len(inserting_chunks)) logger.info("[Entity and Relation Extraction]...") _add_entities_and_relations = await extract_kg( llm_client=self.synthesizer_llm_client, kg_instance=self.graph_storage, tokenizer_instance=self.tokenizer_instance, chunks=[Chunk(id=k, content=v['content']) for k, v in inserting_chunks.items()], progress_bar = self.progress_bar, ) if not _add_entities_and_relations: logger.warning("No entities or relations extracted") return logger.info("[Wiki Search] is %s", 'enabled' if self.if_web_search else 'disabled') if self.if_web_search: logger.info("[Wiki Search]...") _add_wiki_data = await search_wikipedia( llm_client= self.synthesizer_llm_client, wiki_search_client=self.wiki_client, knowledge_graph_instance=_add_entities_and_relations ) await self.wiki_storage.upsert(_add_wiki_data) await self._insert_done() async def _insert_done(self): tasks = [] for storage_instance in [self.full_docs_storage, self.text_chunks_storage, self.graph_storage, self.wiki_storage]: if storage_instance is None: continue tasks.append(cast(StorageNameSpace, storage_instance).index_done_callback()) await asyncio.gather(*tasks) def quiz(self, max_samples=1): loop = create_event_loop() loop.run_until_complete(self.async_quiz(max_samples)) async def async_quiz(self, max_samples=1): await quiz(self.synthesizer_llm_client, self.graph_storage, self.rephrase_storage, max_samples) await self.rephrase_storage.index_done_callback() def judge(self, re_judge=False, skip=False): loop = create_event_loop() loop.run_until_complete(self.async_judge(re_judge, skip)) async def async_judge(self, re_judge=False, skip=False): if skip: _update_relations = await skip_judge_statement(self.graph_storage) else: _update_relations = await judge_statement(self.trainee_llm_client, self.graph_storage, self.rephrase_storage, re_judge) await _update_relations.index_done_callback() def traverse(self): loop = create_event_loop() loop.run_until_complete(self.async_traverse()) async def async_traverse(self): if self.traverse_strategy.qa_form == "atomic": results = await traverse_graph_atomically(self.synthesizer_llm_client, self.tokenizer_instance, self.graph_storage, self.traverse_strategy, self.text_chunks_storage, self.progress_bar) elif self.traverse_strategy.qa_form == "multi_hop": results = await traverse_graph_for_multi_hop(self.synthesizer_llm_client, self.tokenizer_instance, self.graph_storage, self.traverse_strategy, self.text_chunks_storage, self.progress_bar) elif self.traverse_strategy.qa_form == "aggregated": results = await traverse_graph_by_edge(self.synthesizer_llm_client, self.tokenizer_instance, self.graph_storage, self.traverse_strategy, self.text_chunks_storage, self.progress_bar) else: raise ValueError(f"Unknown qa_form: {self.traverse_strategy.qa_form}") await self.qa_storage.upsert(results) await self.qa_storage.index_done_callback() def clear(self): loop = create_event_loop() loop.run_until_complete(self.async_clear()) async def async_clear(self): await self.full_docs_storage.drop() await self.text_chunks_storage.drop() await self.wiki_storage.drop() await self.graph_storage.clear() await self.rephrase_storage.drop() await self.qa_storage.drop() logger.info("All caches are cleared")