yqkcn Kevin Hu commited on
Commit
5720dea
·
1 Parent(s): 5f0746a

Simplify the usage of dict (#2681)

Browse files

### What problem does this PR solve?
Simplify the usage of dictionaries

### Type of change

- [x] Refactoring

---------

Co-authored-by: Kevin Hu <[email protected]>

Files changed (1) hide show
  1. graphrag/entity_resolution.py +6 -10
graphrag/entity_resolution.py CHANGED
@@ -13,7 +13,7 @@
13
  # See the License for the specific language governing permissions and
14
  # limitations under the License.
15
  #
16
-
17
  import logging
18
  import re
19
  import traceback
@@ -93,16 +93,12 @@ class EntityResolution:
93
  node_clusters[graph.nodes[node]['entity_type']].append(node)
94
 
95
  candidate_resolution = {entity_type: [] for entity_type in entity_types}
96
- for node_cluster in node_clusters.items():
97
  candidate_resolution_tmp = []
98
- for a in node_cluster[1]:
99
- for b in node_cluster[1]:
100
- if a == b:
101
- continue
102
- if self.is_similarity(a, b) and (b, a) not in candidate_resolution_tmp:
103
- candidate_resolution_tmp.append((a, b))
104
- if candidate_resolution_tmp:
105
- candidate_resolution[node_cluster[0]] = candidate_resolution_tmp
106
 
107
  gen_conf = {"temperature": 0.5}
108
  resolution_result = set()
 
13
  # See the License for the specific language governing permissions and
14
  # limitations under the License.
15
  #
16
+ import itertools
17
  import logging
18
  import re
19
  import traceback
 
93
  node_clusters[graph.nodes[node]['entity_type']].append(node)
94
 
95
  candidate_resolution = {entity_type: [] for entity_type in entity_types}
96
+ for k, v in node_clusters.items():
97
  candidate_resolution_tmp = []
98
+ for a, b in itertools.permutations(v, 2):
99
+ if self.is_similarity(a, b) and (b, a) not in candidate_resolution_tmp:
100
+ candidate_resolution_tmp.append((a, b))
101
+ candidate_resolution[k] = candidate_resolution_tmp
 
 
 
 
102
 
103
  gen_conf = {"temperature": 0.5}
104
  resolution_result = set()