Source file ovo.ml

(**************************************************************************)
(*                                                                        *)
(*  This file is part of CAISAR.                                          *)
(*                                                                        *)
(*  Copyright (C) 2022                                                    *)
(*    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 = Caml.Format
module Sys = Caml.Sys
module Filename = Caml.Filename
module Fun = Caml.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)

(* Skip the header part, ie comments, of the OVO format. *)
let skip_ovo_header filename in_channel =
  let exception End_of_header in
  let pos_in = ref (Stdlib.pos_in in_channel) in
  try
    while true do
      let line = Stdlib.input_line in_channel in
      if not (Str.string_match (Str.regexp "//") line 0)
      then raise End_of_header
      else pos_in := Stdlib.pos_in in_channel
    done;
    assert false
  with
  | End_of_header ->
    (* At this point we have read one line past the header part: seek back. *)
    Stdlib.seek_in in_channel !pos_in;
    Ok ()
  | End_of_file ->
    Error (Format.sprintf "OVO model not found in file `%s'." filename)

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

(* Skip unused flag. *)
let handle_ovo_unused_flag in_channel =
  try
    let _ = Csv.next in_channel in
    Ok ()
  with End_of_file -> ovo_format_error "second"

(* Retrieves [filename] OVO model metadata and weights wrt OVO format
   specification, which is described here:
   https://github.com/abstract-machine-learning/saver#classifier-format. *)
let parse_in_channel filename in_channel =
  let open Result in
  try
    skip_ovo_header filename in_channel >>= fun () ->
    let in_channel = Csv.of_channel in_channel in
    handle_ovo_unused_flag in_channel >>= fun () ->
    handle_ovo_basic_info in_channel >>= fun n_is ->
    handle_ovo_basic_info in_channel >>= fun 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 filename in_channel)