Fix code quality tests (#174)

* Revert "Filter examples by `num_chars` to include in a batch (#137)"

This reverts commit 8021ce7.

* fix test (code quality)

* fix code_quality

* fix

* fix style

Author: manandey

bsmetadata/experiments/sample.py

@@ -2,9 +2,9 @@
 
 from datasets import load_dataset
 from torch.utils.data import DataLoader
+from transformers import PreTrainedTokenizerBase
 
 from bsmetadata.input_pipeline import DataConfig
-from transformers import PreTrainedTokenizerBase
 
 
 @dataclass

bsmetadata/experiments/with_metadata.py

@@ -6,9 +6,9 @@
 from datasets import config, load_dataset
 from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
+from transformers import default_data_collator
 
 from bsmetadata.metadata_utils import add_metadata_and_chunk_examples, get_metadata_types, random_sample_metadata
-from transformers import default_data_collator
 
 
 logger = logging.getLogger(__name__)

bsmetadata/experiments/with_metadata_datasetv2.py

@@ -9,12 +9,12 @@
 from datasets import DatasetDict
 from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
+from transformers import default_data_collator
 
 from bsmetadata.experiments.datasetv2 import get_files, load_dataset_by_files
 from bsmetadata.experiments.without_metadata import preprocess_no_metadata
 from bsmetadata.metadata_processors import PROCESSORS
 from bsmetadata.metadata_utils import add_metadata_and_chunk_examples, random_sample_metadata_v2
-from transformers import default_data_collator
 
 
 logger = logging.getLogger(__name__)
@@ -123,7 +123,7 @@ def remove_num_proc_kwarg(kwargs):
 
             if args.metadata_config.random_sample_metadata_weights is not None:
                 metadata_type_sample_weights = args.metadata_config.random_sample_metadata_weights
-                logger.info(f"using metadata_type_sample_weights proviced in args")
+                logger.info(f"using {metadata_type_sample_weights} proviced in args")
             else:
 
                 def get_metadata_types(example):

bsmetadata/experiments/without_metadata.py

@@ -3,7 +3,6 @@
 
 from datasets import config, load_dataset
 from torch.utils.data import DataLoader
-
 from transformers import default_data_collator
 
 

bsmetadata/metadata_utils.py

@@ -20,9 +20,9 @@
 from typing import Any, DefaultDict, Dict, List, Optional, Tuple
 
 import numpy as np
+from transformers import PreTrainedTokenizerFast
 
 from bsmetadata.metadata_processors import PROCESSORS, MetadataConfig, MetadataProcessor
-from transformers import PreTrainedTokenizerFast
 
 
 logger = logging.getLogger(__name__)

bsmetadata/train.py

@@ -13,6 +13,7 @@
 import hydra
 import torch
 import torch.nn.functional as F
+import wandb
 from accelerate import Accelerator
 from accelerate.utils import DistributedType, DummyOptim, DummyScheduler
 from hydra.core.config_store import ConfigStore
@@ -22,7 +23,6 @@
 from transformers import AddedToken, AutoConfig, AutoModelForCausalLM, AutoTokenizer, get_scheduler, set_seed
 from transformers.trainer_utils import IntervalStrategy
 
-import wandb
 from bsmetadata.input_pipeline import DataConfig, get_dataloaders
 
 
@@ -297,7 +297,10 @@ def main(args: CFG) -> None:
             model, optimizer, dummy_dataloader, scheduler
         )
     else:
-        format_fn = lambda x: x
+
+        def format_fn(x):
+            return x
+
         train_dataloader, eval_dataloaders = get_dataloaders(tokenizer, args.data_config)
 
         # Prepare everything
@@ -409,7 +412,7 @@ def save(path):
     step = 0
     model.train()
     # for epoch in range(args.num_train_epochs):
-    finished = False
+    # finished = False
     if not args.data_config.streaming:
         metrics_logger.log({"train_dataloader_length": len(train_dataloader)})
 
@@ -486,7 +489,7 @@ def get_data_iter():
             evaluate_multiple_dateloaders(eval_dataloaders)
 
         if completed_steps >= args.max_train_steps:
-            finished = True
+            # finished = True
             break
     metrics_logger.close()
     logger.info("Training finished")