Update DSIP/preprocess.py

DSIP/preprocess.py

@@ -1,7 +1,43 @@
 import argparse
+import pandas as pd
+from sklearn.model_selection import train_test_split
+import os
+
+
 
 def parse(csv_path):
     print(f"Location of the file: {csv_path}")
+    
+    # Step 1: Load the dataset
+    # file_path = "dataset.csv"  # Path to the original dataset
+    data = pd.read_csv(csv_path)
+    
+    # Step 2: Define the feature columns (X) and target column (y)
+    X = data[["name", "attendance percentage", "average sleep time", "average screen time"]]  # Feature columns
+    y = data["grade"]  # Target column
+    
+    # Step 3: Split the dataset into training and testing sets
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+    
+    # Step 4: Combine X and y back into dataframes for train and test
+    train_data = pd.concat([X_train, y_train], axis=1)  # Combine features and target for training data
+    test_data = pd.concat([X_test, y_test], axis=1)    # Combine features and target for testing data
+    
+    # Step 5: Create the 'data' folder if it doesn't exist
+    output_folder = "data"
+    os.makedirs(output_folder, exist_ok=True)
+    
+    # Step 6: Save the train and test sets as CSV files
+    train_file_path = os.path.join(output_folder, "train.csv")
+    test_file_path = os.path.join(output_folder, "test.csv")
+    
+    train_data.to_csv(train_file_path, index=False)
+    test_data.to_csv(test_file_path, index=False)
+    
+    print(f"Train and test datasets saved in '{output_folder}' folder.")
+    
+
+
 
 
 if __name__ == '__main__':
@@ -10,4 +46,7 @@ if __name__ == '__main__':
     
     
     args = parser.parse_args()
-    parse(args.csv_path)
+    parse(args.csv_path)
+
+
+