diff options
-rw-r--r-- | main.py | 12 |
1 files changed, 8 insertions, 4 deletions
@@ -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() |