diff --git a/fast_api.py b/fast_api.py index 52bd174..fa1fa05 100644 --- a/fast_api.py +++ b/fast_api.py @@ -161,7 +161,7 @@ async def extract_information(input_data: InputText): converted_id = text_to_number(target_value) - schema = ["姓名", "卡号", "身份证号"] + schema = ["姓名", '嫌疑人', '涉案人员', "身份证号", "交易证件号", "卡号", "交易卡号", "银行卡号", ] ie = Taskflow('information_extraction', schema=schema, model='uie-base') extracted_info = ie(converted_id) @@ -170,9 +170,11 @@ async def extract_information(input_data: InputText): for key, value in item.items(): result[key.lower()] = value[0]['text'] - extracted_result = ExtractedInfo(name=result.get('姓名', ''), - cardNumber=result.get('卡号', ''), - idNnumber=result.get('身份证号', '')) + extracted_result = ExtractedInfo( + name=result.get('姓名', '') or result.get('嫌疑人', '') or result.get('涉案人员', ''), + cardNumber=result.get('卡号', '') or result.get('交易卡号', '') or result.get('银行卡号', ''), + idNnumber=result.get('身份证号', '') or result.get('交易证件号', '') or result.get('交易证件号', '') + ) return extracted_result