refactor gender country matching

app.py

@@ -1,14 +1,17 @@
 
 '''Streamlit app for Student Name Detection models.'''
-from piilo.analyzer import prepare_analyzer
-from piilo.anonymizer import surrogate_anonymizer
-import pandas as pd
-from annotated_text import annotated_text
-from json import JSONEncoder
 import json
+import os
 import warnings
+from json import JSONEncoder
+
+import pandas as pd
 import streamlit as st
-import os
+from annotated_text import annotated_text
+
+from piilo.engines.analyzer import prepare_analyzer
+from piilo.engines.anonymizer import surrogate_anonymizer
+
 os.environ['TOKENIZERS_PARALLELISM'] = 'false'
 warnings.filterwarnings('ignore')
 
@@ -56,7 +59,7 @@ st.set_page_config(page_title='Student Name Detector (English)', layout='wide')
 
 # Side bar
 st.sidebar.markdown(
-    '''Detect and anonymize PII in text using an [NLP model](https://huggingface.co/langdonholmes/en_student_name_detector) [trained](https://github.com/aialoe/deidentification-pipeline) on student-generated text collected by Coursera.
+    '''Detect and anonymize PII in text using an [NLP model](https://huggingface.co/langdonholmes/en_student_name_detector) [trained](https://github.com/aialoe/deidentification-pipeline) on student-generated text collected from a massive online open-enrollment course.
 '''
 )
 
@@ -74,7 +77,7 @@ st_return_decision_process = st.sidebar.checkbox(
     'Add analysis explanations in json')
 
 st.sidebar.info(
-    'This is part of a deidentification project for student-generated text.'
+    'This is part of a project to develop new anonymization systems that are appropriate for student-generated text.'
 )
 
 # Main panel

piilo/{anonymizer.py → engines/anonymizer.py}

@@ -10,11 +10,14 @@ from presidio_anonymizer import AnonymizerEngine
 from presidio_anonymizer.entities import OperatorConfig
 from presidio_anonymizer.operators import OperatorType
 
-name_table = Path(__file__).parent.parent / 'data' / 'ascii_names.parquet'
+data = Path(__file__).parent.parent.parent / 'data' 
+name_table = data / 'ascii_names.parquet'
 
 logger = logging.getLogger('anonymizer')
-
 class NameDatabase(NameDataset):
+    '''A wrapper around the names_dataset.NameDataset class.
+    '''
+    
     def __init__(self) -> None:
         super().__init__()
     
@@ -29,27 +32,31 @@ class NameDatabase(NameDataset):
         return {'first_name': fn, 'last_name': ln}
        
     def get_gender(self, first_names: str) -> str:
-        '''Return the most frequent gender code for a specific last name,
+        '''Return the most frequent gender code for the provided first name,
         or None if a match cannot be found.
         '''
         gender = NameWrapper(self.search(first_names)).gender
         return gender if gender else None
 
     def get_country(self, last_names: str) -> str:
-        '''Return the most frequent country code for a specific last name,
+        '''Return the most frequent country code for a the provided last name,
         or None if a match cannot be found.
         '''
         country = NameWrapper(self.search(last_names)).country
         return country if country else None
     
 class surrogate_anonymizer(AnonymizerEngine):
+    '''A wrapper around the presidio_anonymizer.AnonymizerEngine class.
+    '''
+    
     def __init__(self):
         super().__init__()
         self.names_db = NameDatabase()
         self.names_df = pd.read_parquet(name_table)
         
         # keep track of names we have seen
-        self.seen_names = dict()
+        self.seen_first_names = dict()
+        self.seen_last_names = dict()
         
     def get_random_name(
             self,
@@ -63,58 +70,74 @@ class surrogate_anonymizer(AnonymizerEngine):
         returns two rows of the names dataframe
         '''
         names_view = self.names_df
+        
         if country:
             names_view = names_view[names_view['country'] == country]
+            
         if gender:
             names_view = names_view[names_view['gender'] == gender]
+            
         if names_view.size < 25:
-            return self.names_df.sample(n=2, weights=self.names_df['count'])
-        return names_view.sample(n=2, weights=names_view['count'])
+            # If we don't have enough names, just return a random sample
+            return self.names_df.sample(n=1, weights=self.names_df['count'])
+        
+        return names_view.sample(n=1, weights=names_view['count'])
 
     def generate_surrogate(self, original_name: str) -> str:
         '''Generate a surrogate name.
         '''
+        
         if original_name == 'PII':
             # Every time we call this function, Presidio will validate it
             # by testing that the function returns a str when the input is
-            # 'PII'. Bypass this test.
+            # 'PII'. We don't need to run below code in this case.
             return 'PII'
         
-        # If we have seen this name before, return the same surrogate
-        if original_name in self.seen_names:
-            return self.seen_names[original_name]
-        
         # Use nameparser to split the name
         name = HumanName(original_name)
+        new_name = HumanName()
+        gender, country = None, None
         
-        gender = self.names_db.get_gender(name.first) if name.first else None
-        logger.info(f'Gender set to {gender}')
-        country = self.names_db.get_country(name.last) if name.last else None
-        logger.info(f'Country set to {country}')
-        
-        surrogate_name = ''
-        
-        name_candidates = self.get_random_name(gender=gender, country=country)
-        
-        surrogate_name += name_candidates.iloc[0]['first']
-        logger.info(f'First name surrogate is {surrogate_name}')
-        
+        # First check if we have seen this name before
         if name.last:
-            logger.info(f'Last name surrogate is {name_candidates.iloc[1]["last"]}')
-            surrogate_name += ' ' + name_candidates.iloc[1]['last']
+            if name.last in self.seen_last_names:
+                new_name.last = self.seen_last_names[name.last]
+            else:
+                # Sample last name, matching country
+                country = self.names_db.get_country(name.last)
+                logger.info(f'Country set to {country}')
+                new_name.last = self.get_random_name(
+                    country=country,
+                    )['last'].iloc[0]
+                logger.info(f'Last name surrogate is {new_name.last}')
+                
+        if name.first:
+            if name.first in self.seen_first_names:
+                new_name.first = self.seen_first_names[name.first]
+            else:
+                # Sample first name matching gender and country, if available.
+                gender = self.names_db.get_gender(name.first)
+                logger.info(f'Gender set to {gender}')
+                new_name.first = self.get_random_name(
+                    gender=gender,
+                    country=country,
+                    )['first'].iloc[0]
+                logger.info(f'First name surrogate is {new_name.first}')
             
-        logger.info(f'Returning surrogate name {surrogate_name}')
+        logger.info(f'Returning surrogate name {new_name}')
         
-        self.seen_names[original_name] = surrogate_name
+        self.seen_first_names[name.first] = new_name.first
+        self.seen_last_names[name.last] = new_name.last
                 
-        return surrogate_name
+        return str(new_name)
 
     def anonymize(
         self,
         text: str,
         analyzer_results: List[RecognizerResult]
         ):
-        '''Anonymize identified input using Presidio Anonymizer.'''
+        '''Anonymize identified input using Presidio Anonymizer.
+        '''
         
         if not text:
             return
@@ -149,6 +172,8 @@ if __name__ == '__main__':
     anonymizer = surrogate_anonymizer()
     
     test_names = ['Nora Wang',
+                  'John Williams',
+                  'John H. Williams',
                   'MJ',
                   '',
                   '(',

main.py → piilo/main.py

@@ -5,31 +5,26 @@ import logging
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 
-from piilo.analyzer import prepare_analyzer
-from piilo.anonymizer import surrogate_anonymizer
-from piilo.models.anonymize import AnonymizeRequest, AnonymizeResponse
+from engines.analyzer import prepare_analyzer
+from engines.anonymizer import surrogate_anonymizer
+from models.anonymize import AnonymizeRequest, AnonymizeResponse
 
-# Define Student Name Detection Model
 configuration = {
     'nlp_engine_name': 'spacy',
     'models': [
         {'lang_code': 'en', 'model_name': 'en_student_name_detector'}],
 }
 
-# set up logger for this module
 logger = logging.getLogger('api')
 logging.basicConfig(level=logging.INFO)
 
-# Load Custom Presidio Analyzer and Anonymizer
-logger.info("Loading Presidio Analyzer and Anonymizer")
+logger.info("Loading Custom Presidio Analyzer and Anonymizer...")
 analyzer = prepare_analyzer(configuration)
 anonymizer = surrogate_anonymizer()
-logger.info("Loaded Presidio Analyzer and Anonymizer")
+logger.info("Loading Successful!")
 
-# Initialize FastAPI
 app = FastAPI()
 
-# Enable CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -44,8 +39,7 @@ def hello():
 
 @app.post("/anonymize")
 def anonymize(anon_req: AnonymizeRequest) -> AnonymizeResponse:
-    '''Anonymize PII in text using a custom Presidio Analyzer and Anonymizer
-    '''
+    
     analyzer_result = analyzer.analyze(anon_req.raw_text,
                                        entities=anon_req.entities,
                                        language=anon_req.language,
@@ -62,7 +56,6 @@ def anonymize(anon_req: AnonymizeRequest) -> AnonymizeResponse:
 
 if __name__ == "__main__":
     import os
-
     import uvicorn
     
     uvicorn.run(