Skip to content

Commit

Permalink
Simplify the usage of dict (#2681)
Browse files Browse the repository at this point in the history
### What problem does this PR solve?
Simplify the usage of dictionaries

### Type of change

- [x] Refactoring

---------

Co-authored-by: Kevin Hu <[email protected]>
  • Loading branch information
yqkcn and KevinHuSh authored Sep 30, 2024
1 parent 9945988 commit ae5a877
Showing 1 changed file with 6 additions and 10 deletions.
16 changes: 6 additions & 10 deletions graphrag/entity_resolution.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#

import itertools
import logging
import re
import traceback
Expand Down Expand Up @@ -93,16 +93,12 @@ def __call__(self, graph: nx.Graph, prompt_variables: dict[str, Any] | None = No
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():
for k, v 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
for a, b in itertools.permutations(v, 2):
if self.is_similarity(a, b) and (b, a) not in candidate_resolution_tmp:
candidate_resolution_tmp.append((a, b))
candidate_resolution[k] = candidate_resolution_tmp

gen_conf = {"temperature": 0.5}
resolution_result = set()
Expand Down

0 comments on commit ae5a877

Please sign in to comment.