Merge pull request #19 from atlasia-ma/feat/preprocessed-dataset

feat(train): load pre-processed dataset instead of processing at runtime

Author: hBouanane

pyproject.toml

@@ -14,11 +14,11 @@ eval = [
     "sacrebleu>=2.6.0",
 ]
 train = [
-    "datasets>=5.0.0",
-    "trl>=1.7.1",
-    "unsloth>=2026.7.2",
-    "wandb>=0.28.0",
-    "python-dotenv>=1.2.2",
+    "datasets",
+    "trl",
+    "unsloth",
+    "wandb",
+    "python-dotenv",
 ]
 
 [project.scripts]

src/darija_translator/cli.py

@@ -13,17 +13,23 @@
 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
 
 
-def prepare_data(dataset_name: str, data_config: DataConfig, tokenizer):
+def prepare_data(dataset_name: str,
+                 data_config: DataConfig,
+                 tokenizer,
+                 remove_columns: bool = True) -> tuple:
     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)
-    dataset = dataset.map(lambda b: format_conversations(b, tokenizer),
-                          batched=True)
-    dataset = dataset.filter(lambda ex: is_within_length(ex, data_config))
+    # dataset = dataset.filter(is_darija_script)
+    # dataset = dataset.map(lambda b: to_conversations(b, data_config),
+    #                       batched=True)
+    # dataset = dataset.map(lambda b: format_conversations(b, tokenizer),
+    #                       batched=True)
+    # dataset = dataset.filter(lambda ex: is_within_length(ex, data_config))
+    if remove_columns:
+        dataset = dataset.remove_columns(
+            [c for c in dataset.column_names if c != "text"])
     return split_dataset(dataset, data_config)
 
 
@@ -43,7 +49,10 @@ def run_train(args):
 def run_evaluate(args):
     model_config, data_config = ModelConfig(), DataConfig()
     model, tokenizer = load_model_and_tokenizer(model_config)
-    _, eval_dataset = prepare_data(args.dataset, data_config, tokenizer)
+    _, eval_dataset = prepare_data(args.dataset,
+                                   data_config,
+                                   tokenizer,
+                                   remove_columns=False)
     predictions = generate_translations(
         model,
         tokenizer,
@@ -59,13 +68,15 @@ def main():
     subparsers = parser.add_subparsers(required=True)
 
     train_parser = subparsers.add_parser("train")
-    train_parser.add_argument("--dataset",
-                              default="atlasia/darija-english-combined")
+    train_parser.add_argument(
+        "--dataset",
+        default="atlasia/english-to-darija-arabic-script-formatted")
     train_parser.set_defaults(func=run_train)
 
     eval_parser = subparsers.add_parser("evaluate")
-    eval_parser.add_argument("--dataset",
-                             default="atlasia/darija-english-combined")
+    eval_parser.add_argument(
+        "--dataset",
+        default="atlasia/english-to-darija-arabic-script-formatted")
     eval_parser.set_defaults(func=run_evaluate)
 
     args = parser.parse_args()

src/darija_translator/config.py

@@ -36,7 +36,7 @@ class TrainConfig:
     gradient_accumulation_steps: int = 2
     per_device_eval_batch_size: int = 8
     num_train_epochs: int = 3
-    warmup_steps: int = 5
+    warmup_ratio: float = 0.1
     learning_rate: float = 2e-4
     logging_steps: int = 100
     weight_decay: float = 0.01

src/darija_translator/train.py

@@ -23,7 +23,7 @@ def build_trainer(model, tokenizer, train_dataset, eval_dataset,
             gradient_accumulation_steps=config.gradient_accumulation_steps,
             packing=config.packing,
             max_seq_length=config.max_seq_length,
-            warmup_steps=config.warmup_steps,
+            warmup_ratio=config.warmup_ratio,
             num_train_epochs=config.num_train_epochs,
             per_device_eval_batch_size=config.per_device_eval_batch_size,
             eval_strategy="epoch",