Fitin for infinity. (#4722)

### What problem does this PR solve?

### Type of change

- [x] Refactoring
This commit is contained in:
Kevin Hu 2025-02-05 16:47:05 +08:00 committed by GitHub
parent 307717b045
commit 283d036cba
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 7 additions and 3 deletions

View File

@ -65,9 +65,12 @@ class KGSearch(Dealer):
def _ent_info_from_(self, es_res, sim_thr=0.3): def _ent_info_from_(self, es_res, sim_thr=0.3):
res = {} res = {}
es_res = self.dataStore.getFields(es_res, ["content_with_weight", "_score", "entity_kwd", "rank_flt", flds = ["content_with_weight", "_score", "entity_kwd", "rank_flt", "n_hop_with_weight"]
"n_hop_with_weight"]) es_res = self.dataStore.getFields(es_res, flds)
for _, ent in es_res.items(): 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: if float(ent.get("_score", 0)) < sim_thr:
continue continue
if isinstance(ent["entity_kwd"], list): if isinstance(ent["entity_kwd"], list):

View File

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