Update main.py

main.py

@@ -38,29 +38,22 @@ def train_the_model(data,page):
                             'weight', 'cod', 'pickup_address', 'client_number', 'destination_city',
                             'status_name']
 
-
-
-
-
-        
         new_data_filled = new_data[selected_columns].fillna('Missing')
-        
-        # Encoding categorical data
         for col, encoder in encoders.items():
             if col in new_data_filled.columns:
                 unseen_categories = set(new_data_filled[col]) - set(encoder.classes_)
                 if unseen_categories:
-                    encoder.classes_ = np.append(encoder.classes_, unseen_categories)
-                new_data_filled[col] = encoder.transform(new_data_filled[col])
-        
-        # Splitting data into features and target
-        X_new = new_data_filled.drop('status_name', axis=1)
-        y_new = new_data_filled['status_name']
-        
-        # Splitting data into training and testing sets
+                    for category in unseen_categories:
+                        encoder.classes_ = np.append(encoder.classes_, category)
+                    new_data_filled[col] = encoder.transform(new_data_filled[col])
+                else:
+                    new_data_filled[col] = encoder.transform(new_data_filled[col])
+        X_new = new_data_filled.drop('status.name', axis=1)
+        y_new = new_data_filled['status.name']
         X_train, X_test, y_train, y_test = train_test_split(X_new, y_new, test_size=0.2, random_state=42)
-        
-        # Setting up parameter grid for hyperparameter tuning
+
+
+        # Setup the hyperparameter grid to search
         param_grid = {
             'max_depth': [3, 4, 5],
             'learning_rate': [0.01, 0.1, 0.4],
@@ -69,10 +62,13 @@ def train_the_model(data,page):
             'colsample_bytree': [0.3, 0.7]
         }
         
-        # Initializing GridSearchCV
+        # Initialize the classifier
+        #xgb = XGBClassifier(use_label_encoder=False, eval_metric='logloss')
+        
+        # Setup GridSearchCV
         grid_search = GridSearchCV(xgb_model, param_grid, cv=2, n_jobs=-1, scoring='accuracy')
         
-        # Fitting GridSearchCV
+        # Fit the grid search to the data
         grid_search.fit(X_train, y_train)
         
         dump(grid_search, 'transexpress_xgb_model.joblib')