Source file ovo.ml

(**************************************************************************)
(*                                                                        *)
(*  This file is part of CAISAR.                                          *)
(*                                                                        *)
(*  Copyright (C) 2024                                                    *)
(*    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).            *)
(*                                                                        *)
(**************************************************************************)

open Base
module Format = Stdlib.Format
module Sys = Stdlib.Sys
module Filename = Stdlib.Filename
module Fun = Stdlib.Fun

type t = {
  n_inputs : int;
  n_outputs : int;
}

(* OVO format handling. *)

let ovo_format_error s =
  Error (Format.sprintf "OVO format error: %s condition not satisfied." s)

(* Parse a single OVO format line: split line wrt CSV format, and convert each
   string into a number by means of converter [f]. *)
let handle_ovo_line ~f in_channel =
  List.filter_map
    ~f:(fun s -> try Some (f (String.strip s)) with _ -> None)
    (Csv.next in_channel)

(* Handle ovo first line: either 'ovo' or 'ovo x y' with 'x' and 'y' positive
   integer numbers. *)
let handle_ovo_first_line in_channel =
  let ovo_format_error_on_first_line = ovo_format_error "first line" in
  match Csv.next in_channel with
  | [ first_line ] -> (
    let in_channel = Csv.of_string ~separator:' ' first_line in
    match Csv.next in_channel with
    | [ "ovo"; n_is; n_os ] ->
      let n_is = Int.of_string (String.strip n_is) in
      let n_os = Int.of_string (String.strip n_os) in
      Ok (Some (n_is, n_os))
    | [ "ovo" ] -> Ok None
    | _ -> ovo_format_error_on_first_line
    | exception End_of_file -> ovo_format_error_on_first_line)
  | _ -> ovo_format_error_on_first_line
  | exception End_of_file -> ovo_format_error_on_first_line

(* Retrieve inputs and outputs size. *)
let handle_ovo_basic_info ~descr in_channel =
  match handle_ovo_line ~f:Int.of_string in_channel with
  | [ dim ] -> Ok dim
  | _ -> ovo_format_error descr
  | exception End_of_file -> ovo_format_error descr

(* Retrieves [filename] OVO model metadata and weights wrt OVO format
   specification, which is described here:
   https://github.com/abstract-machine-learning/saver#classifier-format. In
   practice, the first line may specify the input/output size values or not. In
   the latter case, we search for input/output size values in the second/third
   lines respectively. *)
let parse_in_channel in_channel =
  let open Result in
  try
    let in_channel = Csv.of_channel in_channel in
    (handle_ovo_first_line in_channel >>= function
     | Some (n_is, n_os) -> Ok (n_is, n_os)
     | None ->
       handle_ovo_basic_info ~descr:"input size" in_channel >>= fun n_is ->
       handle_ovo_basic_info ~descr:"output size" in_channel >>= fun n_os ->
       Ok (n_is, n_os))
    >>= fun (n_is, n_os) ->
    Csv.close_in in_channel;
    Ok { n_inputs = n_is; n_outputs = n_os }
  with
  | Csv.Failure (_nrecord, _nfield, msg) -> Error msg
  | Sys_error s -> Error s
  | Failure msg -> Error (Format.sprintf "Unexpected error: %s." msg)

let parse filename =
  let in_channel = Stdlib.open_in filename in
  Fun.protect
    ~finally:(fun () -> Stdlib.close_in in_channel)
    (fun () -> parse_in_channel in_channel)