Merge pull request #21 from atlasia-ma/fix/suppress-warnings-dataloader-workers

fix: suppress transformers warnings, add dataloader workers config

Author: hBouanane

src/darija_translator/__init__.py

@@ -0,0 +1,7 @@
+import warnings
+
+warnings.filterwarnings("ignore",
+                        message=".*is_flash_linear_attention_available.*")
+import logging
+
+logging.getLogger("transformers").setLevel(logging.ERROR)

src/darija_translator/cli.py

@@ -13,7 +13,9 @@
 from darija_translator.evaluate import compute_translation_metrics, generate_translations
 from darija_translator.model import attach_lora, load_model_and_tokenizer
 from darija_translator.train import build_trainer, save_model
-# from dotenv import load_dotenv
+from dotenv import load_dotenv
+
+load_dotenv()
 
 
 def prepare_data(dataset_name: str,

src/darija_translator/config.py

@@ -32,7 +32,8 @@ class ModelConfig:
 
 @dataclass(frozen=True)
 class TrainConfig:
-    per_device_train_batch_size: int = 16
+    dataloader_num_workers: int = 4
+    per_device_train_batch_size: int = 64
     gradient_accumulation_steps: int = 2
     per_device_eval_batch_size: int = 8
     num_train_epochs: int = 3

src/darija_translator/train.py

@@ -15,7 +15,7 @@ def prepare_data(dataset_name: str,
                  data_config: DataConfig,
                  tokenizer,
                  remove_columns: bool = True) -> tuple:
-    dataset = load_dataset(dataset_name, split="train[:10]")
+    dataset = load_dataset(dataset_name, split="train")
     # dataset = dataset.filter(is_darija_script)
     # dataset = dataset.map(lambda b: to_conversations(b, data_config),
     #                       batched=True)
@@ -36,6 +36,7 @@ def build_trainer(model, tokenizer, train_dataset, eval_dataset,
 
     sft_args = SFTConfig(
         dataset_text_field="text",
+        dataloader_num_workers=config.dataloader_num_workers,
         per_device_train_batch_size=config.per_device_train_batch_size,
         gradient_accumulation_steps=config.gradient_accumulation_steps,
         packing=False,
@@ -52,6 +53,7 @@ def build_trainer(model, tokenizer, train_dataset, eval_dataset,
         seed=config.seed,
         report_to=config.report_to,
         output_dir=config.output_dir,
+        padding_free=False,
         save_strategy="steps",  # Save checkpoints at step intervals
         save_steps=400,
         save_total_limit=3,