Added prediction feature of year

1 files changed, 8 insertions, 4 deletions

diff --git a/main.py b/main.py
index 3355b34..a8bdec8 100644
--- a/main.py
+++ b/main.py
@@ -17,6 +17,9 @@ def load_data(file_path):
 
 # Preprocessing
 def preprocess_data(df):
+    # Save the 'Year' column before calculating car age
+    df['Original_Year'] = df['Year']
+
     # Calculate Car Age
     df['Car_Age'] = 2024 - df['Year']
     df.drop(columns=['Year'], inplace=True)
@@ -42,7 +45,7 @@ def preprocess_data(df):
 # Train Model
 def train_model(df, target, model_name):
     # Features and target
-    X = df.drop(columns=['Mileage Km/L', 'Price', 'Name'])
+    X = df.drop(columns=['Mileage Km/L', 'Price', 'Name', 'Original_Year'])
     y = df[target]
 
     # Split data
@@ -52,9 +55,7 @@ def train_model(df, target, model_name):
     categorical_cols = ['Fuel_Type', 'Transmission', 'Owner_Type', 'Location']
     numerical_cols = ['Kilometers_Driven', 'Engine CC', 'Power', 'Seats', 'Car_Age']
 
-    preprocessor = ColumnTransformer(
-        transformers=[
-            ('num', StandardScaler(), numerical_cols),
+    preprocessor = ColumnTransformer( transformers=[ ('num', StandardScaler(), numerical_cols),
             ('cat', OneHotEncoder(drop='first'), categorical_cols)
         ]
     )
@@ -106,5 +107,8 @@ def main():
     print("Training price prediction model...")
     train_model(df, target='Price', model_name='price_predictor')
 
+    print("Training year prediction model...")
+    train_model(df, target='Original_Year', model_name='year_predictor')
+
 if __name__ == "__main__":
     main()