Merge pull request #15 from atlasia-ma/feat/wandb-tracking

Feat/wandb tracking

Author: hBouanane

.env.example

@@ -0,0 +1 @@
+WANDB_API_KEY=your-key-here

pyproject.toml

@@ -17,6 +17,8 @@ train = [
     "datasets>=4.3.0",
     "trl>=0.24.0",
     "unsloth>=2026.6.9",
+    "wandb>=0.19",
+    "python-dotenv>=1.0",
 ]
 
 [project.scripts]

src/darija_translator/cli.py

@@ -13,6 +13,7 @@
 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
 
 
 def prepare_data(dataset_name: str, data_config: DataConfig, tokenizer):

src/darija_translator/config.py

@@ -47,3 +47,5 @@ class TrainConfig:
     group_by_length: bool = True
     max_seq_length: int = 2048
     output_dir: str = "lora_model"
+    report_to: str = "wandb"
+    wandb_project: str = "darija-translator"

src/darija_translator/train.py

@@ -1,3 +1,5 @@
+import os
+
 from trl import SFTConfig, SFTTrainer
 from unsloth.chat_templates import train_on_responses_only
 
@@ -6,6 +8,9 @@
 
 def build_trainer(model, tokenizer, train_dataset, eval_dataset,
                   config: TrainConfig):
+    if config.report_to == "wandb" or (isinstance(config.report_to, list) and "wandb" in config.report_to):
+        os.environ["WANDB_PROJECT"] = config.wandb_project
+
     trainer = SFTTrainer(
         model=model,
         tokenizer=tokenizer,
@@ -27,7 +32,7 @@ def build_trainer(model, tokenizer, train_dataset, eval_dataset,
             weight_decay=config.weight_decay,
             lr_scheduler_type=config.lr_scheduler_type,
             seed=config.seed,
-            report_to="none",
+            report_to=config.report_to,
             group_by_length=config.group_by_length,
         ),
     )

tests/test_config.py

@@ -58,3 +58,9 @@ def test_train_config_defaults():
     assert cfg.num_train_epochs == 3
     assert cfg.learning_rate == 2e-4
     assert cfg.seed == 3407
+
+
+def test_train_config_wandb_defaults():
+    cfg = TrainConfig()
+    assert cfg.report_to == "wandb"
+    assert cfg.wandb_project == "darija-translator"