# -*- 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
|
|
def __init__(self):
|
Log.info("初始化EntityHelper")
|
current_dir = os.path.dirname(__file__)
|
self.doc_prompt_map = {}
|
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']:
|
_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
|
|
def get_entity_helper():
|
global _entity_helper
|
return _entity_helper
|
|
def init_entity_helper():
|
global _entity_helper
|
_entity_helper = EntityHelper()
|
