Kevin Hu commited on
Commit
20174e9
·
1 Parent(s): e2220ea

add callback to entity extraction (#1843)

Browse files

### What problem does this PR solve?

### Type of change

- [x] Refactoring
- [ ]

Files changed (1) hide show
  1. graphrag/index.py +4 -2
graphrag/index.py CHANGED
@@ -86,13 +86,15 @@ def build_knowlege_graph_chunks(tenant_id: str, chunks: List[str], callback, ent
86
  for i in range(len(chunks)):
87
  tkn_cnt = num_tokens_from_string(chunks[i])
88
  if cnt+tkn_cnt >= left_token_count and texts:
89
- threads.append(exe.submit(ext, texts, {"entity_types": entity_types}, callback))
 
90
  texts = []
91
  cnt = 0
92
  texts.append(chunks[i])
93
  cnt += tkn_cnt
94
  if texts:
95
- threads.append(exe.submit(ext, texts))
 
96
 
97
  callback(0.5, "Extracting entities.")
98
  graphs = []
 
86
  for i in range(len(chunks)):
87
  tkn_cnt = num_tokens_from_string(chunks[i])
88
  if cnt+tkn_cnt >= left_token_count and texts:
89
+ for b in range(0, len(texts), 16):
90
+ threads.append(exe.submit(ext, ["\n".join(texts[b:b+16])], {"entity_types": entity_types}, callback))
91
  texts = []
92
  cnt = 0
93
  texts.append(chunks[i])
94
  cnt += tkn_cnt
95
  if texts:
96
+ for b in range(0, len(texts), 16):
97
+ threads.append(exe.submit(ext, ["\n".join(texts[b:b+16])], {"entity_types": entity_types}, callback))
98
 
99
  callback(0.5, "Extracting entities.")
100
  graphs = []