From 8dc0b2b0b046002454475095c2db3344cbe0fca1 Mon Sep 17 00:00:00 2001
From: yangyaming <mxscmxsc@gmail.com>
Date: Wed, 23 Aug 2017 14:41:41 +0800
Subject: [PATCH] Make setup.py to support parallel processing.

---
 deploy/README.md  |  2 +-
 deploy/scorer.cpp |  7 +++--
 deploy/setup.py   | 70 +++++++++++++++++++++++++++++++++++++++++++----
 3 files changed, 70 insertions(+), 9 deletions(-)

diff --git a/deploy/README.md b/deploy/README.md
index 9bd55dd9a..90809ad35 100644
--- a/deploy/README.md
+++ b/deploy/README.md
@@ -25,7 +25,7 @@ git clone https://github.com/progschj/ThreadPool.git
 Then run the setup
 
 ```shell
-python setup.py install
+python setup.py install --num_processes 4
 cd ..
 ```
 
diff --git a/deploy/scorer.cpp b/deploy/scorer.cpp
index 17bb6e10d..233b4766d 100644
--- a/deploy/scorer.cpp
+++ b/deploy/scorer.cpp
@@ -1,5 +1,8 @@
 #include <iostream>
 #include <unistd.h>
+#include "lm/config.hh"
+#include "lm/state.hh"
+#include "lm/model.hh"
 #include "scorer.h"
 #include "decoder_utils.h"
 
@@ -24,7 +27,7 @@ void Scorer::load_LM(const char* filename) {
         exit(1);
     }
     RetriveStrEnumerateVocab enumerate;
-    Config config;
+    lm::ngram::Config config;
     config.enumerate_vocab = &enumerate;
     _language_model = lm::ngram::LoadVirtual(filename, config);
     _max_order = static_cast<lm::base::Model*>(_language_model)->Order();
@@ -43,7 +46,7 @@ void Scorer::load_LM(const char* filename) {
 double Scorer::get_log_cond_prob(const std::vector<std::string>& words) {
     lm::base::Model* model = static_cast<lm::base::Model*>(_language_model);
     double cond_prob;
-    State state, tmp_state, out_state;
+    lm::ngram::State state, tmp_state, out_state;
     // avoid to inserting <s> in begin
     model->NullContextWrite(&state);
     for (size_t i = 0; i < words.size(); ++i) {
diff --git a/deploy/setup.py b/deploy/setup.py
index 1342478b2..7a4b7e02c 100644
--- a/deploy/setup.py
+++ b/deploy/setup.py
@@ -1,17 +1,75 @@
-from setuptools import setup, Extension
+"""Script to build and install decoder package."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from setuptools import setup, Extension, distutils
 import glob
 import platform
-import os
+import os, sys
+import multiprocessing.pool
+import argparse
+
+parser = argparse.ArgumentParser(description=__doc__)
+parser.add_argument(
+    "--num_processes",
+    default=1,
+    type=int,
+    help="Number of cpu processes to build package. (default: %(default)d)")
+args = parser.parse_known_args()
+
+# reconstruct sys.argv to pass to setup below
+sys.argv = [sys.argv[0]] + args[1]
+
+
+# monkey-patch for parallel compilation
+# See: https://stackoverflow.com/a/13176803
+def parallelCCompile(self,
+                     sources,
+                     output_dir=None,
+                     macros=None,
+                     include_dirs=None,
+                     debug=0,
+                     extra_preargs=None,
+                     extra_postargs=None,
+                     depends=None):
+    # those lines are copied from distutils.ccompiler.CCompiler directly
+    macros, objects, extra_postargs, pp_opts, build = self._setup_compile(
+        output_dir, macros, include_dirs, sources, depends, extra_postargs)
+    cc_args = self._get_cc_args(pp_opts, debug, extra_preargs)
+
+    # parallel code
+    def _single_compile(obj):
+        try:
+            src, ext = build[obj]
+        except KeyError:
+            return
+        self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
+
+    # convert to list, imap is evaluated on-demand
+    thread_pool = multiprocessing.pool.ThreadPool(args[0].num_processes)
+    list(thread_pool.imap(_single_compile, objects))
+    return objects
 
 
 def compile_test(header, library):
     dummy_path = os.path.join(os.path.dirname(__file__), "dummy")
-    command = "bash -c \"g++ -include " + header + " -l" + library + " -x c++ - <<<'int main() {}' -o " + dummy_path + " >/dev/null 2>/dev/null && rm " + dummy_path + " 2>/dev/null\""
+    command = "bash -c \"g++ -include " + header \
+                + " -l" + library + " -x c++ - <<<'int main() {}' -o " \
+                + dummy_path + " >/dev/null 2>/dev/null && rm " \
+                + dummy_path + " 2>/dev/null\""
     return os.system(command) == 0
 
 
-FILES = glob.glob('kenlm/util/*.cc') + glob.glob('kenlm/lm/*.cc') + glob.glob(
-    'kenlm/util/double-conversion/*.cc')
+# hack compile to support parallel compiling
+distutils.ccompiler.CCompiler.compile = parallelCCompile
+
+FILES = glob.glob('kenlm/util/*.cc') \
+        + glob.glob('kenlm/lm/*.cc') \
+        + glob.glob('kenlm/util/double-conversion/*.cc')
+
+FILES += glob.glob('openfst-1.6.3/src/lib/*.cc')
+
 FILES = [
     fn for fn in FILES if not (fn.endswith('main.cc') or fn.endswith('test.cc'))
 ]
@@ -40,7 +98,7 @@ decoders_module = [
     Extension(
         name='_swig_decoders',
         sources=FILES + glob.glob('*.cxx') + glob.glob('*.cpp'),
-        language='C++',
+        language='c++',
         include_dirs=['.', 'kenlm', 'openfst-1.6.3/src/include', 'ThreadPool'],
         libraries=LIBS,
         extra_compile_args=ARGS)