(* Simple neural network example *)

type activation = Relu | Sigmoid | Tanh | Silu | Softmax [@@deriving show]

let parse_activation = function 
    "relu" -> Relu | "sigmoid" -> Sigmoid | "tanh" -> Tanh | "silu" -> Silu | "softmax" -> Softmax
    | s -> raise (Invalid_argument s)

let show_activation = function
    | Relu -> "relu" | Sigmoid -> "sigmoid" | Tanh -> "tanh" | Silu -> "silu" | Softmax -> "softmax"


type t = {
    (* [@short] for short form of the argument *)
    nlayers        : int   [@short "-l"]
        (** the number of layers *);

    (* with [@set_false], "-use-dropout" flag turns the value false. *)
    use_dropout    : bool  [@set_false]
        (** true if use dropout *);

    dropout_rate   : float option
        (** dropout rate *);

    test : int list option
        (** this is an test argument *);

    test2 : int option list;

    (* arguments with custom type *)
    activation : activation [@print show_activation] [@parse parse_activation];
        (** activation function in feed forward layers *)

    activation2 : activation [@parse parse_activation];
        (** activation function in feed forward layers *)
} [@@deriving show, argparse {
    (* You can pass positional argument, which is (string * string) list.  *)
    positional = 
        ["train", "train file";
        "eval", "some file"];

    (* also some description. *)
    description = "some neural networks";
}]

(*
   You can omit either of positional and description 
   as in [@@deriving argparse { description = "some ..." }]
   or entirely: [@@deriving argparse].
*)


(* default arguments *)
let default = {
    nlayers = 2;
    use_dropout = true;
    dropout_rate = Some 0.1;
    test = Some [1;2;3];
    test2 = [Some 1;Some 2;Some 3; None];
    activation = Relu;
    activation2 = Sigmoid;
}

(*
   [@@argparse] generates the following functions:
   val argparse : t -> string -> string array -> t * string array
   val prerr_argparse : string -> t -> unit

   while for types whose name is different from t (say config):
   val argparse_config : config -> string -> string array -> config * string array
   val prerr_config_argparse : string -> config -> unit
*)

let () =
    let cfg, rest = argparse default "example" Sys.argv in
    prerr_argparse "example" cfg;
    print_endline (show cfg);
    Array.iter print_endline rest

(* print_endline (show cfg) *)
{ Example.src_vocab_size = 3; tgt_vocab_size = 100; num_units = 512;
  nheads = 8; nlayers = 6; use_dropout = false; dropout_rate = (Some 0.1);
  test = (Some [1; 1; 1]); test2 = [(Some 1); (Some 2); (Some 3); None];
  activation = Example.Relu; activation2 = Example.Sigmoid } 
  
(* Array.iter print_endline rest *)
train.txt
test.txt

[@@argparse { positional = ["files ...", "a list of files to process"] }]