<commit_before>open Core_kernel.Std

type t =
  { session : Wrapper.Session_with_graph.t
  ; graph : Wrapper.Graph.t
  ; nodes : Wrapper.Graph.operation Node.Id.Table.t
  }

let create () =
  let graph = Wrapper.Graph.create () in
  match Wrapper.Session_with_graph.create graph with
  | Error status ->
    failwithf "Unable to generate session: %s" (Wrapper.Status.message status) ()
  | Ok session ->
    { session
    ; graph
    ; nodes = Node.Id.Table.create ()
    }

let default_session = lazy (create ())

let add_attribute operation_description ~attr_name attr =
  match (attr : Node.attr) with
  | String str ->
    Wrapper.Graph.set_attr_string operation_description ~attr_name str
  | Type dtype ->
    let dtype = Node.Type.to_data_type dtype in
    Wrapper.Graph.set_attr_type operation_description ~attr_name dtype
  | Tensor_float tensor_float ->
    let tensor = Tensor.create Float32 (Array.of_list tensor_float.shape) in
    Tensor.copy_elt_list tensor tensor_float.values;
    Wrapper.Graph.set_attr_tensor operation_description ~attr_name (Tensor.P tensor)
    |> Wrapper.Status.ok_exn
  | Tensor_int tensor_int ->
    let tensor = Tensor.create Int32 (Array.of_list tensor_int.shape) in
    Tensor.copy_elt_list tensor (List.map tensor_int.values ~f:Int32.of_int_exn);
    Wrapper.Graph.set_attr_tensor operation_description ~attr_name (Tensor.P tensor)
    |> Wrapper.Status.ok_exn
  | Int i ->
    Wrapper.Graph.set_attr_int operation_description ~attr_name i
  | Float f ->
    Wrapper.Graph.set_attr_float operation_description ~attr_name f
  | Bool b ->
    Wrapper.Graph.set_attr_bool operation_description ~attr_name b
  | List _ -> assert false
  | Tensor_string _ -> assert false
  | Shape _ -> assert false

let rec build t node =
  let id = Node.packed_id node in
  match Hashtbl.find t.nodes id with
  | Some op -> op
  | None ->
    let Node.P u_node = node in
    (* TODO: [Var.get_init]. *)
    let operation_description =
      Wrapper.Graph.new_operation t.graph
        ~op_name:(Node.op_name u_node |> Node.Op_name.to_string)
        ~name:(Node.unique_name u_node)
    in
    List.iter (Node.inputs u_node) ~f:(function
      | `single input ->
        Wrapper.Graph.add_input
          operation_description
          (build t input)
          ~index:(Node.packed_output_idx input |> Option.value ~default:0)
      | `multi inputs ->
        let inputs =
          List.map inputs ~f:(fun input ->
            let index = Node.packed_output_idx input |> Option.value ~default:0 in
            build t input, index)
        in
        Wrapper.Graph.add_inputs operation_description inputs);
    List.iter (Node.attributes u_node) ~f:(fun (attr_name, attr) ->
      add_attribute operation_description ~attr_name attr);
    let operation =
      Wrapper.Graph.finish_operation operation_description<commit_msg>Proper error messages.<commit_after>open Core_kernel.Std

type t =
  { session : Wrapper.Session_with_graph.t
  ; graph : Wrapper.Graph.t
  ; nodes : Wrapper.Graph.operation Node.Id.Table.t
  }

let create () =
  let graph = Wrapper.Graph.create () in
  match Wrapper.Session_with_graph.create graph with
  | Error status ->
    failwithf "Unable to generate session: %s" (Wrapper.Status.message status) ()
  | Ok session ->
    { session
    ; graph
    ; nodes = Node.Id.Table.create ()
    }

let default_session = lazy (create ())

let add_attribute operation_description ~attr_name attr =
  match (attr : Node.attr) with
  | String str ->
    Wrapper.Graph.set_attr_string operation_description ~attr_name str
  | Type dtype ->
    let dtype = Node.Type.to_data_type dtype in
    Wrapper.Graph.set_attr_type operation_description ~attr_name dtype
  | Tensor_float tensor_float ->
    let tensor = Tensor.create Float32 (Array.of_list tensor_float.shape) in
    Tensor.copy_elt_list tensor tensor_float.values;
    Wrapper.Graph.set_attr_tensor operation_description ~attr_name (Tensor.P tensor)
    |> Wrapper.Status.ok_exn
  | Tensor_int tensor_int ->
    let tensor = Tensor.create Int32 (Array.of_list tensor_int.shape) in
    Tensor.copy_elt_list tensor (List.map tensor_int.values ~f:Int32.of_int_exn);
    Wrapper.Graph.set_attr_tensor operation_description ~attr_name (Tensor.P tensor)
    |> Wrapper.Status.ok_exn
  | Int i ->
    Wrapper.Graph.set_attr_int operation_description ~attr_name i
  | Float f ->
    Wrapper.Graph.set_attr_float operation_description ~attr_name f
  | Bool b ->
    Wrapper.Graph.set_attr_bool operation_description ~attr_name b
  | List _ -> failwith "List attributes are not supported yet."
  | Tensor_string _ -> failwith "Tensor_string attributes are not supported yet."
  | Shape _ -> failwith "Shape attributes are not supported yet."

let rec build t node =
  let id = Node.packed_id node in
  match Hashtbl.find t.nodes id with
  | Some op -> op
  | None ->
    let Node.P u_node = node in
    (* TODO: [Var.get_init]. *)
    let operation_description =
      Wrapper.Graph.new_operation t.graph
        ~op_name:(Node.op_name u_node |> Node.Op_name.to_string)
        ~name:(Node.unique_name u_node)
    in
    List.iter (Node.inputs u_node) ~f:(function
      | `single input ->
        Wrapper.Graph.add_input
          operation_description
          (build t input)
          ~index:(Node.packed_output_idx input |> Option.value ~default:0)
      | `multi inputs ->
        let inputs =
          List.map inputs ~f:(fun input ->
            let index = Node.packed_output_idx input |> Option.value ~default:0 in
            build t input, index)
        in
        Wrapper.Graph.add_inputs operation_description inputs);
    List.iter (Node.attributes u_node) ~f:(fun (attr_name, attr) ->
      add_attribute operation_description ~attr_name attr);
    let operation =
      Wrapper.Graph.finish_operation operation_description
{'fraction_non_alphanumeric': 0.09324648611587247, 'fraction_numerical': 0.002742543709290367, 'mean_word_length': 3.1441761363636362, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 49, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '879014', 'n_tokens_mistral': 1836, 'n_tokens_neox': 1787, 'n_words': 456}