BIDI_EXTRACTION_AI/BiddingKG/dl/pipeline/pipeline.py

# -*- coding: utf-8 -*-
"""Pipeline 主编排器。

按 ARCHITECTURE.md Phase 8 设计，统一同步模式与队列模式的调用入口。
``api/single_server.py`` 和 ``api/queue_worker.py`` 都调用 ``Pipeline.run(ctx)``，
确保两条路径走完全相同的业务逻辑。

编排步骤（原 extract.py / run_single_server.py / run_model_server.py 的共有流程）::

    1. preprocess  — Preprocessing.get_preprocessed
    2. codename    — CodeNamePredict
    3. prem        — PREMPredict
    4. rule        — RoleRulePredictor
    5. epc         — EPCPredict
    6. entity_link — entityLink.link_entitys
    7. assemble    — getAttributes.getPREMs
    8. union       — Preprocessing.union_result

使用方式::

    from BiddingKG.dl.pipeline.pipeline import Pipeline
    from BiddingKG.dl.pipeline.context import PipelineContext

    ctx = PipelineContext(doc_id="123", html=content, title=title)
    pipeline = Pipeline()
    result = pipeline.run(ctx)
    # result == ctx.result

约束（ARCHITECTURE.md §4.1）:
    pipeline 只负责编排，不写复杂抽取逻辑。
"""
from __future__ import absolute_import

import time
import logging

from BiddingKG.dl.pipeline.context import PipelineContext

logger = logging.getLogger(__name__)


class Pipeline(object):
    """Pipeline 主编排器。

    同步模式（single_server）和队列模式（queue_worker）共用此类。
    模型实例在 ``__init__`` 时惰性创建（首次调用 ``run()`` 时才加载），
    避免导入时的副作用。
    """

    def __init__(self):
        self._predictors = None
        self._preprocessing = None
        self._getAttributes = None
        self._entityLink = None

    # ------------------------------------------------------------------------
    # 惰性加载业务模块
    # ------------------------------------------------------------------------
    def _ensure_loaded(self):
        """惰性加载模型和业务模块（仅首次调用时）。"""
        if self._predictors is not None:
            return

        logger.info("Pipeline: loading models and modules...")

        import BiddingKG.dl.interface.predictor as predictor
        import BiddingKG.dl.interface.Preprocessing as Preprocessing
        import BiddingKG.dl.interface.getAttributes as getAttributes
        import BiddingKG.dl.entityLink.entityLink as entityLink

        self._predictors = {
            "codeName": predictor.CodeNamePredict(),
            "prem": predictor.PREMPredict(),
            "epc": predictor.EPCPredict(),
            "roleRule": predictor.RoleRulePredictor(),
        }
        self._preprocessing = Preprocessing
        self._getAttributes = getAttributes
        self._entityLink = entityLink

        logger.info("Pipeline: models loaded")

    # ------------------------------------------------------------------------
    # 主入口
    # ------------------------------------------------------------------------
    def run(self, ctx):
        """执行完整 Pipeline。

        :param ctx: PipelineContext，输入字段（doc_id/html/title）必须已填充
        :return: dict — 最终结果（同时写入 ctx.result）
        """
        self._ensure_loaded()
        t_total = time.time()

        # ---- 1. 预处理 ----
        t0 = time.time()
        k = ctx.doc_id or "pipeline"
        content = ctx.html
        ContentIDs = [[k, content, time.time(), ctx.doc_id, ctx.title]]
        list_articles, list_sentences, list_entitys, _cost_time = \
            self._preprocessing.get_preprocessed(ContentIDs, useselffool=True)
        ctx.articles = list_articles
        ctx.sentences = list_sentences
        ctx.entities = list_entitys
        ctx.cost_time["preprocess"] = time.time() - t0
        ctx.cost_time.update(_cost_time)

        # ---- 2-7. 预测 + 装配 ----
        self._run_predict(ctx, list_articles, list_sentences, list_entitys)

        # ---- 8. 合并结果 ----
        t0 = time.time()
        codeName = ctx.model_results.get("codeName")
        prem = ctx.model_results.get("prem")
        data = self._preprocessing.union_result(codeName, prem)[0][1]
        data["cost_time"] = ctx.cost_time
        data["success"] = True
        ctx.result = data
        ctx.cost_time["union"] = time.time() - t0

        ctx.cost_time["total"] = time.time() - t_total
        logger.info("Pipeline done for doc_id=%s, total=%.2fs" % (
            ctx.doc_id, ctx.cost_time["total"]))

        return ctx.result

    def run_predict(self, ctx, list_articles, list_sentences, list_entitys):
        """仅执行预测+装配（跳过预处理）。

        用于队列模式中预处理已由其他 worker 完成的场景。
        """
        self._ensure_loaded()
        ctx.articles = list_articles
        ctx.sentences = list_sentences
        ctx.entities = list_entitys
        self._run_predict(ctx, list_articles, list_sentences, list_entitys)

        # 合并结果
        codeName = ctx.model_results.get("codeName")
        prem = ctx.model_results.get("prem")
        data = self._preprocessing.union_result(codeName, prem)[0][1]
        data["cost_time"] = ctx.cost_time
        data["success"] = True
        ctx.result = data
        return ctx.result

    # ------------------------------------------------------------------------
    # 内部：预测步骤
    # ------------------------------------------------------------------------
    def _run_predict(self, ctx, list_articles, list_sentences, list_entitys):
        """执行步骤 2-7：codename → prem → rule → epc → entityLink → assemble。"""
        p = self._predictors

        # 2. 项目编号/名称
        t0 = time.time()
        codeName = p["codeName"].predict(list_articles)
        ctx.model_results["codeName"] = codeName
        ctx.cost_time["codename"] = time.time() - t0

        # 3. 角色/金额候选
        t0 = time.time()
        p["prem"].predict(list_sentences, list_entitys)
        ctx.cost_time["prem"] = time.time() - t0

        # 4. 规则修正
        t0 = time.time()
        p["roleRule"].predict(list_articles, list_sentences, list_entitys, codeName)
        ctx.cost_time["rule"] = time.time() - t0

        # 5. 人员/EPC
        t0 = time.time()
        p["epc"].predict(list_sentences, list_entitys)
        ctx.cost_time["person"] = time.time() - t0

        # 6. 实体链接
        t0 = time.time()
        self._entityLink.link_entitys(list_entitys)
        ctx.cost_time["entity_link"] = time.time() - t0

        # 7. 业务装配
        t0 = time.time()
        prem = self._getAttributes.getPREMs(list_sentences, list_entitys, list_articles)
        ctx.model_results["prem"] = prem
        ctx.cost_time["attrs"] = time.time() - t0