BIDI_EXTRACTION_AI/BiddingKG/dl/template_extract/templates/shared_template/template1.py

# coding:utf8

import re
from typing import Dict, Any
from BiddingKG.dl.template_extract.abstract_template import AbstractTemplate, table2list
import os
import json


class TransactionTemplate(AbstractTemplate):
    """
    成交信息提取模板（适配多列表格场景）
    适配包含成交供应商、产品明细的表格提取
    """

    def __init__(self):
        super().__init__(
            template_id=os.path.abspath(__file__),
            priority=3
        )
        self.header1 = None  # 成交供应商表格表头
        self.data1 = None    # 成交供应商表格数据

    def check_call_timing(self, preprocessed_data: Dict[str, Any]) -> bool:
        """检查是否符合模板调用条件"""
        if "表格列表" not in preprocessed_data:
            return False

        tables = preprocessed_data["表格列表"]
        if len(tables) < 1:
            return False

        try:
            # 验证第一张表格（成交供应商表格）
            supplier_table = tables[0]
            if len(supplier_table) < 2:  # 至少包含表头+1行数据
                return False
            # 成交供应商表头正则匹配
            supplier_header_patterns = [
                '序号',
                '成交供应商名称|中标供应商名称|成交单位名称|中标单位名称',
                '成交供应商地址|中标供应商地址|成交单位地址|中标单位地址',
                '成交金额（元）|中标金额（元）|成交价（元）|中标价（元）'
            ]
            for i, pattern in enumerate(supplier_header_patterns):
                if not re.fullmatch(pattern, supplier_table[0][i]):
                    # print(f'成交供应商表头未验证: {pattern} vs {supplier_table[0][i]}')
                    return False

            # 赋值表格数据
            self.header1 = supplier_table[0]
            self.data1 = supplier_table[1:]
            return True
        except Exception as e:
            # print(f"表格验证错误: {e}")
            return False

    def extract(self, preprocessed_data: Dict[str, Any]) -> Dict[str, Any]:
        """提取表格中的中标信息、产品信息
        """
        extract_result = {
            "项目编号": "",
            "项目名称": "",
            "招标信息": [],
            "中标信息": [],
            "候选人信息": [],
            "产品信息": []
        }

        # 提取中标信息（成交供应商表格）
        for row in self.data1:
            if len(row) != 4:
                continue
            seq, supplier_name, supplier_addr, deal_amount = row

            # 补充金额单位
            if re.search('[万亿美欧日]?元', self.header1[3]) and not re.search('[万亿美欧日]?元', deal_amount):
                deal_amount += re.search('[万亿美欧日]?元', self.header1[3]).group(0)

            # 没有的字段不用返回
            extract_result["中标信息"].append({
                "中标人": supplier_name,
                "中标价": deal_amount,
                "地址": supplier_addr,
            })

        # 招标信息校验：无招标人/预算则置空
        extract_result["招标信息"] = []

        return extract_result


# 使用示例
if __name__ == "__main__":
    # 读取并解析JSON文件
    with open("data.json", "r", encoding="utf8") as f:
        preprocessed_data = json.load(f)

    # 初始化模板并提取信息
    template = TransactionTemplate()
    if template.check_call_timing(preprocessed_data):
        result = template.extract(preprocessed_data)
        print("提取结果：", json.dumps(result, ensure_ascii=False, indent=4))
    else:
        print("当前数据不满足模板调用条件")