diff --git a/explicability/shap_vals.py b/explicability/shap_vals.py
index f390653609ae1492313a4c443f1db6b00f4211ec..5e5c1214fcd3d31faa438f356ed34d0dcf5e8ba6 100644
--- a/explicability/shap_vals.py
+++ b/explicability/shap_vals.py
@@ -15,7 +15,7 @@ from sklearn.tree import DecisionTreeClassifier
 
 # Reading test and training data
 # --------------------------------------------------------------------------------------------------------
-def read_data():
+def read_data(attribute_names):
     # Load test data
     X_test_pre = np.load('../gen_train_data/data/output/pre/X_test_pre.npy', allow_pickle=True)
     y_test_pre = np.load('../gen_train_data/data/output/pre/y_test_pre.npy', allow_pickle=True)
@@ -40,25 +40,25 @@ def read_data():
     X_train_under_post = np.load('../gen_train_data/data/output/post/X_train_under_post.npy', allow_pickle=True)
     y_train_under_post = np.load('../gen_train_data/data/output/post/y_train_under_post.npy', allow_pickle=True)
 
+    # Type conversion needed    
     data_dic = {
-        "X_test_pre": X_test_pre,
+        "X_test_pre": pd.DataFrame(X_test_pre, columns=attribute_names).convert_dtypes(),
         "y_test_pre": y_test_pre,
-        "X_test_post": X_test_post,
+        "X_test_post": pd.DataFrame(X_test_post, columns=attribute_names).convert_dtypes(),
         "y_test_post": y_test_post,
-        "X_train_pre": X_train_pre,
+        "X_train_pre": pd.DataFrame(X_train_pre, columns=attribute_names).convert_dtypes(),
         "y_train_pre": y_train_pre,
-        "X_train_post": X_train_post,
+        "X_train_post": pd.DataFrame(X_train_post, columns=attribute_names).convert_dtypes(),
         "y_train_post": y_train_post,
-        "X_train_over_pre": X_train_over_pre,
+        "X_train_over_pre": pd.DataFrame(X_train_over_pre, columns=attribute_names).convert_dtypes(),
         "y_train_over_pre": y_train_over_pre,
-        "X_train_over_post": X_train_over_post,
+        "X_train_over_post": pd.DataFrame(X_train_over_post, columns=attribute_names).convert_dtypes(),
         "y_train_over_post": y_train_over_post,
-        "X_train_under_pre": X_train_under_pre,
+        "X_train_under_pre": pd.DataFrame(X_train_under_pre, columns=attribute_names).convert_dtypes(),
         "y_train_under_pre": y_train_under_pre,
-        "X_train_under_post": X_train_under_post,
+        "X_train_under_post": pd.DataFrame(X_train_under_post, columns=attribute_names).convert_dtypes(),
         "y_train_under_post": y_train_under_post,
     }
-
     return data_dic
 # --------------------------------------------------------------------------------------------------------
 
@@ -120,49 +120,47 @@ if __name__ == "__main__":
 
     # Setup
     # --------------------------------------------------------------------------------------------------------
+    # Retrieve attribute names in order
+    attribute_names = list(np.load('../gen_train_data/data/output/attributes.npy', allow_pickle=True))
     # Reading data
-    data_dic = read_data()
+    data_dic = read_data(attribute_names)
     method_names = {
         0: "ORIG",
         1: "ORIG_CW",
         2: "OVER",
         3: "UNDER"
     }
-    
     model_choices = {
         "ORIG": "XGB",
         "ORIG_CW": "RF",
         "OVER": "XGB",
         "UNDER": "XGB"
     }
-    # Retrieve attribute names in order
-    attribute_names = list(np.load('../gen_train_data/data/output/attributes.npy', allow_pickle=True))
     # --------------------------------------------------------------------------------------------------------
 
     # Shap value generation
     # --------------------------------------------------------------------------------------------------------
     for i, group in enumerate(['pre', 'post']):
         # Get test dataset based on group, add column names
-        X_test = pd.DataFrame(data_dic['X_test_' + group], columns=attribute_names)
+        X_test = data_dic['X_test_' + group]
         y_test = data_dic['y_test_' + group]
         for j, method in enumerate(['', '', 'over_', 'under_']):
             print(f"{group}-{method_names[j]}")
             # Get train dataset based on group and method
-            X_train = pd.DataFrame(data_dic['X_train_' + method + group], columns=attribute_names)
+            X_train = data_dic['X_train_' + method + group]
             y_train = data_dic['y_train_' + method + group]
             method_name = method_names[j]
             # Get chosen tuned model for this group and method context
             model, is_tree = get_chosen_model(group_str=group, method_str=method_name, model_name=model_choices[method_name])
             # --------------------------------------------------------------------------------------------------------
-            # Fit model with training data
-            fitted_model = model.fit(X_train[:500], y_train[:500])
+            fitted_model = model.fit(X_train[:50], y_train[:50])
             # # Check if we are dealing with a tree vs nn model
             if is_tree:
                  explainer = shap.TreeExplainer(fitted_model)
             # else:
             #     explainer = shap.KernelExplainer(fitted_model.predict_proba, X_test[:500])
             # Compute shap values
-            shap_vals = explainer.shap_values(X_test[:500], check_additivity=False) # Change to true for final results
+            shap_vals = explainer.shap_values(X_test[:50], check_additivity=False) # Change to true for final results
             # ---------------------------------------------------------------------------------------------------------
             # Save results
             np.save(f"./output/shap_values/{group}_{method_names[j]}", shap_vals)