lyg
2025-05-22 e60d75228fb161e464ca59fa2526bf0765f4d902
knowledgebase/doc/entity_recognition.py
@@ -11,6 +11,13 @@
import json
from knowledgebase import utils
from knowledgebase.db.doc_db_helper import doc_dbh
from knowledgebase.log import Log
llm = ChatOpenAI(temperature=0,
                 model="qwen2.5-72b-instruct",
                 base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
                 api_key="sk-15ecf7e273ad4b729c7f7f42b542749e")
class EntityRecognition:
@@ -19,28 +26,35 @@
    使用langchain构建实体抽取流程。
    """
    use_cache = False
    cache_file = "entity_recognition.cache"
    def __init__(self):
        llm = ChatOpenAI(temperature=0,
                         model="qwen2.5-72b-instruct",
                         base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
                         api_key="sk-15ecf7e273ad4b729c7f7f42b542749e")
        msg = HumanMessagePromptTemplate.from_template(template="""
    def __init__(self, doc_type: str):
        # 实体词列表
        entities = doc_dbh.get_entities_by_doc_type(doc_type)
        entity_list = ','.join([entity.name for entity in entities]) + "。"
        entity_rules = ";\n".join([f"- {entity.name}:{entity.prompts}" for entity in entities]) + "。"
        tpl = """
# 指令
请从给定的文本中提取实体词列表。
请根据实体词判断规则从给定的文本中判断是否有下列实体词相关内容,如果有则输出相关的实体词,没有则不输出,实体词列表定义如下:
""" + entity_list + """
## 实体词判断规则:
""" + entity_rules + """
# 约束
- 输出格式为JSON格式;
- 提取的实体词必须是:""" + entity_list + """;
- 如果没有复合上述规则的实体词则不要输出任何实体词;
- 输出数据结构为字符串数组。
# 示例
```json
["实体1","实体2"]
[\"""" + entities[0].name + """\"]
```
# 文本如下:
{text}
"""
                                                       )
        Log.info(tpl)
        msg = HumanMessagePromptTemplate.from_template(template=tpl)
        prompt = ChatPromptTemplate.from_messages([msg])
        parser = JsonOutputParser(pydantic_object=list[str])
        self.chain = prompt | llm | parser
@@ -60,15 +74,16 @@
        保存缓存。
        """
        text = json.dumps(self.cache)
        utils.save_to_file(text, self.cache_file)
        utils.save_text_to_file(text, self.cache_file)
    def run(self, in_text: str) -> list[str]:
        """
        运行实体识别抽取。
        :param in_text: str - 输入文本
        """
        # 缓存命中
        text_md5 = utils.generate_md5(in_text)
        if text_md5 in self.cache:
        text_md5 = utils.generate_text_md5(in_text)
        if self.use_cache and text_md5 in self.cache:
            return self.cache[text_md5]
        result = self.chain.invoke({"text": in_text})
        self.cache[text_md5] = result