Fixedfile path to written dataset, add option to save to csv

category_classification/datasets/datasets_common.py

@@ -51,7 +51,10 @@ def train_test_split(X, y, test_size=0.25):
     return sk_train_test_split(X, y, test_size=test_size, stratify=y)
 
 
-def write_dataset(dest_dir, X, y, filename):
+def write_dataset(dest_dir, X, y, filename, to_json : bool = True):
     dest_dir = Path(dest_dir)
     df = pd.concat((X, y), axis=1)
-    df.to_json(filename, orient="records", lines=True)
+    if to_json:
+        df.to_json(dest_dir / filename, orient="records", lines=True)
+    else:
+        df.to_csv(dest_dir / filename)