code: # -*- coding: UTF-8 -*-
  """
  Main module of preprocessor package. Can be executed by `python -m preprocessor`.
  """
  import os
  import pickle
  from random import shuffle
  
  import numpy as np
  
  from chicksexer.constant import POSITIVE_CLASS, NEGATIVE_CLASS, NEUTRAL_CLASS, CLASS2DEFAULT_CUTOFF
  from chicksexer.util import get_logger
  from preprocessor import PACKAGE_ROOT
  from preprocessor.dbpedia import gen_triples_from_file
  from preprocessor.gender_csv import gen_name_gender_from_csv
  from preprocessor.us_stats import compute_gender_probas
  from preprocessor.util import Name2Proba
  
  __author__ = 'kensk8er'
  
  _DATA_ROOT = os.path.join(PACKAGE_ROOT, os.path.pardir, 'data') 
  #_DATA_ROOT = os.path.join(PACKAGE_ROOT, 'data')
  
  
  _RAW_DATA_ROOT = os.path.join(_DATA_ROOT, 'raw')
  _PROCESSED_DATA_PATH = os.path.join(_DATA_ROOT, 'name2proba_{}.pkl')
  _NEUTRAL_NAME_AUGMENTATION_NUM = 100000
  _FEMALE_NAME_AUGMENTATION_NUM = 85000
  _TEST_DATA_SIZE = 10000  # the size of the whole dataset is ~400,000
  _LOGGER = get_logger(__name__)
  
  _CLASS2PROB = {
      POSITIVE_CLASS: 1.,
      NEUTRAL_CLASS: 0.5,
      NEGATIVE_CLASS: 0.,
  }
  
  
  def _process_csv(name2probfa):
      """Process csv files that list names and their gender."""
      file_names = ['Black-Female-Names.csv', 'Black-Male-Names.csv', 'Hispanic-Female-Names.csv',
                    'Hispanic-Male-Names.csv', 'Indian-Female-Names.csv', 'Indian-Male-Names.csv',
                    'White-Female-Names.csv', 'White-Male-Names.csv']
  
      for file_name in file_names:
          for name, gender in gen_name_gender_from_csv(os.path.join(_RAW_DATA_ROOT, file_name)):
              proba = _CLASS2PROB[gender]
              name2probfa[name] = proba
      return name2probfa
  
  
  def _process_dbpedia(name2proba):
      """Process genders_en.ttl downloaded from dbpedia dump."""
      file_name = 'genders_en.ttl'
      for name, gender in gen_triples_from_file(os.path.join(_RAW_DATA_ROOT, file_name)):
          proba = _CLASS2PROB[gender]
          name2proba[name] = proba
      return name2proba
  
  
  def _process_us_stats(name2proba, start_year=1940):
      """Process yobxxxx.txt files that list first names and their gender."""
      dir_path = os.path.join(_RAW_DATA_ROOT, 'US-Baby-Name-Stats')
      name2proba_stats = compute_gender_probas(dir_path, start_year)
      for name, proba in name2proba_stats.items():
          name2proba.set_fix_item(name, proba)
      return name2proba
  
  
  def _process_common_names(name2proba):
      """Process male/female.txt files that list common male/female names."""
  
      def process_common_names(file_name, gender, name2prob):
          with open(os.path.join(_RAW_DATA_ROOT, file_name), encoding='utf8') as file_:
              for line in file_:
                  if line.startswith('#') or line.startswith('\n'):
                      continue
                  name = line.strip()
                  name2prob[name] = _CLASS2PROB[gender]
          return name2prob
  
      file_name2gender = {
          'male.txt': POSITIVE_CLASS,
          'female.txt': NEGATIVE_CLASS,
      }
      for file_name, gender in file_name2gender.items():
          name2proba = process_common_names(file_name, gender, name2proba)
  
      return name2proba
  
  
  def _augment_full_names(name2proba, gender):
      """Augment neutral names"""
      if gender == 'neutral':
          augmentation_num = _NEUTRAL_NAME_AUGMENTATION_NUM
          low_proba = CLASS2DEFAULT_CUTOFF[NEGATIVE_CLASS]
          high_proba = CLASS2DEFAULT_CUTOFF[POSITIVE_CLASS]
      elif gender == 'female':
          augmentation_num = _FEMALE_NAME_AUGMENTATION_NUM
          low_proba = float('-inf')
          high_proba = CLASS2DEFAULT_CUTOFF[NEGATIVE_CLASS]
      else:
          raise ValueError('Invalid argument gender={}'.format(gender))
  
      neutral_names = [name for name, prob in name2proba.items()
                       if low_proba < prob < high_proba and ' ' not in name]
      multiple = augmentation_num // len(neutral_names)
  
      with open(os.path.join(_DATA_ROOT, 'surname2proba.pkl'), 'rb') as pickle_file:
          surname2proba = pickle.load(pickle_file)
          surnames, surname_probas = list(), list()
          for surname, proba in surname2proba.items():
              surnames.append(surname)
              surname_probas.append(proba)
  
      for neutral_name in neutral_names:
          proba = name2proba[neutral_name]
          sampled_surnames = np.random.choice(surnames, multiple, p=surname_probas)
          for surname in sampled_surnames:
              full_name = '{} {}'.format(neutral_name, surname)
              name2proba[full_name] = proba
  
      return name2proba
  
  
  def main():
      name2proba = Name2Proba()
      _LOGGER.info('Processing Dbpedia...')
      name2proba = _process_dbpedia(name2proba)
      _LOGGER.info('Processing CSVs...')
      name2proba = _process_csv(name2proba)
      _LOGGER.info('Processing US Stats...')
      name2proba = _process_us_stats(name2proba)
      _LOGGER.info('Processing Common Names...')
      name2proba = _process_common_names(name2proba)
      
      # these require surname2proba.pkl
      #_LOGGER.info('Augmenting Neutral Names...')
      #name2proba = _augment_full_names(name2proba, 'neutral')  
      #_LOGGER.info('Augmenting Female Names...')
      #name2proba = _augment_full_names(name2proba, 'female')
      #_LOGGER.info('Saving to the pickle files...')
  
      # randomly split into train/test set
      name2proba = dict(name2proba)
      assert len(name2proba) > _TEST_DATA_SIZE, 'Whole dataset size is not larger than test set size.'
      ids = list(range(len(name2proba)))
      shuffle(ids)
      test_ids = set(ids[:_TEST_DATA_SIZE])
      name2proba_train = dict()
      name2proba_test = dict()
  
      for id_, (name, proba) in enumerate(name2proba.items()):
          if id_ in test_ids:
              name2proba_test[name] = proba
          else:
              name2proba_train[name] = proba
  
      # write to pickle files
      with open(_PROCESSED_DATA_PATH.format('train'), 'wb') as train_file:
          pickle.dump(name2proba_train, train_file)
      with open(_PROCESSED_DATA_PATH.format('test'), 'wb') as test_file:
          pickle.dump(name2proba_test, test_file)
      with open(_PROCESSED_DATA_PATH.format('all'), 'wb') as all_file:
          pickle.dump(name2proba, all_file)
  
apis: numpy.random.choice