feat(operate): 添加实体类型配置并优化提示生成

- 在全局配置中添加 entity_types 参数,用于自定义实体类型
- 在生成实体提取和关系提取的提示时,使用配置的实体类型替代默认值
- 优化了提示生成逻辑,提高了代码的可配置性和灵活性
This commit is contained in:
Magic_yuan
2024-12-11 13:53:05 +08:00
parent a09e1ba190
commit b89041b5b3

View File

@@ -260,6 +260,9 @@ async def extract_entities(
language = global_config["addon_params"].get(
"language", PROMPTS["DEFAULT_LANGUAGE"]
)
entity_types = global_config["addon_params"].get(
"entity_types", PROMPTS["DEFAULT_ENTITY_TYPES"]
)
example_number = global_config["addon_params"].get("example_number", None)
if example_number and example_number < len(PROMPTS["entity_extraction_examples"]):
examples = "\n".join(
@@ -272,7 +275,7 @@ async def extract_entities(
tuple_delimiter=PROMPTS["DEFAULT_TUPLE_DELIMITER"],
record_delimiter=PROMPTS["DEFAULT_RECORD_DELIMITER"],
completion_delimiter=PROMPTS["DEFAULT_COMPLETION_DELIMITER"],
entity_types=",".join(PROMPTS["DEFAULT_ENTITY_TYPES"]),
entity_types=",".join(entity_types),
language=language,
)
# add example's format
@@ -283,7 +286,7 @@ async def extract_entities(
tuple_delimiter=PROMPTS["DEFAULT_TUPLE_DELIMITER"],
record_delimiter=PROMPTS["DEFAULT_RECORD_DELIMITER"],
completion_delimiter=PROMPTS["DEFAULT_COMPLETION_DELIMITER"],
entity_types=",".join(PROMPTS["DEFAULT_ENTITY_TYPES"]),
entity_types=",".join(entity_types),
examples=examples,
language=language,
)