瀏覽代碼

Fitin for infinity. (#4722)

### What problem does this PR solve?

### Type of change

- [x] Refactoring
tags/v0.16.0
Kevin Hu 8 月之前
父節點
當前提交
283d036cba
沒有連結到貢獻者的電子郵件帳戶。
共有 2 個檔案被更改,包括 7 行新增3 行删除
  1. 5
    2
      graphrag/search.py
  2. 2
    1
      rag/utils/infinity_conn.py

+ 5
- 2
graphrag/search.py 查看文件

@@ -65,9 +65,12 @@ class KGSearch(Dealer):

def _ent_info_from_(self, es_res, sim_thr=0.3):
res = {}
es_res = self.dataStore.getFields(es_res, ["content_with_weight", "_score", "entity_kwd", "rank_flt",
"n_hop_with_weight"])
flds = ["content_with_weight", "_score", "entity_kwd", "rank_flt", "n_hop_with_weight"]
es_res = self.dataStore.getFields(es_res, flds)
for _, ent in es_res.items():
for f in flds:
if f in ent and ent[f] is None:
del ent[f]
if float(ent.get("_score", 0)) < sim_thr:
continue
if isinstance(ent["entity_kwd"], list):

+ 2
- 1
rag/utils/infinity_conn.py 查看文件

@@ -311,6 +311,7 @@ class InfinityConnection(DocStoreConnection):
if matchExprs:
selectFields.append(score_func)
selectFields.append(PAGERANK_FLD)
selectFields = [f for f in selectFields if f != "_score"]

# Prepare expressions common to all tables
filter_cond = None
@@ -405,7 +406,7 @@ class InfinityConnection(DocStoreConnection):
if matchExprs:
res = res.sort(pl.col(score_column) + pl.col(PAGERANK_FLD), descending=True, maintain_order=True)
if score_column and score_column != "SCORE":
res = res.rename({score_column: "SCORE"})
res = res.rename({score_column: "_score"})
res = res.limit(limit)
logger.debug(f"INFINITY search final result: {str(res)}")
return res, total_hits_count

Loading…
取消
儲存