Update main.py

main.py

@@ -37,6 +37,12 @@ def train_the_model(data,page):
         selected_columns = ['customer_name', 'customer_address', 'customer_phone_no',
                             'weight', 'cod', 'pickup_address', 'client_number', 'destination_city',
                             'status_name']
+
+
+
+
+
+        
         new_data_filled = new_data[selected_columns].fillna('Missing')
         
         # Encoding categorical data
@@ -64,17 +70,15 @@ def train_the_model(data,page):
         }
         
         # Initializing GridSearchCV
-        grid_search = GridSearchCV(xgb_model, param_grid, cv=10, n_jobs=-1, scoring='accuracy')
+        grid_search = GridSearchCV(xgb_model, param_grid, cv=2, n_jobs=-1, scoring='accuracy')
         
         # Fitting GridSearchCV
         grid_search.fit(X_train, y_train)
         
-        # Updating the model with the best estimator
-        best_model = grid_search.best_estimator_
-        dump(best_model, 'transexpress_xgb_model.joblib')
+        dump(grid_search, 'transexpress_xgb_model.joblib')
         
         # Making predictions and evaluating the model
-        y_pred = best_model.predict(X_test)
+        y_pred = grid_search.predict(X_test)
         accuracy = accuracy_score(y_test, y_pred)
         classification_rep = classification_report(y_test, y_pred)
         
@@ -117,7 +121,7 @@ def train_the_model(data,page):
         xgb = XGBClassifier(use_label_encoder=False, eval_metric='logloss')
         
         # Setup GridSearchCV
-        grid_search = GridSearchCV(xgb, param_grid, cv=10, n_jobs=-1, scoring='accuracy')
+        grid_search = GridSearchCV(xgb, param_grid, cv=2, n_jobs=-1, scoring='accuracy')
         
         # Fit the grid search to the data
         grid_search.fit(X_train, y_train)