Fixes : {TheAlgorithms#6551}

Author: Moddy2024

machine_learning/xgboostregressor.py

@@ -5,6 +5,21 @@
 from xgboost import XGBRegressor
 
 
+def dataset(datatype:dict) -> tuple:
+    # Split dataset into train and test data
+    x = (datatype["data"],datatype["target"])  # features
+    return x
+
+
+def xgboost(features:list, target:list,test_features:list) -> list:
+    xgb = XGBRegressor()
+    xgb.fit(features, target)
+    # Predict target for test data
+    predictions = xgb.predict(test_features)
+    predictions = predictions.reshape(len(predictions), 1)
+    return predictions
+
+
 def main() -> None:
 
     """
@@ -15,25 +30,21 @@ def main() -> None:
     # Load Boston house price dataset
     boston = load_boston()
     print(boston.keys())
-
-    # Split dataset into train and test data
-    x = boston["data"]  # features
-    y = boston["target"]
+    
+    features, target = dataset(boston)
     x_train, x_test, y_train, y_test = train_test_split(
-        x, y, test_size=0.25, random_state=1
+        features, target, test_size=0.25, random_state=1
     )
-    # XGBoost Regressor
-    xgb = XGBRegressor()
-    xgb.fit(x_train, y_train)
-
-    # Predict target for test data
-    predictions = xgb.predict(x_test)
-    predictions = predictions.reshape(len(predictions), 1)
+    predictions = xgboost(x_train, y_train, x_test)
 
     # Error printing
     print(f"Mean Absolute Error:\t {mean_absolute_error(y_test, predictions)}")
     print(f"Mean Square Error  :\t {mean_squared_error(y_test, predictions)}")
 
 
 if __name__ == "__main__":
+    import doctest
+    doctest.testmod(name="main", verbose=True)
+    doctest.testmod(name="dataset", verbose=True)
+    doctest.testmod(name="xgboost", verbose=True)
     main()