graphrag/entity_resolution.py

#
# Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import logging
import re
import traceback
from dataclasses import dataclass
from typing import Any

import networkx as nx
from rag.nlp import is_english
import editdistance
from graphrag.entity_resolution_prompt import ENTITY_RESOLUTION_PROMPT
from rag.llm.chat_model import Base as CompletionLLM
from graphrag.utils import ErrorHandlerFn, perform_variable_replacements

DEFAULT_RECORD_DELIMITER = "##"
DEFAULT_ENTITY_INDEX_DELIMITER = "<|>"
DEFAULT_RESOLUTION_RESULT_DELIMITER = "&&"


@dataclass
class EntityResolutionResult:
 """Entity resolution result class definition."""

 output: nx.Graph


class EntityResolution:
 """Entity resolution class definition."""

 _llm: CompletionLLM
 _resolution_prompt: str
 _output_formatter_prompt: str
 _on_error: ErrorHandlerFn
 _record_delimiter_key: str
 _entity_index_delimiter_key: str
 _resolution_result_delimiter_key: str

 def __init__(
 self,
 llm_invoker: CompletionLLM,
 resolution_prompt: str | None = None,
 on_error: ErrorHandlerFn | None = None,
 record_delimiter_key: str | None = None,
 entity_index_delimiter_key: str | None = None,
 resolution_result_delimiter_key: str | None = None,
 input_text_key: str | None = None
 ):
 """Init method definition."""
 self._llm = llm_invoker
 self._resolution_prompt = resolution_prompt or ENTITY_RESOLUTION_PROMPT
 self._on_error = on_error or (lambda _e, _s, _d: None)
 self._record_delimiter_key = record_delimiter_key or "record_delimiter"
 self._entity_index_dilimiter_key = entity_index_delimiter_key or "entity_index_delimiter"
 self._resolution_result_delimiter_key = resolution_result_delimiter_key or "resolution_result_delimiter"
 self._input_text_key = input_text_key or "input_text"

 def __call__(self, graph: nx.Graph, prompt_variables: dict[str, Any] | None = None) -> EntityResolutionResult:
 """Call method definition."""
 if prompt_variables is None:
 prompt_variables = {}

 # Wire defaults into the prompt variables
 prompt_variables = {
 **prompt_variables,
 self._record_delimiter_key: prompt_variables.get(self._record_delimiter_key)
 or DEFAULT_RECORD_DELIMITER,
 self._entity_index_dilimiter_key: prompt_variables.get(self._entity_index_dilimiter_key)
 or DEFAULT_ENTITY_INDEX_DELIMITER,
 self._resolution_result_delimiter_key: prompt_variables.get(self._resolution_result_delimiter_key)
 or DEFAULT_RESOLUTION_RESULT_DELIMITER,
 }

 nodes = graph.nodes
 entity_types = list(set(graph.nodes[node]['entity_type'] for node in nodes))
 node_clusters = {entity_type: [] for entity_type in entity_types}

 for node in nodes:
 node_clusters[graph.nodes[node]['entity_type']].append(node)

 candidate_resolution = {entity_type: [] for entity_type in entity_types}
 for node_cluster in node_clusters.items():
 candidate_resolution_tmp = []
 for a in node_cluster[1]:
 for b in node_cluster[1]:
 if a == b:
 continue
 if self.is_similarity(a, b) and (b, a) not in candidate_resolution_tmp:
 candidate_resolution_tmp.append((a, b))
 if candidate_resolution_tmp:
 candidate_resolution[node_cluster[0]] = candidate_resolution_tmp

 gen_conf = {"temperature": 0.5}
 resolution_result = set()
 for candidate_resolution_i in candidate_resolution.items():
 if candidate_resolution_i[1]:
 try:
 pair_txt = [
 f'When determining whether two {candidate_resolution_i[0]}s are the same, you should only focus on critical properties and overlook noisy factors.\n']
 for index, candidate in enumerate(candidate_resolution_i[1]):
 pair_txt.append(
 f'Question {index + 1}: name of{candidate_resolution_i[0]} A is {candidate[0]} ,name of{candidate_resolution_i[0]} B is {candidate[1]}')
 sent = 'question above' if len(pair_txt) == 1 else f'above {len(pair_txt)} questions'
 pair_txt.append(
 f'\nUse domain knowledge of {candidate_resolution_i[0]}s to help understand the text and answer the {sent} in the format: For Question i, Yes, {candidate_resolution_i[0]} A and {candidate_resolution_i[0]} B are the same {candidate_resolution_i[0]}./No, {candidate_resolution_i[0]} A and {candidate_resolution_i[0]} B are different {candidate_resolution_i[0]}s. For Question i+1, (repeat the above procedures)')
 pair_prompt = '\n'.join(pair_txt)

 variables = {
 **prompt_variables,
 self._input_text_key: pair_prompt
 }
 text = perform_variable_replacements(self._resolution_prompt, variables=variables)

 response = self._llm.chat(text, [], gen_conf)
 result = self._process_results(len(candidate_resolution_i[1]), response,
 prompt_variables.get(self._record_delimiter_key,
 DEFAULT_RECORD_DELIMITER),
 prompt_variables.get(self._entity_index_dilimiter_key,
 DEFAULT_ENTITY_INDEX_DELIMITER),
 prompt_variables.get(self._resolution_result_delimiter_key,
 DEFAULT_RESOLUTION_RESULT_DELIMITER))
 for result_i in result:
 resolution_result.add(candidate_resolution_i[1][result_i[0] - 1])
 except Exception as e:
 logging.exception("error entity resolution")
 self._on_error(e, traceback.format_exc(), None)

 connect_graph = nx.Graph()
 connect_graph.add_edges_from(resolution_result)
 for sub_connect_graph in nx.connected_components(connect_graph):
 sub_connect_graph = connect_graph.subgraph(sub_connect_graph)
 remove_nodes = list(sub_connect_graph.nodes)
 keep_node = remove_nodes.pop()
 for remove_node in remove_nodes:
 remove_node_neighbors = graph[remove_node]
 graph.nodes[keep_node]['description'] += graph.nodes[remove_node]['description']
 graph.nodes[keep_node]['weight'] += graph.nodes[remove_node]['weight']
 remove_node_neighbors = list(remove_node_neighbors)
 for remove_node_neighbor in remove_node_neighbors:
 if remove_node_neighbor == keep_node:
 graph.remove_edge(keep_node, remove_node)
 continue
 if graph.has_edge(keep_node, remove_node_neighbor):
 graph[keep_node][remove_node_neighbor]['weight'] += graph[remove_node][remove_node_neighbor][
 'weight']
 graph[keep_node][remove_node_neighbor]['description'] += \
 graph[remove_node][remove_node_neighbor]['description']
 graph.remove_edge(remove_node, remove_node_neighbor)
 else:
 graph.add_edge(keep_node, remove_node_neighbor,
 weight=graph[remove_node][remove_node_neighbor]['weight'],
 description=graph[remove_node][remove_node_neighbor]['description'],
 source_id="")
 graph.remove_edge(remove_node, remove_node_neighbor)
 graph.remove_node(remove_node)

 for node_degree in graph.degree:
 graph.nodes[str(node_degree[0])]["rank"] = int(node_degree[1])

 return EntityResolutionResult(
 output=graph,
 )

 def _process_results(
 self,
 records_length: int,
 results: str,
 record_delimiter: str,
 entity_index_delimiter: str,
 resolution_result_delimiter: str
 ) -> list:
 ans_list = []
 records = [r.strip() for r in results.split(record_delimiter)]
 for record in records:
 pattern_int = f"{re.escape(entity_index_delimiter)}(\d+){re.escape(entity_index_delimiter)}"
 match_int = re.search(pattern_int, record)
 res_int = int(str(match_int.group(1) if match_int else '0'))
 if res_int > records_length:
 continue

 pattern_bool = f"{re.escape(resolution_result_delimiter)}([a-zA-Z]+){re.escape(resolution_result_delimiter)}"
 match_bool = re.search(pattern_bool, record)
 res_bool = str(match_bool.group(1) if match_bool else '')

 if res_int and res_bool:
 if res_bool.lower() == 'yes':
 ans_list.append((res_int, "yes"))

 return ans_list

 def is_similarity(self, a, b):
 if is_english(a) and is_english(b):
 if editdistance.eval(a, b) <= min(len(a), len(b)) // 2:
 return True

 if len(set(a) & set(b)) > 0:
 return True

 return False