🥅 Disallow empty train streams

Signed-off-by: Evaline Ju <[email protected]>
flaviabeo · Nov 2, 2023 · 414fa1e · 414fa1e
1 parent e512728
commit 414fa1e
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/caikit_nlp/modules/text_generation/peft_prompt_tuning.py b/caikit_nlp/modules/text_generation/peft_prompt_tuning.py
@@ -345,6 +345,7 @@ def train(
  PeftPromptTuning
  Instance of this class with tuned prompt vectors.
  """
+ error.value_check("<NLP46653367E>", len(train_stream) > 0, "train_stream cannot be empty")
 
  # Configure random seed
  transformers.set_seed(seed)
@@ -388,6 +389,8 @@ def train(
 
  train_stream = train_stream.map(convert_to_generation_record)
  if val_stream:
+ error.value_check("<NLP63201425E>", len(val_stream) > 0, "val_stream cannot be empty")
+
  val_stream = val_stream.map(convert_to_generation_record)
 
  # Convert our datastreams -> data loaders by disguising them as PyTorch iterable datasets

diff --git a/caikit_nlp/modules/text_generation/text_generation_local.py b/caikit_nlp/modules/text_generation/text_generation_local.py
@@ -224,6 +224,7 @@ def train(
  TextGeneration
  Instance of this class with fine-tuned models.
  """
+ error.value_check("<NLP96406893E>", len(train_stream) > 0, "train_stream cannot be empty")
 
  torch_dtype = get_torch_dtype(torch_dtype)