# -*- coding: utf-8 -*-
|
#
|
# @author:
|
# @date:
|
# @version:
|
# @description:
|
import json
|
import os.path
|
|
from knowledgebase.db.doc_db_helper import doc_dbh
|
from knowledgebase.db.doc_db_models import TEntity
|
from knowledgebase.log import Log
|
|
|
class EntityHelper:
|
# 文档类型和识别提示词map
|
doc_prompt_map: dict
|
# 所有实体
|
entities: list[TEntity]
|
|
def __init__(self):
|
Log.info("初始化EntityHelper")
|
current_dir = os.path.dirname(__file__)
|
self.entities = doc_dbh.get_all_entities()
|
self.doc_prompt_map = {}
|
entity_names = [entity.name for entity in self.entities]
|
with open(f'{current_dir}/../../tpl/entities.json', 'r', encoding='utf-8') as f:
|
text = f.read()
|
obj = json.loads(text)
|
for ty in obj:
|
obj2 = obj[ty]
|
for doc_ty in obj2:
|
prompts = obj2[doc_ty]['prompts']
|
self.doc_prompt_map[doc_ty] = prompts
|
for entity in obj2[doc_ty]['entities']:
|
if entity in entity_names:
|
continue
|
_entity = TEntity(name=entity, type=ty, doc_type=doc_ty,
|
prompts=obj2[doc_ty]['entities'][entity])
|
doc_dbh.add_entity(_entity)
|
self.entities.append(_entity)
|
Log.info(f"新增Entity：{entity}，id：{_entity.id}")
|
|
|
entity_helper = EntityHelper()
|
