Merge remote-tracking branch 'upstream/master'

1 files changed, 9 insertions, 1 deletions

diff --git a/python/pkg/cdec/sa/extract.py b/python/pkg/cdec/sa/extract.py
index 782bed8b..b6502c52 100644
--- a/python/pkg/cdec/sa/extract.py
+++ b/python/pkg/cdec/sa/extract.py
@@ -8,6 +8,7 @@ import logging
 import signal
 import multiprocessing as mp
 import cdec.sa
+from cdec.sa._sa import monitor_cpu
 
 extractor, prefix = None, None
 online, compress = False, False
@@ -62,6 +63,7 @@ def extract(inp):
     return '<seg grammar="{}" id="{}">{}</seg>{}'.format(grammar_file, i, sentence, suffix)
 
 def main():
+    global online
     logging.basicConfig(level=logging.INFO)
     parser = argparse.ArgumentParser(description='Extract grammars from a compiled corpus.')
     parser.add_argument('-c', '--config', required=True,
@@ -87,7 +89,10 @@ def main():
             sys.stderr.write('Error: feature definition file <{}>'
                     ' should be a python module\n'.format(featdef))
             sys.exit(1)
-    
+
+    online = args.online
+
+    start_time = monitor_cpu()
     if args.jobs > 1:
         logging.info('Starting %d workers; chunk size: %d', args.jobs, args.chunksize)
         pool = mp.Pool(args.jobs, make_extractor, (args,))
@@ -101,5 +106,8 @@ def main():
         for output in map(extract, enumerate(sys.stdin)):
             print(output)
 
+    stop_time = monitor_cpu()
+    logging.info("Overall extraction step took %f seconds", stop_time - start_time)
+
 if __name__ == '__main__':
     main()