add sample_rate to accelrate training

README.md

@@ -43,7 +43,7 @@ The template can be configured using the following parameters:
 | Deploy to HuggingFace | Whether to deploy to HuggingFace Hub | False |
 | Deploy to SkyPilot | Whether to deploy to SkyPilot | False |
 | Dataset | The dataset to use from HuggingFace Datasets | airline_reviews |
-| Model | The model to use from HuggingFace Models | roberta-base |
+| Model | The model to use from HuggingFace Models | distilbert-base-uncased |
 | Cloud Provider | The cloud provider to use (AWS or GCP) | aws |
 | Metric-Based Promotion | Whether to promote models based on metrics | True |
 | Notifications on Failure | Whether to notify about pipeline failures | True |

copier.yml

@@ -65,6 +65,10 @@ accelerator:
         - gpu
         - cpu
     default: gpu
+sample_rate:
+    type: bool
+    help: "Whether to use a sample of the dataset for quick iteration"
+    default: False
 deploy_locally:
     type: bool
     help: "Whether to deploy locally"
@@ -91,8 +95,8 @@ model:
     choices:
         - bert-base-uncased
         - roberta-base
-        - distilbert-base-cased
-    default: roberta-base
+        - distilbert-base-uncased
+    default: distilbert-base-uncased
 cloud_of_choice:
     type: str
     help: "Whether to use AWS cloud provider or GCP"

template/steps/dataset_loader/data_loader.py

@@ -4,6 +4,9 @@
 from datasets import load_dataset, DatasetDict
 from zenml import step
 from zenml.logger import get_logger
+{%- if sample_rate %}
+import numpy as np
+{%- endif %}
 
 logger = get_logger(__name__)
 
@@ -41,6 +44,19 @@ def data_loader(
     dataset = dataset.remove_columns(["airline_sentiment_confidence","negativereason_confidence"])
     {%- endif %}
 
+    {%- if sample_rate %}
+    # Sample 20% of the data randomly for the demo
+    def sample_dataset(dataset, sample_rate=0.2):
+        sampled_dataset = DatasetDict()
+        for split in dataset.keys():
+            split_size = len(dataset[split])
+            indices = np.random.choice(split_size, int(split_size * sample_rate), replace=False)
+            sampled_dataset[split] = dataset[split].select(indices)
+        return sampled_dataset
+
+    dataset = sample_dataset(dataset)
+    {%- endif %}
+
     # Log the dataset and sample examples
     logger.info(dataset)
     logger.info(f"Sample Example 1 : {dataset['train'][0]['text']} with label {dataset['train'][0]['label']}")

template/steps/training/model_trainer.py

@@ -105,7 +105,7 @@ def model_trainer(
         evaluation_strategy='steps',
         save_strategy='steps',
         save_steps=1000,
-        eval_steps=200,
+        eval_steps=100,
         logging_steps=logging_steps,
         save_total_limit=5,
         report_to="mlflow",