style code

parallelformers/parallel/process.py

@@ -16,17 +16,18 @@
 import io
 import os
 import pickle
+import random
 import traceback
 import types
 from contextlib import suppress
 from dataclasses import _is_dataclass_instance, asdict
-from time import time
 from inspect import signature
+from time import time
 from typing import Any, List, Union
-import torch.distributed as dist
+
 import numpy as np
-import random
 import torch
+import torch.distributed as dist
 import torch.multiprocessing as mp
 import torch.nn as nn
 from transformers.file_utils import ModelOutput

parallelformers/parallelize.py

@@ -77,9 +77,9 @@ def __init__(
         master_port: int = 29500,
         backend="nccl",
         verbose: str = None,
-        init_method="spawn",
+        init_method: str = "spawn",
         daemon: bool = True,
-        seed: int = None
+        seed: int = None,
     ):
         self.init_environments(
             num_gpus,

tests/causal_lm.py

@@ -13,10 +13,11 @@
 # limitations under the License.
 
 import os
+import random
 import unittest
 from argparse import ArgumentParser
+
 import numpy as np
-import random
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -33,7 +34,7 @@ def test_generation(self, model, tokens, tokenizer):
             max_length=40,
             no_repeat_ngram_size=4,
             do_sample=True,
-            top_p=0.7
+            top_p=0.7,
         )
 
         gen = tokenizer.batch_decode(output)[0]