Fix failing test - test_no_stored_outputs AND format fix

Author: tarun-etikala

examples/knowledge-tuning/04_Knowledge_Mixing/Knowledge_Mixing.ipynb

@@ -160,9 +160,11 @@
     "\n",
     "    # Filter out problematic questions\n",
     "    ds = ds.filter(\n",
-    "        lambda x: \"...\" not in x[\"question\"]\n",
-    "        and \"<question>\" not in x[\"question\"]\n",
-    "        and \"<Insert question here>\" not in x[\"question\"]\n",
+    "        lambda x: (\n",
+    "            \"...\" not in x[\"question\"]\n",
+    "            and \"<question>\" not in x[\"question\"]\n",
+    "            and \"<Insert question here>\" not in x[\"question\"]\n",
+    "        )\n",
     "    )\n",
     "\n",
     "    # Clean response text\n",

examples/knowledge-tuning/04_Knowledge_Mixing/utils/knowledge_utils.py

@@ -100,7 +100,8 @@ def sample_doc_qa(
 def _clean_response_text(df: pl.DataFrame) -> pl.DataFrame:
     """Clean response text by removing markers and whitespace."""
     return df.with_columns(
-        pl.col("response")
+        pl
+        .col("response")
         .str.replace_all(r"\[END\]", "")
         .str.replace_all(r"\[ANSWER\]", "")
         .str.strip_chars()
@@ -111,7 +112,8 @@ def _clean_response_text(df: pl.DataFrame) -> pl.DataFrame:
 def _create_metadata(df: pl.DataFrame) -> pl.Expr:
     """Create metadata JSON structure."""
     return (
-        pl.struct([
+        pl
+        .struct([
             pl.col("document").alias("sdg_document"),
             pl.lit("document_knowledge_qa").alias("dataset"),
             pl.col("raw_document"),
@@ -232,7 +234,8 @@ def generate_knowledge_qa_dataset(
             "reasoning",
         ]
         messages_expr = (
-            pl.struct(message_columns)
+            pl
+            .struct(message_columns)
             .map_elements(_create_messages_with_reasoning_no_document)
             .alias("messages")
         )
@@ -245,21 +248,24 @@ def generate_knowledge_qa_dataset(
             "reasoning",
         ]
         messages_expr = (
-            pl.struct(message_columns)
+            pl
+            .struct(message_columns)
             .map_elements(_create_messages_with_reasoning)
             .alias("messages")
         )
     elif keep_document_in_context:
         message_columns = ["question", "response", "document", "document_outline"]
         messages_expr = (
-            pl.struct(message_columns)
+            pl
+            .struct(message_columns)
             .map_elements(_create_messages_without_reasoning)
             .alias("messages")
         )
     else:
         message_columns = ["question", "response", "document", "document_outline"]
         messages_expr = (
-            pl.struct(message_columns)
+            pl
+            .struct(message_columns)
             .map_elements(_create_messages_without_reasoning_no_document)
             .alias("messages")
         )
@@ -307,7 +313,8 @@ def count_tokens(text: str) -> int:
         return len(tokenizer.encode(text))
 
     return df.with_columns(
-        pl.col(column_name)
+        pl
+        .col(column_name)
         .map_elements(apply_chat_template, return_dtype=pl.String)
         .map_elements(count_tokens, return_dtype=pl.Int32)
         .alias("token_length")

tests/validation/test_notebook_content.py

@@ -43,13 +43,22 @@ def test_no_execution_counts(notebook_path, relative_path):
 
 
 def test_no_stored_outputs(notebook_path, relative_path):
-    """Test that notebooks have no stored outputs (should be cleared)."""
+    """Test that notebooks have no stored outputs (should be cleared).
+
+    Cells with 'keep_output' tag in metadata are ignored.
+    """
     with open(notebook_path, encoding="utf-8") as f:
         nb = json.load(f)
 
     cells_with_outputs = []
     for i, cell in enumerate(nb.get("cells", [])):
         if cell.get("cell_type") == "code":
+            # Check if cell has keep_output tag
+            metadata = cell.get("metadata", {})
+            tags = metadata.get("tags", [])
+            if "keep_output" in tags:
+                continue
+
             outputs = cell.get("outputs", [])
             if len(outputs) > 0:
                 cells_with_outputs.append((i, len(outputs)))