Skip to content

Commit

Permalink
Check correct dump format for gblinear. (#10831)
Browse files Browse the repository at this point in the history
  • Loading branch information
trivialfis authored Sep 20, 2024
1 parent 24241ed commit 2a37a88
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 10 deletions.
4 changes: 3 additions & 1 deletion src/gbm/gblinear_model.h
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,7 @@ class GBLinearModel : public Model {
}
}
fo << std::endl << " ]" << std::endl << " }";
} else {
} else if (format == "text") {
fo << "bias:\n";
for (int gid = 0; gid < ngroup; ++gid) {
fo << this->Bias()[gid] << std::endl;
Expand All @@ -136,6 +136,8 @@ class GBLinearModel : public Model {
fo << (*this)[i][gid] << std::endl;
}
}
} else {
LOG(FATAL) << "Dump format `" << format << "` is not supported by the gblinear model.";
}
std::vector<std::string> v;
v.push_back(fo.str());
Expand Down
26 changes: 17 additions & 9 deletions tests/cpp/gbm/test_gblinear.cc
Original file line number Diff line number Diff line change
@@ -1,21 +1,18 @@
/*!
* Copyright 2019 by Contributors
/**
* Copyright 2019-2024, XGBoost Contributors
*/
#include <gtest/gtest.h>
#include <xgboost/feature_map.h> // for FeatureMap

#include <memory>
#include <sstream>

#include "../helpers.h"
#include "xgboost/context.h"
#include "xgboost/gbm.h"
#include "xgboost/json.h"
#include "xgboost/learner.h"
#include "xgboost/logging.h"

namespace xgboost {
namespace gbm {

namespace xgboost::gbm {
TEST(GBLinear, JsonIO) {
size_t constexpr kRows = 16, kCols = 16;

Expand All @@ -40,5 +37,16 @@ TEST(GBLinear, JsonIO) {
ASSERT_EQ(weights.size(), 17);
}
}
} // namespace gbm
} // namespace xgboost

TEST(GBLinear, Dump) {
Context ctx;
size_t constexpr kRows = 16, kCols = 16;
LearnerModelParam mparam{MakeMP(kCols, .5, 1)};

std::unique_ptr<GradientBooster> gbm{
CreateTrainedGBM("gblinear", Args{}, kRows, kCols, &mparam, &ctx)};
FeatureMap fmap;
ASSERT_THAT([&] { [[maybe_unused]] auto vec = gbm->DumpModel(fmap, true, "dot"); },
GMockThrow(R"(`dot` is not supported)"));
}
} // namespace xgboost::gbm

0 comments on commit 2a37a88

Please sign in to comment.