Repository: oumi-ai/oumi
Pull request: Add text_col param that's required for SFTTrainer
- Replaced the usage of `data.trainer_kwargs["dataset_text_field"]` with `data.text_col` since it's a required field for SFTTrainer
- Changed the default TrainerType to HF. IMO this is more logical as a default value. More practically, I made `text_col` a required field for SFTTrainer, making it not possible to do `TrainingConfig()` in tests.

Problem statement:
trainer_kwargs: avoid listing in config.data
The `trainer_kwargs` seems better fitting to go under training and not data. There is a subtle param listed there `dataset_text_field` which is applicable also on the inference side since it affects the loaders of the data.

Hints from issue discussion:
I've been wondering about this as well. What about putting dataset_text_field under DataParams? That feels more logical to me.

Commit message: Add text_col param that's required for SFTTrainer

Patch to review:
diff --git a/.vscode/launch.json b/.vscode/launch.json
index 669b90841..5af7366cc 100644
--- a/.vscode/launch.json
+++ b/.vscode/launch.json
@@ -9,10 +9,12 @@
             "args": [
                 "data.dataset_name=yahma/alpaca-cleaned",
                 "data.preprocessing_function_name=alpaca",
-                "data.trainer_kwargs.dataset_text_field=prompt",
+                "data.text_col=prompt",
                 "model.model_name=openai-community/gpt2",
                 "model.trust_remote_code=true",
-                "training.logging_steps=10",
+                "training.trainer_type=TRL_SFT",
+                "training.max_steps=3",
+                "training.logging_steps=3",
                 "training.enable_wandb=false",
                 "training.enable_tensorboard=false",
                 "training.output_dir=tmp"
diff --git a/configs/lema/gpt2.pt.yaml b/configs/lema/gpt2.pt.yaml
index 0aa8dbc72..25536318a 100644
--- a/configs/lema/gpt2.pt.yaml
+++ b/configs/lema/gpt2.pt.yaml
@@ -10,10 +10,10 @@ data:
   dataset_config: "wikitext-2-raw-v1"
   stream: True
   pack: True
-  trainer_kwargs:
-    dataset_text_field: "text"
+  text_col: "text"
 
 training:
+  trainer_type: TRL_SFT
   max_steps: 2
   enable_tensorboard: False
   output_dir: "output/gpt2.pt"
diff --git a/configs/lema/phi3.lora.yaml b/configs/lema/phi3.lora.yaml
index 6694b2267..967c01d50 100644
--- a/configs/lema/phi3.lora.yaml
+++ b/configs/lema/phi3.lora.yaml
@@ -5,12 +5,12 @@ model:
 data:
   dataset_name: "yahma/alpaca-cleaned"
   preprocessing_function_name: "alpaca"
-  trainer_kwargs:
-    dataset_text_field: "prompt"
+  text_col: "prompt"
 
 training:
   optimizer: "adamw_torch"
   use_peft: true
+  trainer_type: TRL_SFT
   output_dir: "output/phi3.lora"
 
 peft:
diff --git a/configs/lema/phi3.sft.nvidia.24g.yaml b/configs/lema/phi3.sft.nvidia.24g.yaml
index 7e87c1306..21cdfa35f 100644
--- a/configs/lema/phi3.sft.nvidia.24g.yaml
+++ b/configs/lema/phi3.sft.nvidia.24g.yaml
@@ -6,8 +6,7 @@ model:
 data:
   dataset_name: "yahma/alpaca-cleaned"
   preprocessing_function_name: "alpaca"
-  trainer_kwargs:
-    dataset_text_field: "prompt"
+  text_col: "prompt"
 
 training:
   optimizer: "adamw_torch"
diff --git a/configs/lema/zephyr.7b/sft/config_full.yaml b/configs/lema/zephyr.7b/sft/config_full.yaml
index 020ed3216..b09c79b88 100644
--- a/configs/lema/zephyr.7b/sft/config_full.yaml
+++ b/configs/lema/zephyr.7b/sft/config_full.yaml
@@ -15,12 +15,11 @@ data:
   split: train_sft
   stream: True
   pack: True
+  text_col: "text"
   preprocessing_function_name: trl_sft_ultrachat_200k
   preprocessing_function_kwargs:
     num_proc: 6
     batched: False
-  trainer_kwargs:
-    dataset_text_field: text
 
 training:
   optimizer: adamw_torch
diff --git a/configs/lema/zephyr.7b/sft/config_qlora.yaml b/configs/lema/zephyr.7b/sft/config_qlora.yaml
index 43278ac18..5252605c8 100644
--- a/configs/lema/zephyr.7b/sft/config_qlora.yaml
+++ b/configs/lema/zephyr.7b/sft/config_qlora.yaml
@@ -15,12 +15,11 @@ data:
   split: train_sft
   stream: True
   pack: True
+  text_col: "text"
   preprocessing_function_name: trl_sft_ultrachat_200k
   preprocessing_function_kwargs:
     num_proc: 6
     batched: false
-  trainer_kwargs:
-    dataset_text_field: text
 
 training:
   optimizer: adamw_torch
diff --git a/notebooks/LeMa - Colab Setup Example.ipynb b/notebooks/LeMa - Colab Setup Example.ipynb
index ced4a397c..3c30618ee 100644
--- a/notebooks/LeMa - Colab Setup Example.ipynb	
+++ b/notebooks/LeMa - Colab Setup Example.ipynb	
@@ -27,12 +27,13 @@
    },
    "source": [
     "#### 1. Setting up read-only github token\n",
-    "Since the Github repository is private, we need to generate a `read-only` user token scoped for the `lema` repo.\n",
-    "1. In Github.com, go to `Settings -> Developer settings -> Personal access tokens -> Generate new token`\n",
-    "2. See example [here](https://drive.google.com/file/d/1zxd8r7qkPfl34mfGK83m_13oLGFGghW1/view?usp=share_link) on how to fill the form. The only permission that should be granted is `Contents`, in `read-only` mode\n",
-    "3. Add the github token to your colab environment secrets (Key icon in the left menu)\n",
+    "Since the Github repository is private, we need to generate a `Read-only` user token scoped for the `lema` repo.\n",
+    "1. In Github.com, go to `Settings -> Developer settings -> Personal access tokens -> Fine-grained tokens -> Generate new token`.\n",
+    "1. See example [here](https://drive.google.com/file/d/1zxd8r7qkPfl34mfGK83m_13oLGFGghW1/view?usp=share_link) on how to fill the form. The only permission that should be granted is `Repository permissions -> Contents -> Read-only`.\n",
+    "1. Click `Generate token`, copy the token, and save it somewhere safe (as you can't access it again).\n",
+    "1. Create a colab environment secret (Key icon in the left menu) with `repo-token` as the name and your token as the value.\n",
     "\n",
-    "This only needs to be done once!"
+    "This only needs to be done once! However, you have to wait for the token request to be approved for this to work."
    ]
   },
   {
@@ -55,7 +56,7 @@
     "from google.colab import userdata\n",
     "\n",
     "github_repo_token = userdata.get(\"repo-token\")  # Setup token in your notebook secrets\n",
-    "github_username = \"<GITHUB_USERNAME>\"  # Change your github username\n",
+    "github_username = \"<GITHUB_USERNAME>\"  # Change to your github username\n",
     "\n",
     "!git clone https://$github_username:$github_repo_token@github.com/openlema/lema.git"
    ]
@@ -77,37 +78,50 @@
    },
    "outputs": [],
    "source": [
-    "%pip install -e lema[all]"
+    "%pip install -e 'lema[all]'"
    ]
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "id": "3ASgNcAx0lZ_"
-   },
+   "metadata": {},
    "source": [
-    "## Training\n",
-    "Make sure to enable GPU runtime for faster training"
+    "#### 4. Importing LeMa"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import lema\n",
+    "from lema.core.types import (\n",
+    "    DataParams,\n",
+    "    EvaluationConfig,\n",
+    "    ModelParams,\n",
+    "    TrainerType,\n",
+    "    TrainingConfig,\n",
+    "    TrainingParams,\n",
+    ")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {
-    "id": "SYSpXqvP0sbT"
+    "id": "3ASgNcAx0lZ_"
    },
    "source": [
-    "#### Using `lema` module"
+    "## Training\n",
+    "Make sure to enable GPU runtime for faster training."
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
+   "cell_type": "markdown",
    "metadata": {
-    "id": "IYBkZmipsQL6"
+    "id": "SYSpXqvP0sbT"
    },
-   "outputs": [],
    "source": [
-    "import lema"
+    "#### Using `lema` module"
    ]
   },
   {
@@ -118,13 +132,22 @@
    },
    "outputs": [],
    "source": [
-    "lema.train(\n",
-    "    model_name=\"microsoft/Phi-3-mini-4k-instruct\",\n",
-    "    dataset_name=\"yahma/alpaca-cleaned\",\n",
-    "    preprocessing_function_name=\"alpaca\",\n",
-    "    output_dir=\"train/\",\n",
-    "    trust_remote_code=True,\n",
-    ")"
+    "config = TrainingConfig(\n",
+    "    data=DataParams(\n",
+    "        dataset_name=\"yahma/alpaca-cleaned\",\n",
+    "        preprocessing_function_name=\"alpaca\",\n",
+    "        text_col=\"prompt\",\n",
+    "    ),\n",
+    "    model=ModelParams(\n",
+    "        model_name=\"microsoft/Phi-3-mini-4k-instruct\",\n",
+    "        trust_remote_code=True,\n",
+    "    ),\n",
+    "    training=TrainingParams(\n",
+    "        trainer_type=TrainerType.TRL_SFT,\n",
+    "        output_dir=\"train/\",\n",
+    "    ),\n",
+    ")\n",
+    "lema.train(config)"
    ]
   },
   {
@@ -147,9 +170,10 @@
     "!lema-train \\\n",
     "    \"data.dataset_name=yahma/alpaca-cleaned\" \\\n",
     "    \"data.preprocessing_function_name=alpaca\" \\\n",
-    "    \"data.trainer_kwargs.dataset_text_field=prompt\" \\\n",
+    "    \"data.text_col=prompt\" \\\n",
     "    \"model.model_name=microsoft/Phi-3-mini-4k-instruct\" \\\n",
     "    \"model.trust_remote_code=true\" \\\n",
+    "    \"training.trainer_type=TRL_SFT/\" \\\n",
     "    \"training.output_dir=train/\""
    ]
   },
@@ -177,13 +201,18 @@
    },
    "outputs": [],
    "source": [
-    "lema.evaluate(\n",
-    "    model_name=\"train/best.pt\",  # model output\n",
-    "    dataset_name=\"yahma/alpaca-cleaned\",\n",
-    "    preprocessing_function_name=\"alpaca\",\n",
-    "    output_dir=\"eval/\",\n",
-    "    trust_remote_code=True,\n",
-    ")"
+    "config = EvaluationConfig(\n",
+    "    data=DataParams(\n",
+    "        dataset_name=\"yahma/alpaca-cleaned\",\n",
+    "        preprocessing_function_name=\"alpaca\",\n",
+    "    ),\n",
+    "    model=ModelParams(\n",
+    "        model_name=\"train/best.pt\",\n",
+    "        trust_remote_code=True,\n",
+    "    ),\n",
+    ")\n",
+    "\n",
+    "lema.evaluate(config)"
    ]
   },
   {
@@ -202,10 +231,9 @@
     "!lema-evaluate \\\n",
     "    \"data.dataset_name=yahma/alpaca-cleaned\" \\\n",
     "    \"data.preprocessing_function_name=alpaca\" \\\n",
-    "    \"data.trainer_kwargs.dataset_text_field=prompt\" \\\n",
+    "    \"data.text_col=prompt\" \\\n",
     "    \"model.model_name=microsoft/Phi-3-mini-4k-instruct\" \\\n",
-    "    \"model.trust_remote_code=true\" \\\n",
-    "    \"training.output_dir=eval/\""
+    "    \"model.trust_remote_code=true\""
    ]
   }
  ],
diff --git a/scripts/finetune_sft_phi3.sh b/scripts/finetune_sft_phi3.sh
index c0bf44fd9..53eb2560f 100755
--- a/scripts/finetune_sft_phi3.sh
+++ b/scripts/finetune_sft_phi3.sh
@@ -5,6 +5,6 @@ python -m lema.train \
     "model.model_name=microsoft/Phi-3-mini-4k-instruct" \
     "data.dataset_name=yahma/alpaca-cleaned" \
     "data.preprocessing_function_name=alpaca" \
-    "data.trainer_kwargs.dataset_text_field=prompt" \
+    "data.text_col=prompt" \
     "training.output_dir=train/" \
     "model.trust_remote_code=true"
diff --git a/src/lema/builders/data.py b/src/lema/builders/data.py
index dbe10961d..b2c02c2aa 100644
--- a/src/lema/builders/data.py
+++ b/src/lema/builders/data.py
@@ -79,12 +79,16 @@ def build_dataset(
         dataset_kwargs = {}
         if config.model.model_max_length:
             dataset_kwargs["seq_length"] = config.model.model_max_length
-
+        # Our preprocessing functions take a dict as input and return a dict as output.
+        # formatting_func must return a str, so we fetch the target str from the dict.
+        if preprocessing_fn:
+            dataset_kwargs["formatting_func"] = lambda x: preprocessing_fn(x)[
+                data_params.text_col
+            ]
         dataset = ConstantLengthDataset(
             tokenizer,
             dataset,
-            dataset_text_field=data_params.get_dataset_text_field(),
-            formatting_func=preprocessing_fn,
+            dataset_text_field=data_params.text_col,
             **dataset_kwargs,
         )
     elif data_params.preprocessing_function_name:
diff --git a/src/lema/core/types.py b/src/lema/core/types.py
index baf2b4e70..0bd7352fa 100644
--- a/src/lema/core/types.py
+++ b/src/lema/core/types.py
@@ -7,7 +7,7 @@
 from omegaconf import MISSING, OmegaConf
 from peft.utils.peft_types import TaskType
 
-_DATASET_TEXT_FIELD = "dataset_text_field"
+from lema.logging import logger
 
 
 #
@@ -30,7 +30,7 @@ class TrainerType(Enum):
 class TrainingParams:
     optimizer: str = "adamw_torch"
     use_peft: bool = False
-    trainer_type: TrainerType = TrainerType.TRL_SFT
+    trainer_type: TrainerType = TrainerType.HF
     enable_gradient_checkpointing: bool = False
     output_dir: str = "output"
     per_device_train_batch_size: int = 8
@@ -137,29 +137,21 @@ def _default_factory_preprocessing_kwargs() -> dict:
         defaults["batched"] = True  # Note the default of huggingface is False.
         return defaults
 
+    # The dataset column name containing the text to train on. Required for SFTTrainer.
+    text_col: Optional[str] = None
     preprocessing_function_name: Optional[str] = None
     preprocessing_function_kwargs: Dict[str, Any] = field(
         default_factory=_default_factory_preprocessing_kwargs
     )
     trainer_kwargs: Dict[str, Any] = field(default_factory=dict)
 
-    def get_dataset_text_field(self) -> Optional[str]:
-        """Get the `dataset_text_field` value if present."""
-        return self.trainer_kwargs.get(_DATASET_TEXT_FIELD)
-
     def __post_init__(self):
-        """Verify params if packing is enabled."""
+        """Verify params."""
         if self.pack:
             if not self.stream:
                 raise ValueError("`stream` must be enabled if `pack` is enabled.")
-            if (
-                not self.preprocessing_function_name
-                and _DATASET_TEXT_FIELD not in self.trainer_kwargs
-            ):
-                raise ValueError(
-                    "Either `trainer_kwargs['dataset_text_field']` "
-                    "or `preprocessing_function_name` must be specified."
-                )
+            if not self.text_col:
+                raise ValueError("`text_col` must be specified if `pack` is enabled.")
 
 
 @dataclass
@@ -269,6 +261,24 @@ class TrainingConfig(BaseConfig):
     training: TrainingParams = field(default_factory=TrainingParams)
     peft: PeftParams = field(default_factory=PeftParams)
 
+    def __post_init__(self):
+        """Verify/populate params."""
+        if self.training.trainer_type == TrainerType.TRL_SFT:
+            if not self.data.text_col:
+                raise ValueError("`text_col` must be specified for TRL_SFT Trainer.")
+            existing_dataset_text_field = self.data.trainer_kwargs.get(
+                "dataset_text_field"
+            )
+            if (
+                existing_dataset_text_field is not None
+                and existing_dataset_text_field != self.data.text_col
+            ):
+                logger.warning(
+                    "Overriding existing `dataset_text_field` value "
+                    f'"{existing_dataset_text_field}" with "{self.data.text_col}"'
+                )
+            self.data.trainer_kwargs["dataset_text_field"] = self.data.text_col
+
 
 @dataclass
 class EvaluationConfig(BaseConfig):
diff --git a/tests/test_evaluate.py b/tests/test_evaluate.py
index c69ef06f1..3683c4480 100644
--- a/tests/test_evaluate.py
+++ b/tests/test_evaluate.py
@@ -9,9 +9,6 @@ def test_basic_evaluate():
         data=DataParams(
             dataset_name="yahma/alpaca-cleaned",
             preprocessing_function_name="alpaca",
-            trainer_kwargs={
-                "dataset_text_field": "prompt",
-            },
         ),
         model=ModelParams(
             model_name="openai-community/gpt2",
diff --git a/tests/test_train.py b/tests/test_train.py
index e9a2719b1..47e57967c 100644
--- a/tests/test_train.py
+++ b/tests/test_train.py
@@ -1,7 +1,13 @@
 import tempfile
 
 from lema import train
-from lema.core.types import DataParams, ModelParams, TrainingConfig, TrainingParams
+from lema.core.types import (
+    DataParams,
+    ModelParams,
+    TrainerType,
+    TrainingConfig,
+    TrainingParams,
+)
 
 
 def test_basic_train():
@@ -11,15 +17,14 @@ def test_basic_train():
         data=DataParams(
             dataset_name="yahma/alpaca-cleaned",
             preprocessing_function_name="alpaca",
-            trainer_kwargs={
-                "dataset_text_field": "prompt",
-            },
+            text_col="prompt",
         ),
         model=ModelParams(
             model_name="openai-community/gpt2",
             trust_remote_code=True,
         ),
         training=TrainingParams(
+            trainer_type=TrainerType.TRL_SFT,
             max_steps=3,
             logging_steps=3,
             enable_wandb=False,
@@ -38,9 +43,7 @@ def test_custom_train():
         data=DataParams(
             dataset_name="yahma/alpaca-cleaned",
             preprocessing_function_name="alpaca",
-            trainer_kwargs={
-                "dataset_text_field": "prompt",
-            },
+            text_col="prompt",
         ),
         model=ModelParams(
             model_name="learning-machines/sample",
@@ -48,6 +51,7 @@ def test_custom_train():
             trust_remote_code=False,
         ),
         training=TrainingParams(
+            trainer_type=TrainerType.TRL_SFT,
             max_steps=3,
             logging_steps=3,
             enable_wandb=False,
@@ -58,3 +62,32 @@ def test_custom_train():
     )
 
     train(config)
+
+
+def test_pack_train():
+    output_temp_dir = tempfile.mkdtemp()
+
+    config: TrainingConfig = TrainingConfig(
+        data=DataParams(
+            dataset_name="Salesforce/wikitext",
+            dataset_config="wikitext-2-raw-v1",
+            stream=True,
+            pack=True,
+            text_col="text",
+        ),
+        model=ModelParams(
+            model_name="openai-community/gpt2",
+            model_max_length=1024,
+            trust_remote_code=True,
+        ),
+        training=TrainingParams(
+            trainer_type=TrainerType.TRL_SFT,
+            max_steps=1,
+            logging_steps=1,
+            enable_wandb=False,
+            enable_tensorboard=False,
+            output_dir=output_temp_dir,
+        ),
+    )
+
+    train(config)


Code review comments:
diff_hunk: @@ -27,12 +27,13 @@
    },
    "source": [
     "#### 1. Setting up read-only github token\n",
-    "Since the Github repository is private, we need to generate a `read-only` user token scoped for the `lema` repo.\n",
-    "1. In Github.com, go to `Settings -> Developer settings -> Personal access tokens -> Generate new token`\n",
-    "2. See example [here](https://drive.google.com/file/d/1zxd8r7qkPfl34mfGK83m_13oLGFGghW1/view?usp=share_link) on how to fill the form. The only permission that should be granted is `Contents`, in `read-only` mode\n",
-    "3. Add the github token to your colab environment secrets (Key icon in the left menu)\n",
+    "Since the Github repository is private, we need to generate a `Read-only` user token scoped for the `lema` repo.\n",
+    "1. In Github.com, go to `Settings -> Developer settings -> Personal access tokens -> Fine-grained tokens -> Generate new token`.\n",
+    "1. See example [here](https://drive.google.com/file/d/1zxd8r7qkPfl34mfGK83m_13oLGFGghW1/view?usp=share_link) on how to fill the form. The only permission that should be granted is `Repository permissions -> Contents -> Read-only`.\n",
+    "1. Click `Generate token`, copy the token, and save it somewhere safe (as you can't access it again).\n",
+    "1. Create a colab environment secret (Key icon in the left menu) with `repo-token` as the name and your token as the value.\n",
     "\n",
-    "This only needs to be done once!"
+    "This only needs to be done once! However, you have to wait for the token request to be approved for this to work."
path: notebooks/LeMa - Colab Setup Example.ipynb
text: @user1:
nit: for some reason I don't get any notifications when someone requests a token, could add a note about slacking me or Nikolai to have that approved quickly?