BIDI_EXTRACTION_AI/BiddingKG/dl/predictors/role_context.py

# -*- coding: utf-8 -*-
"""角色规则上下文预计算（角色分类流程优化 Phase A）。

背景
====
角色分类链路中 ``prem`` / ``roleRule`` / ``roleRuleFinal`` /
``tendereeRuleRecall`` 四个 Predictor 各自遍历
``list_entitys × list_sentences`` 做实体-句子配对，并重复构建同一实体的
上下文切片：

    - ``prem.search_role_data``     — 字符切片（前 23 / 后 25），while 线性扫描
    - ``prem.search_money_data``    — 字符切片（前 13 / 后 15），while 线性扫描
    - ``RoleRulePredictor``         — ``get_context``（前后各 20 字）
    - ``TendereeRuleRecall``        — ``spanWindow``（前后各 40 词）

本模块把「配对 + 上下文构建」收敛为一次预计算：

    - 句子索引 ``{(doc_id, sentence_index): sentence}``，O(1) 定位实体所属句
    - ``RoleContext`` 惰性缓存三类上下文切片，按需计算、只算一次
    - 概率阶梯常量（命名原代码中的魔法概率，Phase C 起统一引用）

行为兼容性
==========
- ``build_role_contexts`` 保留原 ``search_role_data`` 的就地排序副作用
  （``list_entity.sort`` / ``list_sentence.sort``）。
- 与原 while 线性扫描的差异：当某个实体的 (doc_id, sentence_index) 在句子
  列表中不存在时，原实现会把扫描指针推到列表末尾导致其后所有实体配对失败；
  本实现按索引独立配对，不受前序缺口影响（属缺陷修正，已记录在优化日志）。

依赖方向
========
CORE（``predictors`` 内部），仅依赖 ``common``，不依赖 ``rules/``，
符合 ARCHITECTURE.md §4.3 依赖方向约束。

源码集成约束（Phase 7 §7.8）
============================
本模块不加载 ``dl/rules/`` 的 YAML；规则外置的 YAML 目前作为
「source of truth」用于校验与文档，predictors 消费 YAML 须待
Phase C pipeline 层规则注入机制就绪后进行。
"""
from __future__ import absolute_import

import re

from BiddingKG.dl.common.context_utils import spanWindow, get_context

__all__ = [
    "RoleContext",
    "build_role_contexts",
    "ROLE_LABELS",
    "P_MODEL_THRESHOLD",
    "P_RULE_OVERRIDE",
    "P_FALLBACK",
    "P_TITLE",
]


# ----------------------------------------------------------------------------
# 角色标签与概率阶梯
# ----------------------------------------------------------------------------

#: 角色标签 → 含义（与 Entity.label 语义一致）
ROLE_LABELS = {
    0: "tenderee",        # 招标人/采购人
    1: "agency",          # 代理机构
    2: "winTenderer",     # 中标人
    3: "secondTenderer",  # 第二候选人
    4: "thirdTenderer",   # 第三候选人
    5: "other",           # 其他/无角色
}

# ---- 概率阶梯（Phase B 命名魔法数）------------------------------------------
# 原代码中 0.501/0.51/0.55 等数值表示「略高于上游某置信度」的隐式约定，
# 此处显式命名。现有 prem.py / role.py 中的字面量保持不变，
# Phase C RoleRuleEngine 统一改用这些常量。

#: 模型置信阈值：模型预测 prob < 0.5 时置 label=5，交给规则召回
P_MODEL_THRESHOLD = 0.5

#: 规则覆盖模型：略高于模型置信（原代码 0.501）
P_RULE_OVERRIDE = 0.501

#: 兜底召回：略高于标题召回基线（原代码 0.51）
P_FALLBACK = 0.51

#: 标题召回基线（原代码 0.55）
P_TITLE = 0.55


#: 参与角色/金额链路上下文预计算的实体类型
_ROLE_CONTEXT_ENTITY_TYPES = ("org", "company", "money")


class RoleContext(object):
    """单个实体的预计算上下文。

    惰性计算三类切片并缓存；同一实体在多个 Predictor 间复用同一实例，
    每种切片整个链路只计算一次。

    属性
    ----
    entity / sentence   : 原实体与所属句子（一次配对，替代 O(E×S) 扫描）
    role_model_text     : prem 角色模型输入 (前文23字, 实体, 后文25字)
    money_model_text    : prem 金额模型输入 (前文13字, 实体, 后文15字)
    span                : RoleRulePredictor 用 get_context（前后各 20 字，
                          [0]=前文 [1]=实体 [2]=后文）
    span_wide           : TendereeRuleRecall 用 spanWindow（前后各 40 词，
                          word_flag=True，实体文本做全半角括号归一）
    span_rule           : RoleRulePredictor.rule_predict 输入用 spanWindow
                          （前后各 40 词，word_flag=True，use_text=False）
    """

    __slots__ = ("entity", "sentence", "_span", "_span_wide", "_span_rule")

    def __init__(self, entity, sentence):
        self.entity = entity
        self.sentence = sentence
        self._span = None
        self._span_wide = None
        self._span_rule = None

    # ---- 便捷透传 ----
    @property
    def entity_type(self):
        return self.entity.entity_type

    @property
    def sentence_text(self):
        return self.sentence.sentence_text

    # ---- prem 模型输入切片 ----
    def _char_slices(self, front_size, behind_size):
        text = self.sentence.sentence_text
        b = self.entity.wordOffset_begin
        e = self.entity.wordOffset_end
        return (text[max(0, b - front_size):b], text[b:e], text[e:e + behind_size])

    @property
    def role_model_text(self):
        """prem.search_role_data 的文本切片：(前文 23 字, 实体, 后文 25 字)。"""
        return self._char_slices(23, 25)

    @property
    def money_model_text(self):
        """prem.search_money_data 的文本切片：(前文 13 字, 实体, 后文 15 字)。"""
        return self._char_slices(13, 15)

    # ---- roleRule 上下文（get_context，前后各 20 字）----
    @property
    def span(self):
        if self._span is None:
            e = self.entity
            self._span = get_context(
                self.sentence.sentence_text,
                e.wordOffset_begin, e.wordOffset_end,
                size=20, center_include=True)
        return self._span

    # ---- tendereeRuleRecall 上下文（spanWindow，前后各 40 词）----
    @property
    def span_wide(self):
        if self._span_wide is None:
            e = self.entity
            self._span_wide = spanWindow(
                tokens=self.sentence.tokens,
                begin_index=e.begin_index, end_index=e.end_index,
                size=40, center_include=True,
                word_flag=True, use_text=True,
                text=re.sub("）", ")", re.sub("（", "(", e.entity_text)))
        return self._span_wide

    # ---- roleRule.rule_predict 输入（spanWindow，前后各 40 词，use_text=False）----
    @property
    def span_rule(self):
        if self._span_rule is None:
            e = self.entity
            self._span_rule = spanWindow(
                tokens=self.sentence.tokens,
                begin_index=e.begin_index, end_index=e.end_index,
                size=40, center_include=True,
                word_flag=True, use_text=False)
        return self._span_rule


def build_role_contexts(list_sentences, list_entitys):
    """一次预计算所有 (org/company/money) 实体的角色链路上下文。

    :param list_sentences: 与 ``list_entitys`` 同构嵌套的句子列表（按文章）
    :param list_entitys:   实体列表（按文章）
    :return: ``List[List[RoleContext]]``，与输入同构嵌套；
             仅含能配对到句子的 org/company/money 实体，
             内层按 sentence_index 升序（与原 prem 排序后遍历顺序一致）。

    副作用（与原 ``prem.search_role_data`` 保持一致）：
        就地对各 ``list_entity`` / ``list_sentence`` 按 sentence_index 排序。
    """
    result = []
    for list_sentence, list_entity in zip(list_sentences, list_entitys):
        # 保持与原实现一致的就地排序副作用（下游 RoleRuleFinalAdd 等依赖此顺序）
        list_entity.sort(key=lambda x: x.sentence_index)
        list_sentence.sort(key=lambda x: x.sentence_index)

        # 句子索引：O(1) 定位，等价于原 while 扫描的首次命中
        index = {}
        for s in list_sentence:
            key = (s.doc_id, s.sentence_index)
            if key not in index:
                index[key] = s

        ctxs = []
        for entity in list_entity:
            if entity.entity_type not in _ROLE_CONTEXT_ENTITY_TYPES:
                continue
            sentence = index.get((entity.doc_id, entity.sentence_index))
            if sentence is None:
                continue
            ctxs.append(RoleContext(entity, sentence))
        result.append(ctxs)
    return result