diff --git a/src/home.py b/src/home.py
index 754dc65..5d45647 100644
--- a/src/home.py
+++ b/src/home.py
@@ -1,6 +1,6 @@
 import streamlit as st
 from io import StringIO
-from ydata_profiling import ProfileReport
+# from ydata_profiling import ProfileReport
 import pandas as pd
 
 def  statistics(df):
@@ -44,8 +44,8 @@ def main():
         st.write("## Statistics")
         statistics(df)
 
-        profile = ProfileReport(df, title='Pandas Profiling Report', explorative=True)
-        profile.to_widgets()
+        # profile = ProfileReport(df, title='Pandas Profiling Report', explorative=True)
+        # profile.to_widgets()
 
         if st.button("Next"):
             st.switch_page("pages/clean.py")
diff --git a/src/pages/prediction.py b/src/pages/prediction.py
index 6892c69..6a87c01 100644
--- a/src/pages/prediction.py
+++ b/src/pages/prediction.py
@@ -6,8 +6,18 @@ sys.path.append('./back/')
 import clustering_csv as cc
 import prediction as p
 
-if 'df' in st.session_state:
+def display_prediction_results(df, targetCol):
+    df_cols.remove(col)
+    original_col = df[col]
+    predicted_col = p.getColumnsForPredictionAndPredict(df, df_cols, "Route Type", "Linear Regression")    
+    
+    new_df = pd.DataFrame()
+    new_df['Original'] = original_col
+    new_df['Predicted'] = predicted_col
+
+    st.dataframe(new_df)    
 
+if 'df' in st.session_state:
     df = st.session_state.df
     df_cols = df.columns.tolist()
 
@@ -21,21 +31,10 @@ if 'df' in st.session_state:
 
     if st.button("Linear Regression"):
         col = "Route Type"
-        df_cols.remove(col)
-        original_col = df[col]
-        predicted_col = p.getColumnsForPredictionAndPredict(df, df_cols, "Route Type", "Linear Regression")
+        display_prediction_results(df, col)
 
     if st.button("Random Forest"):
         col = "Route Type"
-        df_cols.remove(col)
-        original_col = df[col]
-        predicted_col = p.getColumnsForPredictionAndPredict(df, df_cols, "Route Type", "Random Forest")
-
-    ndf = pd.DataFrame()
-    ndf['Original'] = original_col
-    ndf['Predicted'] = predicted_col
-
-    st.dataframe(ndf)
-
+        display_prediction_results(df, col)
 else:
     st.write("Please clean your dataset.")