<commit_before>open Core_kernel.Std
open Tensorflow_core
open Tensorflow

module Image : sig
  val load : string -> (float, Bigarray.float32_elt) Tensorflow_core.Tensor.t * int * int
  val save : (float, Bigarray.float32_elt) Tensorflow_core.Tensor.t -> string -> unit
end = struct
  let imagenet_mean = function
    | `blue -> 103.939
    | `green -> 116.779
    | `red -> 123.68

  let normalize x ~channel =
    float x -. imagenet_mean channel

  let unnormalize x ~channel =
    min 255 (int_of_float (x +. imagenet_mean channel))
    |> max 0

  let load filename =
    let image = OImages.load filename [] in
    let image = OImages.rgb24 image in
    let img_w = image # width in
    let img_h = image # height in
    let tensor = Tensor.create3 Float32 img_h img_w 3 in
    for i = 0 to img_h - 1 do
      for j = 0 to img_w - 1 do
        let { Color.r; g; b } = image # get j i in
        Tensor.set tensor [| i; j; 0 |] (normalize r ~channel:`red);
        Tensor.set tensor [| i; j; 1 |] (normalize g ~channel:`green);
        Tensor.set tensor [| i; j; 2 |] (normalize b ~channel:`blue);
      done;
    done;
    tensor, img_w, img_h

  let save tensor filename =
    let img_w, img_h =
      match Tensor.dims tensor with
      | [| img_h; img_w; 3 |] -> img_w, img_h
      | _ -> failwith "Improper tensor dimensions"
    in
    let image = new OImages.rgb24 img_w img_h in
    for i = 0 to img_h - 1 do
      for j = 0 to img_w - 1 do
        let r = Tensor.get tensor [| i; j; 0 |] |> unnormalize ~channel:`red in
        let g = Tensor.get tensor [| i; j; 1 |] |> unnormalize ~channel:`green in
        let b = Tensor.get tensor [| i; j; 2 |] |> unnormalize ~channel:`blue in
        image # set j i { Color.r; g; b }
      done;
    done;
    image # save filename None []
end

let conv2d node ~in_channels ~out_channels =
  let w = Var.f_or_d [ 3; 3; in_channels; out_channels ] ~type_:Float 0. in
  let b = Var.f_or_d [ out_channels ] ~type_:Float 0. in
  let conv2d = Ops.conv2D ~strides:[ 1; 1; 1; 1 ] ~padding:"SAME" node w in
  Ops.(+) conv2d b, w, b

let max_pool node =
  Ops.maxPool node ~ksize:[ 1; 2; 2; 1 ] ~strides:[ 1; 2; 2; 1 ] ~padding:"SAME"

let load var_and_names ~npz_filename =
  let npz = Npy.Npz.open_in npz_filename in
  let inputs, targets =
    List.map var_and_names ~f:(fun (var_name, var) ->
      let Npy.P tensor = Npy.Npz.read npz var_name in
      match Bigarray.Genarray.kind tensor, Bigarray.Genarray.layout tensor with
      | Bigarray.Float32, Bigarray.C_layout ->
        let tensor = Tensor.of_bigarray tensor ~scalar:false in
        let shape = Tensor.dims tensor |> Array.to_list in
        let ph = Ops.placeholder ~type_:Float shape in
        Session.Input.float ph tensor, Node.P (Ops.assign var (Ops.Placeholder.to_node ph))
      | _ -> failwith "Improper weight types or layout")
    |> List.unzip
  in
  Session.run ~inputs ~targets Session.Output.empty;
  Npy.Npz.close_in npz

let style_grams_and_content_nodes input ~img_h ~img_w ~npz_filename =
  let var_by_name = String.Table.create () in
  let block iter ~block_idx ~in_channels ~out_channels (node, acc) =
    List.init iter ~f:Fn.id
    |> List.fold ~init:(node, []) ~f:(fun (node, acc_relus) idx ->
      let name = sprintf "conv%d_%d" block_idx (idx+1) in
      let in_channels = if idx = 0 then in_channels else out_channels in
      let conv2d, w, b = conv2d node ~in_channels ~out_channels in
      Hashtbl.set var_by_name ~key:(name ^ "W") ~data:w;
      Hashtbl.set var_by_name ~key:(name ^ "b") ~data:b;
      let relu = Ops.relu conv2d in
      relu, (relu, (`block_idx block_idx, `out_channels out_channels)) :: acc_relus)
    |> fun (node, acc_relus) -> max_pool node, List.rev acc_relus :: acc
  in
  let _model, acc_relus =
    (Ops.reshape input (Ops.const_int ~type_:Int32 [ 1; img_h; img_w; 3 ]), [])
    |> block 2 ~block_idx:1 ~in_channels:3   ~out_channels:64
    |> block 2 ~block_idx:2 ~in_channels:64  ~out_channels:128
    |> block 4 ~block_idx:3 ~in_channels:128 ~out_channels:256
    |> block 4 ~block_idx:4 ~in_channels:256 ~out_channels:512
    |> block 4 ~block_idx:5 ~in_channels:512 ~out_channels:512
  in
  let acc_relus = List.rev acc_relus in
  load (Hashtbl.to_alist var_by_name) ~npz_filename;
  let style_grams =
    List.map acc_relus ~f:(fun relus ->
      let node, (`block_idx block_idx, `out_channels out_channels) = List.hd_exn relus in
      let node = Ops.reshape node (Ops.const_int ~type_:Int32 [ -1; out_channels ]) in
      let two = Int.pow 2 (block_idx - 1) in
      let size_ = float (out_channels * (img_h/two) * (img_w/two)) in
      Ops.(matMul ~transpose_a:true node node / f size_))
  in
  let content_nodes =
    (* Block 4, conv 2. *)
    List.map [ 4, 2 ] ~f:(fun (block_idx, conv_idx) ->
      let block = List.nth_exn acc_relus (block_idx - 1) in
      List.nth_exn block (conv_idx - 1) |> fst)
  in
  style_grams, content_nodes

let compute_grams ~filename ~npz_filename =
  let input_tensor, img_w, img_h = Image.load filename in
  let input_placeholder = Ops.placeholder ~type_:Float [ img_h; img_w; 3 ] in
  let style_grams, _ =
    style_grams_and_content_nodes (Ops.Placeholder.to_node input_placeholder)
      ~img_h ~img_w ~npz_filename
  in
  List.map style_grams ~f:(fun node ->
    Session.run
      ~inputs:[ Session.Input.float input_placeholder input_tensor ]
      ~targets:[ Node.P node ]
      (Session.Output.float node))

let total_variation_loss input ~img_h ~img_w =
  let input =
    Ops.reshape input (Ops.const_int ~type_:Int32 [ img_h; img_w; 3 ])
  in
  let axis1_diff =
    Ops.(-)
      (Ops.slice input (Ops.ci32 [ 0; 0; 0 ]) (Ops.ci32 [ img_h-1; img_w; 3 ]))
      (Ops.slice input (Ops.ci32 [ 1; 0; 0 ]) (Ops.ci32 [ img_h-1; img_w; 3 ]))
  in
  let axis2_diff =
    Ops.(-)
      (Ops.slice input (Ops.ci32 [ 0; 0; 0 ]) (Ops.ci32 [ img_h; img_w-1; 3 ]))
      (Ops.slice input (Ops.ci32 [ 0; 1; 0 ]) (Ops.ci32 [ img_h; img_w-1; 3 ]))
  in
  Ops.(reduce_sum (axis1_diff * axis1_diff) + reduce_sum (axis2_diff * axis2_diff))

let create_and_set_var tensor =
  let dims = Tensor.dims tensor |> Array.to_list in
  let input_var = Var.f_or_d dims ~type_:Float 0. in
  let placeholder = Ops.placeholder dims ~type_:Float in
  let assign = Ops.assign input_var (Ops.Placeholder.to_node placeholder) in
  Session.run
    ~inputs:[ Session.Input.float placeholder tensor ]
    ~targets:[ Node.P assign ]
    Session.Output.empty;
  input_var

let run epochs learning_rate content_weight style_weight tv_weight npz_filename input_filename style_filename =
  let suffix =
    String.rsplit2 input_filename ~on:'.'
    |> Option.value_map ~f:snd ~default:"jpg"
  in
  let input_tensor, img_w, img_h = Image.load input_filename in
  printf "Computing target features...\n%!";
  let target_grams = compute_grams ~filename:style_filename ~npz_filename in
  printf "Done computing target features.\n%!";
  let input_var = create_and_set_var input_tensor in
  let style_grams, content_nodes =
    style_grams_and_content_nodes input_var ~img_h ~img_w ~npz_filename
  in
  let style_losses, style_inputs =
    List.map2_exn style_grams target_grams ~f:(fun gram_node target_gram ->
      let dims = Tensor.dims target_gram |> Array.to_list in
      let size_ = List.reduce_exn dims ~f:( * ) in
      let placeholder = Ops.placeholder ~type_:Float dims in
      let diff = Ops.(-) gram_node (Ops.Placeholder.to_node placeholder) in
      Ops.(reduce_sum (diff * diff) / f (float size_)),
      Session.Input.float placeholder target_gram)
    |> List.unzip
  in<commit_msg>Factorize the code setting the variables.<commit_after>open Core_kernel.Std
open Tensorflow_core
open Tensorflow

module Image : sig
  val load : string -> (float, Bigarray.float32_elt) Tensorflow_core.Tensor.t * int * int
  val save : (float, Bigarray.float32_elt) Tensorflow_core.Tensor.t -> string -> unit
end = struct
  let imagenet_mean = function
    | `blue -> 103.939
    | `green -> 116.779
    | `red -> 123.68

  let normalize x ~channel =
    float x -. imagenet_mean channel

  let unnormalize x ~channel =
    min 255 (int_of_float (x +. imagenet_mean channel))
    |> max 0

  let load filename =
    let image = OImages.load filename [] in
    let image = OImages.rgb24 image in
    let img_w = image # width in
    let img_h = image # height in
    let tensor = Tensor.create3 Float32 img_h img_w 3 in
    for i = 0 to img_h - 1 do
      for j = 0 to img_w - 1 do
        let { Color.r; g; b } = image # get j i in
        Tensor.set tensor [| i; j; 0 |] (normalize r ~channel:`red);
        Tensor.set tensor [| i; j; 1 |] (normalize g ~channel:`green);
        Tensor.set tensor [| i; j; 2 |] (normalize b ~channel:`blue);
      done;
    done;
    tensor, img_w, img_h

  let save tensor filename =
    let img_w, img_h =
      match Tensor.dims tensor with
      | [| img_h; img_w; 3 |] -> img_w, img_h
      | _ -> failwith "Improper tensor dimensions"
    in
    let image = new OImages.rgb24 img_w img_h in
    for i = 0 to img_h - 1 do
      for j = 0 to img_w - 1 do
        let r = Tensor.get tensor [| i; j; 0 |] |> unnormalize ~channel:`red in
        let g = Tensor.get tensor [| i; j; 1 |] |> unnormalize ~channel:`green in
        let b = Tensor.get tensor [| i; j; 2 |] |> unnormalize ~channel:`blue in
        image # set j i { Color.r; g; b }
      done;
    done;
    image # save filename None []
end

let set_vars var_and_tensors =
  let inputs, targets =
    List.map var_and_tensors ~f:(fun (var, tensor) ->
      let dims = Tensor.dims tensor |> Array.to_list in
      let placeholder = Ops.placeholder dims ~type_:Float in
      let assign = Ops.assign var (Ops.Placeholder.to_node placeholder) in
      Session.Input.float placeholder tensor, Node.P assign)
    |> List.unzip
  in
  Session.run ~inputs ~targets Session.Output.empty

let load var_and_names ~npz_filename =
  let npz = Npy.Npz.open_in npz_filename in
  List.map var_and_names ~f:(fun (var_name, var) ->
    let Npy.P tensor = Npy.Npz.read npz var_name in
    match Bigarray.Genarray.kind tensor, Bigarray.Genarray.layout tensor with
    | Bigarray.Float32, Bigarray.C_layout ->
      var, (Tensor.of_bigarray tensor ~scalar:false : (float, Bigarray.float32_elt) Tensor.t)
    | _ -> failwith "Improper weight types or layout")
  |> set_vars;
  Npy.Npz.close_in npz

let conv2d node ~in_channels ~out_channels =
  let w = Var.f_or_d [ 3; 3; in_channels; out_channels ] ~type_:Float 0. in
  let b = Var.f_or_d [ out_channels ] ~type_:Float 0. in
  let conv2d = Ops.conv2D ~strides:[ 1; 1; 1; 1 ] ~padding:"SAME" node w in
  Ops.(+) conv2d b, w, b

let max_pool node =
  Ops.maxPool node ~ksize:[ 1; 2; 2; 1 ] ~strides:[ 1; 2; 2; 1 ] ~padding:"SAME"

let style_grams_and_content_nodes input ~img_h ~img_w ~npz_filename =
  let var_by_name = String.Table.create () in
  let block iter ~block_idx ~in_channels ~out_channels (node, acc) =
    List.init iter ~f:Fn.id
    |> List.fold ~init:(node, []) ~f:(fun (node, acc_relus) idx ->
      let name = sprintf "conv%d_%d" block_idx (idx+1) in
      let in_channels = if idx = 0 then in_channels else out_channels in
      let conv2d, w, b = conv2d node ~in_channels ~out_channels in
      Hashtbl.set var_by_name ~key:(name ^ "W") ~data:w;
      Hashtbl.set var_by_name ~key:(name ^ "b") ~data:b;
      let relu = Ops.relu conv2d in
      relu, (relu, (`block_idx block_idx, `out_channels out_channels)) :: acc_relus)
    |> fun (node, acc_relus) -> max_pool node, List.rev acc_relus :: acc
  in
  let _model, acc_relus =
    (Ops.reshape input (Ops.const_int ~type_:Int32 [ 1; img_h; img_w; 3 ]), [])
    |> block 2 ~block_idx:1 ~in_channels:3   ~out_channels:64
    |> block 2 ~block_idx:2 ~in_channels:64  ~out_channels:128
    |> block 4 ~block_idx:3 ~in_channels:128 ~out_channels:256
    |> block 4 ~block_idx:4 ~in_channels:256 ~out_channels:512
    |> block 4 ~block_idx:5 ~in_channels:512 ~out_channels:512
  in
  let acc_relus = List.rev acc_relus in
  load (Hashtbl.to_alist var_by_name) ~npz_filename;
  let style_grams =
    List.map acc_relus ~f:(fun relus ->
      let node, (`block_idx block_idx, `out_channels out_channels) = List.hd_exn relus in
      let node = Ops.reshape node (Ops.const_int ~type_:Int32 [ -1; out_channels ]) in
      let two = Int.pow 2 (block_idx - 1) in
      let size_ = float (out_channels * (img_h/two) * (img_w/two)) in
      Ops.(matMul ~transpose_a:true node node / f size_))
  in
  let content_nodes =
    (* Block 4, conv 2. *)
    List.map [ 4, 2 ] ~f:(fun (block_idx, conv_idx) ->
      let block = List.nth_exn acc_relus (block_idx - 1) in
      List.nth_exn block (conv_idx - 1) |> fst)
  in
  style_grams, content_nodes

let compute_grams ~filename ~npz_filename =
  let input_tensor, img_w, img_h = Image.load filename in
  let input_placeholder = Ops.placeholder ~type_:Float [ img_h; img_w; 3 ] in
  let style_grams, _ =
    style_grams_and_content_nodes (Ops.Placeholder.to_node input_placeholder)
      ~img_h ~img_w ~npz_filename
  in
  List.map style_grams ~f:(fun node ->
    Session.run
      ~inputs:[ Session.Input.float input_placeholder input_tensor ]
      ~targets:[ Node.P node ]
      (Session.Output.float node))

let total_variation_loss input ~img_h ~img_w =
  let input =
    Ops.reshape input (Ops.const_int ~type_:Int32 [ img_h; img_w; 3 ])
  in
  let axis1_diff =
    Ops.(-)
      (Ops.slice input (Ops.ci32 [ 0; 0; 0 ]) (Ops.ci32 [ img_h-1; img_w; 3 ]))
      (Ops.slice input (Ops.ci32 [ 1; 0; 0 ]) (Ops.ci32 [ img_h-1; img_w; 3 ]))
  in
  let axis2_diff =
    Ops.(-)
      (Ops.slice input (Ops.ci32 [ 0; 0; 0 ]) (Ops.ci32 [ img_h; img_w-1; 3 ]))
      (Ops.slice input (Ops.ci32 [ 0; 1; 0 ]) (Ops.ci32 [ img_h; img_w-1; 3 ]))
  in
  Ops.(reduce_sum (axis1_diff * axis1_diff) + reduce_sum (axis2_diff * axis2_diff))

let create_and_set_var tensor =
  let dims = Tensor.dims tensor |> Array.to_list in
  let input_var = Var.f_or_d dims ~type_:Float 0. in
  set_vars [ input_var, tensor ];
  input_var

let run epochs learning_rate content_weight style_weight tv_weight npz_filename input_filename style_filename =
  let suffix =
    String.rsplit2 input_filename ~on:'.'
    |> Option.value_map ~f:snd ~default:"jpg"
  in
  let input_tensor, img_w, img_h = Image.load input_filename in
  printf "Computing target features...\n%!";
  let target_grams = compute_grams ~filename:style_filename ~npz_filename in
  printf "Done computing target features.\n%!";
  let input_var = create_and_set_var input_tensor in
  let style_grams, content_nodes =
    style_grams_and_content_nodes input_var ~img_h ~img_w ~npz_filename
  in
  let style_losses, style_inputs =
    List.map2_exn style_grams target_grams ~f:(fun gram_node target_gram ->
      let dims = Tensor.dims target_gram |> Array.to_list in
      let size_ = List.reduce_exn dims ~f:( * ) in
      let placeholder = Ops.placeholder ~type_:Float dims in
      let diff = Ops.(-) gram_node (Ops.Placeholder.to_node placeholder) in
      Ops.(reduce_sum (diff * diff) / f (float size_)),
      Session.Input.float placeholder target_gram)
    |> List.unzip
  in
{'fraction_non_alphanumeric': 0.11401362343760334, 'fraction_numerical': 0.023278883671714833, 'mean_word_length': 3.112591786782703, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 75, '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': '1672198', 'n_tokens_mistral': 5865, 'n_tokens_neox': 5491, 'n_words': 1605}