BIDI_EXTRACTION_AI/BiddingKG/dl/template_extract/templates/XX0182/template8.py

# coding:utf8

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


class ProductInfoTemplate(AbstractTemplate):
    """
    产品信息提取模板（适配data.json中的第二个表格结构）
    从第二个表格提取产品相关信息
    """

    def __init__(self):
        super().__init__(
            template_id=os.path.abspath(__file__),
            priority=3
        )
        self.target_table = 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:
            print("表格数量不足，至少需要2个表格")
            return False

        try:
            self.target_table = tables[0]  # 只处理第二个表格

            # 验证表格关键表头（按data.json第二个表格实际结构泛化匹配）
            header_patterns = [
                (0, 0, r'序号|编号'),  # 第1行第1列
                (0, 1, r'品目|采购品目'),  # 第1行第2列
                (0, 2, r'商品名称|产品名称|名称'),  # 第1行第3列
                (0, 3, r'品牌|品牌名称'),  # 第1行第4列
                (0, 4, r'型号|规格型号|规格'),  # 第1行第5列
                (0, 5, r'产地及生产商|产地|生产商'),  # 第1行第6列（用于结构验证）
                (0, 6, r'数量|采购数量'),  # 第1行第7列
                (0, 7, r'计量单位|单位|数量单位'),  # 第1行第8列
                (0, 8, r'预算单价|单价|单价金额'),  # 第1行第9列
            ]

            # 验证表头行存在
            if len(self.target_table) < 1:
                return False
            header_row = self.target_table[0]

            # 匹配所有关键表头
            for row, col, pattern in header_patterns:
                if col >= len(header_row):
                    print(f"表头列索引超出范围: 第{row+1}行第{col+1}列")
                    return False
                header_text = header_row[col].strip()
                if not re.search(pattern, header_text, re.IGNORECASE):
                    print(f"表头不匹配: {header_text} 不符合 {pattern}")
                    return False

            return True

        except Exception as e:
            print(f"表格验证错误: {e}")
            return False

    def extract(self, preprocessed_data: Dict[str, Any]) -> Dict[str, Any]:
        """提取表格中的产品信息"""
        result = {"产品信息": []}

        # 提取产品信息（只处理数据行，跳过表头行）
        for row_idx in range(1, len(self.target_table)):
            # 过滤技术参数行（根据data.json内容判断）
            if self._get_value(self.target_table, row_idx, 1) == "技术参数":
                continue

            product_info = self._extract_product_info(row_idx)
            if product_info:
                result["产品信息"].append(product_info)

        return result

    def _get_value(self, table: list, row: int, col: int) -> str:
        """获取表格指定位置的值"""
        if row < len(table) and col < len(table[row]):
            return table[row][col].strip()
        return ""

    def _format_amount(self, value: str, header_text: str) -> str:
        """格式化金额，补充单位"""
        if not value:
            return ""
        # 检查值中是否已包含单位
        if re.search(r'[万亿]?元|￥|人民币', value):
            return value
        # 从表头获取单位补充
        if re.search(r'[万亿]?元|￥|人民币', header_text):
            unit = re.search(r'[万亿]?元|￥|人民币', header_text).group(0)
            return f"{value}{unit}"
        return value

    def _extract_product_info(self, data_row: int) -> Dict[str, Any]:
        """提取指定数据行的产品详细信息"""
        product = {}
        header_row = 0  # 表头行固定为第1行（索引0）

        # 提取产品名称（商品名称）
        product_name = self._get_value(self.target_table, data_row, 2)
        if product_name:
            product["产品"] = product_name

        # 提取品目
        paimu = self._get_value(self.target_table, data_row, 1)
        if paimu:
            product["品目"] = paimu

        # 提取品牌
        brand = self._get_value(self.target_table, data_row, 3)
        if brand and brand != "见附件":  # 过滤无效值
            product["品牌"] = brand

        # 提取规格（型号）
        spec = self._get_value(self.target_table, data_row, 4)
        if spec and spec != "见附件":  # 过滤无效值
            product["规格"] = spec

        # 提取数量
        quantity = self._get_value(self.target_table, data_row, 6)
        if quantity and quantity != "见附件":  # 过滤无效值
            product["数量"] = quantity

        # 提取单位
        unit = self._get_value(self.target_table, data_row, 7)
        if unit and unit != "见附件":  # 过滤无效值
            product["单位"] = unit

        # 提取单价（带单位格式化）
        price_value = self._get_value(self.target_table, data_row, 8)
        price_header = self._get_value(self.target_table, header_row, 8)
        if price_value and price_value != "见附件":  # 过滤无效值
            product["单价"] = self._format_amount(price_value, price_header)

        # 过滤空值字段
        return {k: v for k, v in product.items() if v}


# 使用示例
if __name__ == "__main__":
    with open("data.json", "r", encoding="utf8") as f:
        preprocessed_data = json.load(f)

    template = ProductInfoTemplate()
    if template.check_call_timing(preprocessed_data):
        result = template.extract(preprocessed_data)
        print("提取结果：", json.dumps(result, ensure_ascii=False, indent=2))
    else:
        print("当前数据不满足模板调用条件")