Black on file

Signed-off-by: Vibhu Jawa <vibhujawa@gmail.com>

examples/distributed_data_classification_examples/domain_api_example.py

@@ -56,7 +56,9 @@ def main(args):
     model_file_name = "/home/nfs/syurick/LLM_domain_classifier_inference/GoogleDebertaAgree_v3b_bce_maxlen512_bs64_noRef_best.pth"
 
     # Input can be a string or list
-    input_file_path = "/home/nfs/syurick/LLM_domain_classifier_inference/4360_results_jsonl_dir/"
+    input_file_path = (
+        "/home/nfs/syurick/LLM_domain_classifier_inference/4360_results_jsonl_dir/"
+    )
     output_file_path = "/raid/vjawa/output_file.parquet"
 
     client = get_client(args, cluster_type=args.device)
@@ -68,7 +70,7 @@ def main(args):
     domain_classifier = DomainClassifier(
         model_file_name=model_file_name,
         labels=labels,
-        #filter_by=["Games", "Sports"],
+        # filter_by=["Games", "Sports"],
     )
     result_dataset = domain_classifier(dataset=input_dataset)
 

nemo_curator/distributed_data_classification/domain_classifier_inference.py

@@ -15,6 +15,7 @@
 import os
 import time
 import warnings
+
 os.environ["RAPIDS_NO_INITIALIZE"] = "1"
 import torch
 from packaging import version
@@ -38,9 +39,9 @@
 from nemo_curator.utils.file_utils import get_remaining_files
 
 
-
 warnings.filterwarnings("ignore")
 
+
 @dataclass
 class Config:
     model = "microsoft/deberta-v3-base"
@@ -49,7 +50,9 @@ class Config:
 
 
 class CustomModel(nn.Module):
-    def __init__(self, config, out_dim, config_path=None, pretrained=False, autocast=False):
+    def __init__(
+        self, config, out_dim, config_path=None, pretrained=False, autocast=False
+    ):
         super().__init__()
         self.config = config
         if config_path is None:
@@ -95,7 +98,7 @@ def forward(self, batch):
             feature = self.feature(batch["input_ids"], batch["attention_mask"])
             output = self.fc(self.fc_dropout(feature))
         return torch.softmax(output[:, 0, :], dim=1)
-    
+
 
 def load_model(config, device, model_path, autocast):
     """
@@ -111,7 +114,9 @@ def load_model(config, device, model_path, autocast):
         The loaded model.
 
     """
-    model = CustomModel(config, out_dim=27, config_path=None, pretrained=True, autocast=autocast)
+    model = CustomModel(
+        config, out_dim=27, config_path=None, pretrained=True, autocast=autocast
+    )
     model = model.to(device)
     if os.path.exists(model_path):
         sd = torch.load(os.path.join(model_path), map_location="cpu")
@@ -127,22 +132,24 @@ class DomainModel(HFModel):
     def __init__(self, config, model_path=None, autocast=False):
         self.config = config
         self.model_path = model_path
-        self.autocast=autocast
+        self.autocast = autocast
         super().__init__(self.config.model)
 
     def load_model(self, device="cuda"):
-        return load_model(self.config, device=device,
-                         model_path=self.model_path or self.path_or_name,
-                         autocast=self.autocast)
-
+        return load_model(
+            self.config,
+            device=device,
+            model_path=self.model_path or self.path_or_name,
+            autocast=self.autocast,
+        )
+
     def load_tokenizer(self):
         return DebertaV2TokenizerFast.from_pretrained(self.config.model)
 
     def load_config(self):
         return AutoConfig.from_pretrained(self.path_or_name)
 
 
-
 def main():
     labels = [
         "Adult",
@@ -205,11 +212,11 @@ def main():
             file_type=args.input_file_type,
             add_filename=add_filename,
         )
-        df['sliced_text'] = df['text'].str.slice(0, max_chars)
+        df["sliced_text"] = df["text"].str.slice(0, max_chars)
         columns_to_keep_list = df.columns.to_list()
-        columns_to_keep_list.remove('sliced_text')
+        columns_to_keep_list.remove("sliced_text")
 
-        model_path ="/home/nfs/syurick/LLM_domain_classifier_inference/GoogleDebertaAgree_v3b_bce_maxlen512_bs64_best.pth"
+        model_path = "/home/nfs/syurick/LLM_domain_classifier_inference/GoogleDebertaAgree_v3b_bce_maxlen512_bs64_best.pth"
         model = DomainModel(Config, model_path=model_path, autocast=args.autocast)
         pipe = op.Sequential(
             op.Tokenizer(model, cols=["sliced_text"], tokenizer_type="sentencepiece"),

nemo_curator/modules/distributed_data_classifier.py

@@ -28,6 +28,7 @@
 
 from nemo_curator.datasets import DocumentDataset
 
+
 @dataclass
 class domain_Config:
     model = "microsoft/deberta-v3-base"
@@ -36,7 +37,9 @@ class domain_Config:
 
 
 class CustomModel(nn.Module):
-    def __init__(self, config, out_dim, config_path=None, pretrained=False, autocast=False):
+    def __init__(
+        self, config, out_dim, config_path=None, pretrained=False, autocast=False
+    ):
         super().__init__()
         self.config = config
         if config_path is None:
@@ -82,7 +85,7 @@ def forward(self, batch):
             feature = self.feature(batch["input_ids"], batch["attention_mask"])
             output = self.fc(self.fc_dropout(feature))
         return torch.softmax(output[:, 0, :], dim=1)
-    
+
 
 def _load_model(model, device, model_path):
     """
@@ -107,7 +110,6 @@ def _load_model(model, device, model_path):
     return model
 
 
-
 class DistributedDataClassifier(ABC):
     """Abstract class for running multi-node multi-GPU data classification"""
 
@@ -144,7 +146,6 @@ def __call__(self, dataset: DocumentDataset):
     def _run_classifier(self):
         pass
 
-
     def _filter_documents(
         self,
         dataset: DocumentDataset,
@@ -167,20 +168,26 @@ def __init__(self, config, out_dim=None, model_path=None, autocast=False):
         self.config = config
         self.out_dim = out_dim
         self.model_path = model_path
-        self.autocast=autocast
+        self.autocast = autocast
         super().__init__(self.config.model)
 
     def load_model(self, device="cuda"):
-        model = CustomModel(self.config, out_dim=self.out_dim, config_path=None, pretrained=True, autocast=self.autocast)
+        model = CustomModel(
+            self.config,
+            out_dim=self.out_dim,
+            config_path=None,
+            pretrained=True,
+            autocast=self.autocast,
+        )
         return _load_model(model, device, self.model_path)
-    
+
     def load_tokenizer(self):
         return DebertaV2TokenizerFast.from_pretrained(self.config.model)
 
     def load_config(self):
         return AutoConfig.from_pretrained(self.path_or_name)
 
-    
+
 class DomainClassifier(DistributedDataClassifier):
     def __init__(
         self,
@@ -197,11 +204,12 @@ def __init__(
         if out_dim is None:
             out_dim = len(labels)
 
-        model = DomainModel(config=domain_Config,
-                            out_dim=out_dim, 
-                            model_path=model_file_name,
-                            autocast=autocast)
-
+        model = DomainModel(
+            config=domain_Config,
+            out_dim=out_dim,
+            model_path=model_file_name,
+            autocast=autocast,
+        )
 
         super().__init__(
             model=model,
@@ -219,13 +227,17 @@ def _run_classifier(self, dataset: DocumentDataset):
         print("Starting domain classifier inference", flush=True)
 
         df = dataset.df
-        df['sliced_text'] = df['text'].str.slice(0, self.max_chars)
+        df["sliced_text"] = df["text"].str.slice(0, self.max_chars)
         columns_to_keep_list = df.columns.to_list()
-        columns_to_keep_list.remove('sliced_text')
+        columns_to_keep_list.remove("sliced_text")
 
         pipe = op.Sequential(
-            op.Tokenizer(self.model, cols=["sliced_text"], tokenizer_type="sentencepiece"),
-            op.Predictor(self.model, sorted_data_loader=True, batch_size=self.batch_size),
+            op.Tokenizer(
+                self.model, cols=["sliced_text"], tokenizer_type="sentencepiece"
+            ),
+            op.Predictor(
+                self.model, sorted_data_loader=True, batch_size=self.batch_size
+            ),
             op.Labeler(self.labels, cols=["preds"]),
             repartition=df.npartitions,
             keep_cols=columns_to_keep_list,