Merge pull request #272 from relyt0925/issue-240

Handle empty dataset from output of sdg leaf node without raising error

Author: aakankshaduggal

src/instructlab/sdg/generate_data.py

@@ -25,7 +25,6 @@
 from instructlab.sdg.pipeline import (
     FULL_PIPELINES_PACKAGE,
     SIMPLE_PIPELINES_PACKAGE,
-    EmptyDatasetError,
     Pipeline,
     PipelineContext,
 )
@@ -359,6 +358,7 @@ def generate_data(
         )
 
     generated_data = None
+    empty_sdg_leaf_nodes = []
     for leaf_node in leaf_nodes.values():
         is_knowledge = False
         leaf_node_path = leaf_node[0]["taxonomy_path"].replace("->", "_")
@@ -382,9 +382,9 @@ def generate_data(
         logger.debug("Dataset: %s", ds)
         new_generated_data = pipe.generate(ds, leaf_node_path)
         if len(new_generated_data) == 0:
-            raise EmptyDatasetError(
-                "Pipeline stopped: Empty dataset after running pipe"
-            )
+            empty_sdg_leaf_nodes.append(leaf_node_path)
+            logger.warning("Empty dataset for qna node: %s", leaf_node_path)
+            continue
         generated_data = (
             [new_generated_data]
             if generated_data is None
@@ -418,3 +418,9 @@ def generate_data(
 
     generate_duration = time.time() - generate_start
     logger.info(f"Generation took {generate_duration:.2f}s")
+    if len(empty_sdg_leaf_nodes) > 0:
+        logger.warning(
+            "Leaf nodes with empty sdg output: {}".format(
+                " ".join(empty_sdg_leaf_nodes)
+            )
+        )

tests/test_generate_data.py

@@ -9,6 +9,7 @@
 import glob
 import json
 import os
+import re
 import shutil
 import tempfile
 import unittest
@@ -267,6 +268,11 @@ def strip_q(q):
     return output
 
 
+def _empty_llmblock_generate(self, samples):
+    """Return an empty set of generated samples."""
+    return []
+
+
 @patch.object(LLMBlock, "_generate", _noop_llmblock_generate)
 class TestGenerateCompositionalData(unittest.TestCase):
     @pytest.fixture(autouse=True)
@@ -452,6 +458,61 @@ def __exit__(self, *args):
         self.teardown()
 
 
+@patch.object(LLMBlock, "_generate", _empty_llmblock_generate)
+class TestGenerateEmptyDataset(unittest.TestCase):
+    @pytest.fixture(autouse=True)
+    def _init_taxonomy(self, taxonomy_dir):
+        self.test_taxonomy = taxonomy_dir
+
+    def setUp(self):
+        self.tmp_path = tempfile.TemporaryDirectory().name
+        test_valid_knowledge_skill_file = os.path.join(
+            TEST_DATA_DIR, "test_valid_knowledge_skill.yaml"
+        )
+        tracked_knowledge_file = os.path.join("knowledge  ", "tracked", "qna.yaml")
+        untracked_knowledge_file = os.path.join("knowledge", "new", "qna.yaml")
+        test_valid_knowledge_skill = load_test_skills(test_valid_knowledge_skill_file)
+        self.test_taxonomy.add_tracked(
+            tracked_knowledge_file, test_valid_knowledge_skill
+        )
+        self.test_taxonomy.create_untracked(
+            untracked_knowledge_file, test_valid_knowledge_skill
+        )
+
+    def test_generate(self):
+        with patch("logging.Logger.info") as mocked_logger:
+            generate_data(
+                client=MagicMock(),
+                logger=mocked_logger,
+                model_family="merlinite",
+                model_name="models/merlinite-7b-lab-Q4_K_M.gguf",
+                num_instructions_to_generate=10,
+                taxonomy=self.test_taxonomy.root,
+                taxonomy_base=TEST_TAXONOMY_BASE,
+                output_dir=self.tmp_path,
+                chunk_word_count=1000,
+                server_ctx_size=4096,
+                pipeline="simple",
+            )
+        mocked_logger.warning.assert_called()
+        assert re.search(
+            "empty sdg output: knowledge_new", mocked_logger.warning.call_args.args[0]
+        )
+
+    def teardown(self) -> None:
+        """Recursively remove the temporary repository and all of its
+        subdirectories and files.
+        """
+        shutil.rmtree(self.tmp_path)
+        return
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.teardown()
+
+
 def test_context_init_batch_size_optional():
     """Test that the _context_init function can handle a missing batch size by
     delegating to the default in PipelineContext.