From afa8a5e741f7fcd2d2c5e96a1830d1855f4e6b89 Mon Sep 17 00:00:00 2001
From: Matt Seddon <37993418+mattseddon@users.noreply.github.com>
Date: Wed, 18 Oct 2023 06:30:07 +1100
Subject: [PATCH] Update DVCLive xgboost snippet (#4847)

---
 extension/snippets/dvclive.code-snippets | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/extension/snippets/dvclive.code-snippets b/extension/snippets/dvclive.code-snippets
index 5e98f822f5..7aae73bd53 100644
--- a/extension/snippets/dvclive.code-snippets
+++ b/extension/snippets/dvclive.code-snippets
@@ -100,9 +100,17 @@
     "body": [
       "from dvclive.xgb import DVCLiveCallback",
       "",
-      "xgboost.train(",
-      "    ${1:param}, ${2:dtrain}, num_round=${3:5}, evals=[${4:(dval, \"eval_data\")}]",
-      "    callbacks=[DVCLiveCallback(${5:\"eval_data\"}, save_dvc_exp=True)],",
+      "model = xgb.XGBClassifier(",
+      "    n_estimators=${1:100},",
+      "    early_stopping_rounds=${2:5},",
+      "    eval_metric=[${3:\"merror\", \"mlogloss\"}],",
+      "    callbacks=[DVCLiveCallback()]",
+      ")",
+      "",
+      "model.fit(",
+      "    X_train,",
+      "    y_train,",
+      "    eval_set=[(X_test, y_test)]",
       ")"
     ],
     "description": "DVCLive XGBoost callback"