change cli and README

README.md

@@ -62,12 +62,12 @@ import pandas as pd
 from rektgbm import RektDataset, RektGBM, RektOptimizer
 
 # Prepare your datasets
-X_train = pd.read_csv("train.csv")
-X_test = pd.read_csv("test.csv")
-y_train = X_train.pop("target")
+train = pd.read_csv("train.csv")
+test = pd.read_csv("test.csv")
+y_train = train.pop("target")
 
-dtrain = RektDataset(data=X_train, label=y_train)
-dtest = RektDataset(data=X_test, reference=dtrain)
+dtrain = RektDataset(data=train, label=y_train)
+dtest = RektDataset(data=test, reference=dtrain)
 
 # Initialize RektOptimizer to automatically detect task type, objective, and metric
 rekt_optimizer = RektOptimizer()

rektgbm/cli.py

@@ -63,6 +63,10 @@ def main(
         raise ValueError(
             f"The specified target column '{target}' does not exist in the training data."
         )
+
+    if target in test_data.columns:
+        test_data.pop(target)
+
     train_label = train_data.pop(target)
     dtrain = RektDataset(data=train_data, label=train_label)
     dtest = RektDataset(data=test_data, reference=dtrain)