Merge pull request #120 from metarank/feature/catboost-options

more options to the catboost booster
metarank · Feb 9, 2023 · f61ebd1 · f61ebd1
2 parents 26b2640 + 38ed15a
commit f61ebd1
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 5 deletions.
diff --git a/src/main/scala/io/github/metarank/ltrlib/booster/CatboostBooster.scala b/src/main/scala/io/github/metarank/ltrlib/booster/CatboostBooster.scala
@@ -64,13 +64,15 @@ object CatboostBooster extends BoosterFactory[String, CatboostBooster, CatboostO
     val modelFile = dir.createChild("model.bin")
     val opts = Map(
       "--learn-set"     -> dataset,
-      "--loss-function" -> "QueryRMSE",
+      "--loss-function" -> options.objective,
       "--eval-metric"   -> s"NDCG:top=${options.ndcgCutoff}",
       "--iterations"    -> options.trees.toString,
       "--depth"         -> options.maxDepth.toString,
       "--learning-rate" -> options.learningRate.toString,
       "--train-dir"     -> dir.toString(),
-      "--model-file"    -> modelFile.toString()
+      "--model-file"    -> modelFile.toString(),
+      "--logging-level" -> "Silent",
+      "--random-seed"   -> options.randomSeed.toString
     ) ++ test.map(t => Map("--test-set" -> t)).getOrElse(Map.empty)
     native_impl.ModeFitImpl(new TVector_TString(opts.flatMap(kv => List(kv._1, kv._2)).toArray))
     val bytes = modelFile.byteArray

diff --git a/src/main/scala/io/github/metarank/ltrlib/booster/CatboostOptions.scala b/src/main/scala/io/github/metarank/ltrlib/booster/CatboostOptions.scala
@@ -9,5 +9,7 @@ case class CatboostOptions(
     learningRate: Double = 0.1,
     ndcgCutoff: Int = 10,
     maxDepth: Int = 8,
-    randomSeed: Int = Random.nextInt()
+    randomSeed: Int = math.abs(Random.nextInt()),
+    objective: String = "QueryRMSE",
+    loggingLevel: String = "Verbose"
 ) extends BoosterOptions
diff --git a/src/main/scala/io/github/metarank/ltrlib/booster/LightGBMOptions.scala b/src/main/scala/io/github/metarank/ltrlib/booster/LightGBMOptions.scala
@@ -8,7 +8,7 @@ case class LightGBMOptions(
     learningRate: Double = 0.1,
     ndcgCutoff: Int = 10,
     maxDepth: Int = 8,
-    randomSeed: Int = Random.nextInt(),
+    randomSeed: Int = math.abs(Random.nextInt()),
     numLeaves: Int = 16,
     featureFraction: Double = 1.0
 ) extends BoosterOptions
diff --git a/src/main/scala/io/github/metarank/ltrlib/booster/XGBoostOptions.scala b/src/main/scala/io/github/metarank/ltrlib/booster/XGBoostOptions.scala
@@ -9,6 +9,6 @@ case class XGBoostOptions(
     learningRate: Double = 0.1,
     ndcgCutoff: Int = 10,
     maxDepth: Int = 8,
-    randomSeed: Int = Random.nextInt(),
+    randomSeed: Int = math.abs(Random.nextInt()),
     subsample: Double = 1.0
 ) extends BoosterOptions