update readme and test

Author: zzsi

README.md

@@ -1,17 +1,25 @@
-`gbt` is a library for gradient boosted trees with minimal coding required.
+`gbt` is a library for gradient boosted trees with minimal coding required. It is a thin wrapper around [`lightgbm`](https://lightgbm.readthedocs.io/en/v3.3.2/). Give it a `pandas.Dataframe`, `gbt.train()` takes care of feature transforms (e.g. scaling for numerical features, label encoding for categorical features) and metrics print outs.
 
 Example usage:
 
 ```
 import pandas as pd
 import gbt
 
-df = pd.read_csv("my_data.csv")
-gbt.train(
+df = pd.DataFrame({
+    "a": [1, 2, 3, 4, 5, 6, 7],
+    "b": ["a", "b", "c", None, "e", "f", "g"],
+    "c": [1, 0, 1, 1, 0, 0, 1],
+    "some_other_column": [0, 0, None, None, None, 3, 3]
+})
+train(
     df,
-    target_column=...,
+    recipe="binary",
+    label_column="c",
     val_size=0.2,  # fraction of the validation split
-    categorical_feature_columns=[],
-    numerical_feature_columns=[],
+    categorical_feature_columns=["b"],
+    numerical_feature_columns=["a"],
 )
 ```
+
+Supported "recipes": mape, l2, l2_rf, binary, multiclass.

gbt/api.py

@@ -56,11 +56,11 @@ def train(
     df_test=None,
     recipe="l2",
     num_classes=None,
+    label_column=None,
     categorical_feature_columns=None,
     numerical_feature_columns=None,
     preprocess_fn=None,
     sort_by_columns=None,
-    label_column=None,
     add_categorical_stats=False,
     pretrain_size=0,
     val_size=0.1,
@@ -100,7 +100,9 @@ def train(
     )
     ds.df = df
     ds.preprocess()
-    assert ds.features.shape[0] > 10
+    if ds.features.shape[0] <= 10:
+        import warnings
+        warnings.warn(f"Too few samples: {ds.features.shape[0]}. Training may not converge.")
     train_ds, val_ds = ds.split(
         pretrain_size=pretrain_size, val_size=val_size, shuffle=False
     )
@@ -163,6 +165,8 @@ def train(
             "lambda_l2": 0.001,
             "num_class": num_classes,
         }
+    else:
+        raise ValueError(f"Unknown recipe: {recipe}. Supported: mape, l2, l2_rf, binary, multiclass")
 
     print(parameters)
     model = LightGBMModel(parameters=parameters, rounds=100)

tests/test_gbt.py

@@ -1,7 +1,9 @@
 from os import path
 
+import pandas as pd
 from gbt.api import train
 
+
 csd = path.dirname(path.realpath(__file__))
 
 
@@ -13,3 +15,20 @@ def test_model_can_train():
         sort_by_columns=None,
         label_column="label",
     )
+
+
+def test_the_readme_example():
+    df = pd.DataFrame({
+        "a": [1, 2, 3, 4, 5, 6, 7],
+        "b": ["a", "b", "c", None, "e", "f", "g"],
+        "c": [1, 0, 1, 1, 0, 0, 1],
+        "some_other_column": [0, 0, None, None, None, 3, 3]
+    })
+    train(
+        df,
+        recipe="binary",
+        label_column="c",
+        val_size=0.2,  # fraction of the validation split
+        categorical_feature_columns=["b"],
+        numerical_feature_columns=["a"],
+    )