change to lower API

Author: donglihe-hub

libmultilabel/nn/data_utils.py

@@ -1,8 +1,9 @@
 import csv
 import gc
 import logging
+import multiprocessing as mp
+import os
 import warnings
-from concurrent.futures import ProcessPoolExecutor
 
 import pandas as pd
 import torch
@@ -176,9 +177,9 @@ def _load_raw_data(data, is_test=False, tokenize_text=True, remove_no_label_data
 
     data["label"] = data["label"].astype(str).map(lambda s: s.split())
     if tokenize_text:
-        # multiprocessing requires serializable objects
-        with ProcessPoolExecutor() as executor:
-            data["text"] = pd.Series(tqdm(executor.map(tokenize, data["text"]), total=len(data["text"])))
+        cpu_count = min(32, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count())
+        with mp.get_context("fork").Pool(processes=cpu_count) as p:
+            data["text"] = pd.Series(tqdm(p.map(tokenize, data["text"]), total=len(data["text"])))
     data = data.to_dict("records")
     if not is_test:
         num_no_label_data = sum(1 for d in data if len(d["label"]) == 0)