lyg
2 天以前 22f370322412074174cde20ecfd14ec03657ab63
knowledgebase/doc/entity_helper.py
@@ -15,15 +15,11 @@
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)
@@ -33,12 +29,19 @@
                    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])
                        if doc_dbh.get_entity(_entity):
                            continue
                        doc_dbh.add_entity(_entity)
                        Log.info(f"新增Entity:{entity},id:{_entity.id}")
_entity_helper:EntityHelper = None
entity_helper = EntityHelper()
def get_entity_helper():
    global _entity_helper
    return _entity_helper
def init_entity_helper():
    global _entity_helper
    _entity_helper = EntityHelper()