| | |
| | | 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) |
| | |
| | | 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) |
| | | self.entities.append(_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() |