import os

from paddlenlp import Taskflow

from config.keys import PATIENT_NAME, ADMISSION_DATE, DISCHARGE_DATE, MEDICAL_EXPENSES, PERSONAL_CASH_PAYMENT, \
    PERSONAL_ACCOUNT_PAYMENT, PERSONAL_FUNDED_AMOUNT, MEDICAL_INSURANCE_TYPE, HOSPITAL, DEPARTMENT, DOCTOR
from config.photo_review import IE_BATCH_SIZE, LAYOUT_ANALYSIS, CUDA_VISIBLE_DEVICES

os.environ["CUDA_VISIBLE_DEVICES"] = CUDA_VISIBLE_DEVICES

# 基本医保结算单
settlement_list_schema = \
    PATIENT_NAME + ADMISSION_DATE + DISCHARGE_DATE + MEDICAL_EXPENSES + PERSONAL_CASH_PAYMENT \
    + PERSONAL_ACCOUNT_PAYMENT + PERSONAL_FUNDED_AMOUNT + MEDICAL_INSURANCE_TYPE
settlement_ie = Taskflow("information_extraction", schema=settlement_list_schema, model="uie-x-base",
                         task_path="config/model/settlement_list_model", layout_analysis=LAYOUT_ANALYSIS,
                         batch_size=IE_BATCH_SIZE)

# 出院记录
discharge_record_schema = \
    HOSPITAL + DEPARTMENT + PATIENT_NAME + ADMISSION_DATE + DISCHARGE_DATE + DOCTOR
discharge_ie = Taskflow("information_extraction", schema=discharge_record_schema, model="uie-x-base",
                        task_path="config/model/discharge_record_model", layout_analysis=LAYOUT_ANALYSIS,
                        batch_size=IE_BATCH_SIZE)

# 费用清单
cost_list_schema = PATIENT_NAME + ADMISSION_DATE + DISCHARGE_DATE + MEDICAL_EXPENSES
cost_ie = Taskflow("information_extraction", schema=cost_list_schema, model="uie-x-base",
                   task_path="config/model/cost_list_model", layout_analysis=LAYOUT_ANALYSIS, batch_size=IE_BATCH_SIZE)