BIDI_EXTRACTION_AI/BiddingKG/dl/money/test_re_money_total_unit.py

import json
import re
import sys, os
import time

import pandas as pd
from bs4 import BeautifulSoup
sys.path.append(os.path.abspath("../.."))
from BiddingKG.dl.interface.extract import predict


def bidi_predict(html_str):
    content = html_str
    # content = "<div>总价：1110</div>"
    result_dict = json.loads(predict("1", content))
    return result_dict


def test_csv(_path):
    start_time = time.time()
    df = pd.read_csv(_path)

    # total money
    predict_list_1 = []
    predict_list_2 = []
    for index, row in df.iterrows():
        # if index >= 1000:
        #     break

        if index % 50 == 0:
            print("="*30, "Loop", index, "="*30)

        html_str = row["dochtmlcon"]
        # html_str = df.loc[75, "dochtmlcon"]
        # print(html_str)

        # 先筛选
        # possible = '((合计.?金额|合.?计|总.?价|单.?价)(（元）)?([:： ]))' \
        #            '|([0-9.，,]+([（(]?元[)）]?)?/)'
        # if not re.search(possible, html_str):
        #     predict_list_1.append(str([]))
        #     predict_list_2.append(str([]))
        #     continue

        # 先经过模型处理
        result_dict = bidi_predict(html_str)

        # 获取总价单价
        word_list_1 = result_dict.get("total_money")
        word_list_2 = result_dict.get("unit_money")

        if word_list_1:
            predict = word_list_1
        else:
            predict = []
        print("predict total money", predict)
        predict_list_1.append(str(predict))

        if word_list_2:
            predict = word_list_2
        else:
            predict = []
        print("predict unit money", predict)
        predict_list_2.append(str(predict))

    predict_df_1 = pd.DataFrame(predict_list_1)
    predict_df_2 = pd.DataFrame(predict_list_2)
    df = pd.concat([df, predict_df_1, predict_df_2], axis=1)
    df.to_csv(_path)
    print("finish write!", time.time()-start_time)


if __name__ == "__main__":
    path = "D:\\BIDI_DOC\\比地_文档\\总价单价_result.csv"
    test_csv(path)