62 lines
2.5 KiB
Python
62 lines
2.5 KiB
Python
import pytest
|
|
from app.core.document_handlers.ner_processor import NerProcessor
|
|
|
|
def test_generate_masked_mapping():
|
|
processor = NerProcessor()
|
|
unique_entities = [
|
|
{'text': '李雷', 'type': '人名'},
|
|
{'text': '李明', 'type': '人名'},
|
|
{'text': '王强', 'type': '人名'},
|
|
{'text': 'Acme Manufacturing Inc.', 'type': '英文公司名', 'industry': 'manufacturing'},
|
|
{'text': 'Google LLC', 'type': '英文公司名'},
|
|
{'text': 'A公司', 'type': '公司名称'},
|
|
{'text': 'B公司', 'type': '公司名称'},
|
|
{'text': 'John Smith', 'type': '英文人名'},
|
|
{'text': 'Elizabeth Windsor', 'type': '英文人名'},
|
|
{'text': '华梦龙光伏项目', 'type': '项目名'},
|
|
{'text': '案号12345', 'type': '案号'},
|
|
{'text': '310101198802080000', 'type': '身份证号'},
|
|
{'text': '9133021276453538XT', 'type': '社会信用代码'},
|
|
]
|
|
linkage = {
|
|
'entity_groups': [
|
|
{
|
|
'group_id': 'g1',
|
|
'group_type': '公司名称',
|
|
'entities': [
|
|
{'text': 'A公司', 'type': '公司名称', 'is_primary': True},
|
|
{'text': 'B公司', 'type': '公司名称', 'is_primary': False},
|
|
]
|
|
},
|
|
{
|
|
'group_id': 'g2',
|
|
'group_type': '人名',
|
|
'entities': [
|
|
{'text': '李雷', 'type': '人名', 'is_primary': True},
|
|
{'text': '李明', 'type': '人名', 'is_primary': False},
|
|
]
|
|
}
|
|
]
|
|
}
|
|
mapping = processor._generate_masked_mapping(unique_entities, linkage)
|
|
# 人名
|
|
assert mapping['李雷'].startswith('李某')
|
|
assert mapping['李明'].startswith('李某')
|
|
assert mapping['王强'].startswith('王某')
|
|
# 英文公司名
|
|
assert mapping['Acme Manufacturing Inc.'] == 'MANUFACTURING'
|
|
assert mapping['Google LLC'] == 'COMPANY'
|
|
# 公司名同组
|
|
assert mapping['A公司'] == mapping['B公司']
|
|
assert mapping['A公司'].endswith('公司')
|
|
# 英文人名
|
|
assert mapping['John Smith'] == 'J*** S***'
|
|
assert mapping['Elizabeth Windsor'] == 'E*** W***'
|
|
# 项目名
|
|
assert mapping['华梦龙光伏项目'].endswith('项目')
|
|
# 案号
|
|
assert mapping['案号12345'] == '***'
|
|
# 身份证号
|
|
assert mapping['310101198802080000'] == 'XXXXXX'
|
|
# 社会信用代码
|
|
assert mapping['9133021276453538XT'] == 'XXXXXXXX' |