passed the test

Author: pengwei715

pii_recognizer/data/config.csv

@@ -0,0 +1,3 @@
+input_file,output_file
+data/pii.txt,data/pii_out.txt
+data/letter.txt,data/letter_out.txt

pii_recognizer/pii_recognizer.py

@@ -18,6 +18,7 @@
 import pathlib
 import tempfile
 import warnings
+import pandas as pd
 from collections.abc import Iterable
 from multiprocessing import Pool, cpu_count
 from typing import Any, Dict, List, Optional, Set, Tuple, Union
@@ -1037,7 +1038,6 @@ def recognize_pii_parallel(
         )
         for _, row in config_df.iterrows()
     ]
-
     # Create a pool of processes and distribute the tasks
     with Pool(processes=num_processes) as pool:
         pool.starmap(recognize_pii_one_file, tasks)

pii_recognizer/test_pii_recognizer.py

@@ -22,6 +22,7 @@
     _get_analyzer_engine,
     _anonymize,
     _annotate,
+    recognize_pii_parallel
 )
 
 
@@ -211,3 +212,31 @@ def test_only_entities(fake_data):
     analyzer = _get_analyzer_engine(entities=list(ENTITIES.keys())[:5])
     res, results = _process(text, analyzer, score_threshold=0.5)
     assert any(entity in res for entity in ENTITIES.keys())
+
+
+
+def test_parallel():
+    ENTITIES = {
+        "LOCATION": "location",
+        "PERSON": "name",
+        "ORGANIZATION": "organization",
+        "MAC_ADDRESS": "mac_address",
+        "US_BANK_NUMBER": "us_bank_number",
+        "IMEI": "imei",
+        "TITLE": "title",
+        "LICENSE_PLATE": "license_plate",
+        "US_PASSPORT": "us_passport",
+        "CURRENCY": "currency",
+        "ROUTING_NUMBER": "routing_number",
+        "US_ITIN": "us_itin",
+        "US_BANK_NUMBER": "us_bank_number",
+        "AGE": "age",
+        "CREDIT_CARD": "credit_card",
+        "SSN": "ssn",
+        "PHONE": "phone",
+        "EMAIL": "email",
+        "PASSWORD": "password",
+        "SWIFT_CODE": "swift_code",
+    }
+    recognize_pii_parallel(config_input_output = "data/config.csv", score_threshold = 0.5, entities = list(ENTITIES.keys()), model = "whole")
+