refactor

delphi-suite · Apr 5, 2024 · fa1f52e · fa1f52e
1 parent 7405f72
commit fa1f52e
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/scripts/tokenize_dataset.py b/scripts/tokenize_dataset.py
@@ -57,17 +57,18 @@
         text_docs = input_dataset[args.column_name]
     else:
         if len(input_dataset.column_names) > 1:
-            raise ValueError("There are more than one column in the specified dataset")
+            raise ValueError("There is more than one column in the specified dataset")
         text_docs = input_dataset[input_dataset.column_names[0]]
 
+    tokenized_dataset = tokenize_dataset(
+        text_docs,
+        tokenizer,
+        context_size=args.context_size,
+        batch_size=args.batch_size,
+    )
     output_dataset = Dataset.from_dict(
         {
-            "tokens": tokenize_dataset(
-                text_docs,
-                tokenizer,
-                context_size=args.context_size,
-                batch_size=args.batch_size,
-            )
+            "tokens": tokenized_dataset,
         }
     )