diff --git a/etl/epc/Dataset.py b/etl/epc/Dataset.py
index e897da78..7d5c3ef8 100644
--- a/etl/epc/Dataset.py
+++ b/etl/epc/Dataset.py
@@ -203,11 +203,11 @@ class TrainingDataset(BaseDataset):
         common_cols = [[col + "_starting", col + "_ending"] for col in common_cols]
 
         self.df = self.df.loc[
-                  :,
-                  no_suffix_cols
-                  + only_ending_cols
-                  + [col for cols in common_cols for col in cols],
-                  ]
+            :,
+            no_suffix_cols
+            + only_ending_cols
+            + [col for cols in common_cols for col in cols],
+        ]
 
     def _remove_abnormal_change_in_floor_area(self):
         """
@@ -509,7 +509,7 @@ class TrainingDataset(BaseDataset):
                     expanded_df["is_sandstone_or_limestone"]
                     == expanded_df["is_sandstone_or_limestone_ending"]
                 )
-                ]
+            ]
         elif component == "floor":
             expanded_df = expanded_df[
                 (expanded_df["is_suspended"] == expanded_df["is_suspended_ending"])
@@ -526,7 +526,7 @@ class TrainingDataset(BaseDataset):
                     expanded_df["is_to_external_air"]
                     == expanded_df["is_to_external_air_ending"]
                 )
-                ]
+            ]
         elif component == "roof":
             expanded_df = expanded_df[
                 (expanded_df["is_pitched"] == expanded_df["is_pitched_ending"])
@@ -539,7 +539,7 @@ class TrainingDataset(BaseDataset):
                     expanded_df["has_dwelling_above"]
                     == expanded_df["has_dwelling_above_ending"]
                 )
-                ]
+            ]
 
         return expanded_df
 
@@ -567,13 +567,12 @@ class TrainingDataset(BaseDataset):
                 "is_system_built_ending",
                 "is_timber_frame_ending",
                 "is_granite_or_whinstone_ending",
-                "is_as_built_ending",
+                # "is_as_built_ending",
                 "is_cob_ending",
-                "is_assumed_ending",
                 "is_sandstone_or_limestone_ending",
                 # Re remove the is_assumed columns
-                "is_assumed",
-                "is_assumed_ending",
+                # "is_assumed",
+                # "is_assumed_ending",
             ],
             "floor": [
                 "original_description",
@@ -698,6 +697,8 @@ class TrainingDataset(BaseDataset):
             # Rename columns to component specific names, if they have not been dropped
             expanded_df = expanded_df.rename(
                 columns={
+                    "is_assumed": f"{component}_is_assumed",
+                    "is_assumed_ending": f"{component}_is_assumed_ending",
                     "insulation_thickness": f"{component}_insulation_thickness",
                     "insulation_thickness_ending": f"{component}_insulation_thickness_ending",
                     "thermal_transmittance": f"{component}_thermal_transmittance",