diff --git a/example/sdk/dataset_example.py b/example/sdk/dataset_example.py index 252ed02c0..3a0504d8d 100644 --- a/example/sdk/dataset_example.py +++ b/example/sdk/dataset_example.py @@ -36,10 +36,8 @@ try: updated_dataset = dataset_instance.update(updated_message) # get the dataset (list datasets) - dataset_list = ragflow_instance.list_datasets(id=dataset_instance.id) - dataset_instance_2 = dataset_list[0] print(dataset_instance) - print(dataset_instance_2) + print(updated_dataset) # delete the dataset (delete datasets) to_be_deleted_datasets = [dataset_instance.id] diff --git a/graphrag/general/entity_embedding.py b/graphrag/general/entity_embedding.py index 746e2442f..78c7fc31d 100644 --- a/graphrag/general/entity_embedding.py +++ b/graphrag/general/entity_embedding.py @@ -21,7 +21,7 @@ class NodeEmbeddings: embeddings: np.ndarray -def embed_nod2vec( +def embed_node2vec( graph: nx.Graph | nx.DiGraph, dimensions: int = 1536, num_walks: int = 10, @@ -44,13 +44,13 @@ def embed_nod2vec( return NodeEmbeddings(embeddings=lcc_tensors[0], nodes=lcc_tensors[1]) -def run(graph: nx.Graph, args: dict[str, Any]) -> NodeEmbeddings: +def run(graph: nx.Graph, args: dict[str, Any]) -> dict: """Run method definition.""" if args.get("use_lcc", True): graph = stable_largest_connected_component(graph) # create graph embedding using node2vec - embeddings = embed_nod2vec( + embeddings = embed_node2vec( graph=graph, dimensions=args.get("dimensions", 1536), num_walks=args.get("num_walks", 10), diff --git a/graphrag/search.py b/graphrag/search.py index d39f482d0..ebc8f4a88 100644 --- a/graphrag/search.py +++ b/graphrag/search.py @@ -23,7 +23,7 @@ import trio from api.utils import get_uuid from graphrag.query_analyze_prompt import PROMPTS -from graphrag.utils import get_entity_type2sampels, get_llm_cache, set_llm_cache, get_relation +from graphrag.utils import get_entity_type2samples, get_llm_cache, set_llm_cache, get_relation from rag.utils import num_tokens_from_string, get_float from rag.utils.doc_store_conn import OrderByExpr @@ -42,7 +42,7 @@ class KGSearch(Dealer): return response def query_rewrite(self, llm, question, idxnms, kb_ids): - ty2ents = trio.run(lambda: get_entity_type2sampels(idxnms, kb_ids)) + ty2ents = trio.run(lambda: get_entity_type2samples(idxnms, kb_ids)) hint_prompt = PROMPTS["minirag_query2kwd"].format(query=question, TYPE_POOL=json.dumps(ty2ents, ensure_ascii=False, indent=2)) result = self._chat(llm, hint_prompt, [{"role": "user", "content": "Output:"}], {}) diff --git a/graphrag/utils.py b/graphrag/utils.py index 6b80d7fe8..6abe5f9a9 100644 --- a/graphrag/utils.py +++ b/graphrag/utils.py @@ -561,7 +561,7 @@ def merge_tuples(list1, list2): return result -async def get_entity_type2sampels(idxnms, kb_ids: list): +async def get_entity_type2samples(idxnms, kb_ids: list): es_res = await trio.to_thread.run_sync(lambda: settings.retrievaler.search({"knowledge_graph_kwd": "ty2ents", "kb_id": kb_ids, "size": 10000, "fields": ["content_with_weight"]}, idxnms, kb_ids)) res = defaultdict(list)