Source file parser.ml

(**************************************************************************)
(*                                                                        *)
(*  This file is part of CAISAR.                                          *)
(*                                                                        *)
(*  Copyright (C) 2025                                                    *)
(*    CEA (Commissariat à l'énergie atomique et aux énergies              *)
(*         alternatives)                                                  *)
(*                                                                        *)
(*  You can redistribute it and/or modify it under the terms of the GNU   *)
(*  Lesser General Public License as published by the Free Software       *)
(*  Foundation, version 2.1.                                              *)
(*                                                                        *)
(*  It is distributed in the hope that it will be useful,                 *)
(*  but WITHOUT ANY WARRANTY; without even the implied warranty of        *)
(*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the          *)
(*  GNU Lesser General Public License for more details.                   *)
(*                                                                        *)
(*  See the GNU Lesser General Public License version 2.1                 *)
(*  for more details (enclosed in the file licenses/LGPLv2.1).            *)
(*                                                                        *)
(**************************************************************************)

(** From https://github.com/dmlc/xgboost/raw/master/doc/model.schema *)

(* {2 Utils } *)

let assoc l (json : Yojson.Safe.t) : _ Ppx_deriving_yojson_runtime.error_or =
  match json with
  | `Assoc json' -> (
    match List.assoc "name" json' with
    | exception Not_found -> Error "Gradient_booster name not found"
    | `String name -> (
      match List.assoc name l with
      | exception Not_found -> Error ("Unknown gradient_booster " ^ name)
      | conv -> conv json json')
    | s -> Error ("Unknown gradient_booster " ^ Yojson.Safe.to_string s))
  | _ -> Error "Gradient_booster name not found"

type int_option = int option [@@deriving show]

let int_option_none = Int32.of_string "2147483647"

let int_option_of_yojson (json : Yojson.Safe.t) =
  match [%derive.of_yojson: int32] json with
  | Ok i ->
    if Int32.equal i int_option_none
    then Ok None
    else Ok (Some (Int32.to_int i))
  | Error _ as error -> error

let int_option_to_yojson i =
  let r = Option.value (Option.map Int32.of_int i) ~default:int_option_none in
  [%derive.to_yojson: int32] r

(* {2 Definitions } *)

type gbtree_model_param = {
  num_trees : string;
  num_parallel_tree : string; [@default ""]
  size_leaf_vector : string;
}

and tree_param = {
  num_nodes : string;
  size_leaf_vector : string;
  num_feature : string;
}

and reg_loss_param = { scale_pos_weight : string [@default ""] }
and pseudo_huber_param = { huber_slope : string [@default ""] }

and aft_loss_param = {
  aft_loss_distribution : string; [@default ""]
  aft_loss_distribution_scale : string; [@default ""]
}

and softmax_multiclass_param = { num_class : string }

and lambda_rank_param = {
  num_pairsample : string;
  fix_list_weight : string;
}

and tree = {
  tree_param : tree_param;
  id : int;
  loss_changes : float array;
  sum_hessian : float array;
  base_weights : float array;
  left_children : int array;
  right_children : int array;
  parents : int_option array;
  split_indices : int array;
  split_conditions : float array;
  split_type : int array;
  default_left : bool array;
  categories : int array;
  categories_nodes : int array;
  categories_segments : int array;
  categories_sizes : int array;
}

and gbtree = {
  gbtree_model_param : gbtree_model_param;
  trees : tree array;
  tree_info : int array;
}

and gblinear = { weights : float array }

and dart = {
  gbtree : gbtree;
  weight_drop : float array;
}

and learner_model_param = {
  base_score : string; [@default ""]
  num_class : string; [@default ""]
  num_feature : string; [@default ""]
}
[@@deriving show, yojson { strict = false }]

type gradient_booster =
  | Gbtree of gbtree
  | Gblinear of gblinear
  | Dart of dart
[@@deriving show]

let gradient_booster_of_yojson :
  Yojson.Safe.t -> gradient_booster Ppx_deriving_yojson_runtime.error_or =
  assoc
    [
      ( "gbtree",
        fun _ json ->
          Ppx_deriving_yojson_runtime.( >|= )
            (gbtree_of_yojson (List.assoc "model" json))
            (fun t -> Gbtree t) );
      ( "gblinear",
        fun _ json ->
          Ppx_deriving_yojson_runtime.( >|= )
            (gblinear_of_yojson (List.assoc "model" json))
            (fun t -> Gblinear t) );
      ( "dart",
        fun json _ ->
          Ppx_deriving_yojson_runtime.( >|= ) (dart_of_yojson json) (fun t ->
            Dart t) );
    ]

let gradient_booster_to_yojson = function
  | Gbtree t ->
    `Assoc [ ("name", `String "gbtree"); ("model", gbtree_to_yojson t) ]
  | Gblinear t ->
    `Assoc [ ("name", `String "gblinear"); ("model", gblinear_to_yojson t) ]
  | Dart t -> (
    match dart_to_yojson t with
    | `Assoc l -> `Assoc (("name", `String "dart") :: l)
    | _ -> assert false)

type objective =
  | Reg_squarederror of reg_loss_param
  | Reg_pseudohubererror of reg_loss_param
  | Reg_squaredlogerror of reg_loss_param
  | Reg_linear of reg_loss_param
  | Binary_logistic of reg_loss_param
(* TODO: many are missing *)
[@@deriving show]

let objective_of_yojson :
  Yojson.Safe.t -> objective Ppx_deriving_yojson_runtime.error_or =
  let map f _ json =
    Ppx_deriving_yojson_runtime.( >|= )
      (reg_loss_param_of_yojson (List.assoc "reg_loss_param" json))
      f
  in
  assoc
    [
      ("reg:squarederror", map @@ fun t -> Reg_squarederror t);
      ("reg:pseudohubererror", map @@ fun t -> Reg_pseudohubererror t);
      ("reg:squaredlogerror", map @@ fun t -> Reg_squaredlogerror t);
      ("reg:linear", map @@ fun t -> Reg_linear t);
      ("binary:logistic", map @@ fun t -> Binary_logistic t);
    ]

let objective_to_yojson t =
  let mk name t =
    `Assoc
      [ ("name", `String name); ("reg_loss_param", reg_loss_param_to_yojson t) ]
  in
  match t with
  | Reg_squarederror t -> mk "reg:squarederror" t
  | Reg_pseudohubererror t -> mk "reg:pseudohuberror" t
  | Reg_squaredlogerror t -> mk "reg:squarelogerror" t
  | Reg_linear t -> mk "reg:linear" t
  | Binary_logistic t -> mk "binary:logistic" t

type learner = {
  feature_names : string array; [@default [||]]
  feature_types : string array; [@default [||]]
  gradient_booster : gradient_booster;
  objective : objective;
  learner_model_param : learner_model_param;
}

and t = {
  version : int * int * int;
  learner : learner;
}
[@@deriving show, yojson { strict = false }]