raise error if token count exceeds 1024 instead of attempting to re-chunk

Signed-off-by: Khaled Sulayman <[email protected]>

Author: khaledsulayman

notebooks/instructlab-knowledge/utils/create_seed_dataset.py

@@ -158,16 +158,15 @@ def add_icls(qna_yaml: Dict[str, str], chunked_document: Dataset) -> Dataset:
                 }
             )
         )
-    chunked_document_all_icl = safe_concatenate_datasets(chunked_document_all_icl)
-    chunked_document_all_icl = chunked_document_all_icl.map(
-        lambda x: {
-            "chunks": chunk_document(
-                [x["document"]], server_ctx_size=4096, chunk_word_count=1024
-            )
-            if get_token_count(x["document"], tokenizer) > 1024
-            else [x["document"]]
-        }
-    )
+    chunked_document_all_icl = []
+    for c in safe_concatenate_datasets(chunked_document_all_icl):
+        if get_token_count(c["document"], tokenizer) > 1024:
+            raise ValueError("Chunk exceeds token count of 1024")
+
+        chunked_document_all_icl.append({
+            "chunks": [c["document"]]
+        })
+
     df = chunked_document_all_icl.to_pandas()
     df_exploded = df.explode("chunks").reset_index(drop=True)
     new_ds = Dataset.from_pandas(df_exploded)