Update main.py

main.py

@@ -28,36 +28,59 @@ app.add_middleware(
 
 def train_the_model(data):
     try:     
+
         new_data = data
         encoders = load('transexpress_encoders.joblib')
         xgb_model = load('transexpress_xgb_model.joblib')
+        
+        # Selecting and filling missing data
         selected_columns = ['customer_name', 'customer_address', 'customer_phone_no',
-                            'weight','cod','pickup_address','client_number','destination_city',
+                            '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:
-                    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])
+                    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']
-
-        X_train, X_test, y_train, y_test = train_test_split(X_new,y_new, test_size=0.2, random_state=42)
         
-        xgb_model.fit(X_new, y_new)
-        dump(xgb_model,'transexpress_xgb_model.joblib')
+        # Splitting data into training and testing sets
+        X_train, X_test, y_train, y_test = train_test_split(X_new, y_new, test_size=0.2, random_state=42)
         
-
-        y_pred = xgb_model.predict(X_test)
+        # Setting up parameter grid for hyperparameter tuning
+        param_grid = {
+            'max_depth': [3, 4, 5],
+            'learning_rate': [0.01, 0.1, 0.4],
+            'n_estimators': [100, 200, 300],
+            'subsample': [0.8, 0.9, 1],
+            'colsample_bytree': [0.3, 0.7]
+        }
+        
+        # Initializing GridSearchCV
+        grid_search = GridSearchCV(estimator=xgb_model, param_grid, cv=50, 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')
+        
+        # Making predictions and evaluating the model
+        y_pred = best_model.predict(X_test)
         accuracy = accuracy_score(y_test, y_pred)
         classification_rep = classification_report(y_test, y_pred)
-        return accuracy,classification_rep,"Model finetuned with new data."
+        
+        # Returning the results
+        return accuracy, classification_rep, "Model finetuned with new data."
+        
 
         
     except:
@@ -94,7 +117,7 @@ def train_the_model(data):
         xgb = XGBClassifier(use_label_encoder=False, eval_metric='logloss')
         
         # Setup GridSearchCV
-        grid_search = GridSearchCV(xgb, param_grid, cv=40, n_jobs=-1, scoring='accuracy')
+        grid_search = GridSearchCV(xgb, param_grid, cv=50, n_jobs=-1, scoring='accuracy')
         
         # Fit the grid search to the data
         grid_search.fit(X_train, y_train)