Adding support for setting number of threads to faster-whisper cpu, reading from cli, yaml or env var.

scraibe/misc.py

@@ -3,6 +3,7 @@ import yaml
 from argparse import Action
 from ast import literal_eval
 from torch.cuda import is_available
+from torch import get_num_threads, set_num_threads
 
 CACHE_DIR = os.getenv(
     "AUTOT_CACHE",
@@ -21,6 +22,8 @@ PYANNOTE_DEFAULT_CONFIG = os.path.join(PYANNOTE_DEFAULT_PATH, "config.yaml") \
 
 SCRAIBE_TORCH_DEVICE =  os.getenv("SCRAIBE_TORCH_DEVICE", "cuda" if is_available() else "cpu")
 
+SCRAIBE_NUM_THREADS = os.getenv("SCRAIBE_NUM_THREADS", min(8, get_num_threads()))
+
 def config_diarization_yaml(file_path: str, path_to_segmentation: str = None) -> None:
     """Configure diarization pipeline from a YAML file.
 
@@ -49,6 +52,28 @@ def config_diarization_yaml(file_path: str, path_to_segmentation: str = None) ->
         yaml.dump(yml, stream)
 
 
+def set_threads(parse_threads=None,
+                yaml_threads=None,
+                env_var_threads=None):
+    global SCRAIBE_NUM_THREADS
+    if parse_threads is not None:
+        if not isinstance(parse_threads, int):
+            # probably covered with int type of parser arg
+            raise ValueError(f"Type of --num-threads must be int, but the type is {type(parse_threads)}")
+        elif parse_threads < 1:
+            raise ValueError(f"Number of threads must be a positive integer, {parse_threads} was given")
+        else:
+            set_num_threads(parse_threads)
+            SCRAIBE_NUM_THREADS = parse_threads
+    elif yaml_threads is not None:
+        if not isinstance(yaml_threads, int):
+            raise ValueError(f"Type of num_threads must be int, but the type is {type(yaml_threads)}")
+        elif yaml_threads < 1:
+            raise ValueError(f"Number of threads must be a positive integer, {yaml_threads} was given")
+        else:
+            set_num_threads(yaml_threads)
+            SCRAIBE_NUM_THREADS = yaml_threads
+
 class ParseKwargs(Action):
     """
     Custom argparse action to parse keyword arguments.