Source file gpt2.ml

open Rune
module Ptree = Kaun.Ptree
module Activations = Kaun.Activations

(* Configuration *)

type config = {
  vocab_size : int;
  n_positions : int;
  n_embd : int;
  n_layer : int;
  n_head : int;
  n_inner : int option;
  activation_function : [ `gelu | `relu | `swish | `gelu_new ];
  resid_pdrop : float;
  embd_pdrop : float;
  attn_pdrop : float;
  layer_norm_epsilon : float;
  initializer_range : float;
  scale_attn_weights : bool;
  use_cache : bool;
  scale_attn_by_inverse_layer_idx : bool;
  reorder_and_upcast_attn : bool;
  bos_token_id : int option;
  eos_token_id : int option;
  pad_token_id : int option;
}

let default_config =
  {
    vocab_size = 50257;
    n_positions = 1024;
    n_embd = 768;
    n_layer = 12;
    n_head = 12;
    n_inner = None;
    (* Defaults to 4 * n_embd *)
    activation_function = `gelu_new;
    resid_pdrop = 0.1;
    embd_pdrop = 0.1;
    attn_pdrop = 0.1;
    layer_norm_epsilon = 1e-5;
    initializer_range = 0.02;
    scale_attn_weights = true;
    use_cache = true;
    scale_attn_by_inverse_layer_idx = false;
    reorder_and_upcast_attn = false;
    bos_token_id = Some 50256;
    eos_token_id = Some 50256;
    pad_token_id = None;
  }

let gpt2_small = default_config

let gpt2_medium =
  { default_config with n_embd = 1024; n_layer = 24; n_head = 16 }

let gpt2_large =
  { default_config with n_embd = 1280; n_layer = 36; n_head = 20 }

let gpt2_xl = { default_config with n_embd = 1600; n_layer = 48; n_head = 25 }

(* Input type *)
type inputs = {
  input_ids : (int32, int32_elt) Rune.t;
  attention_mask : (int32, int32_elt) Rune.t option;
  position_ids : (int32, int32_elt) Rune.t option;
}

(* Removed dtype inference; dtype is provided by caller/policy *)

(* GPT-2 specific tokenizer with BPE *)
module Tokenizer = struct
  type t = {
    tokenizer : Saga.Tokenizer.t;
        (* Store the actual tokenizer for encoding/decoding *)
    vocab_size : int;
    bos_token_id : int;
    eos_token_id : int;
    pad_token_id : int option;
  }

  let download_vocab_and_merges model_id =
    (* Download vocab and merges files from HuggingFace if not present *)
    let model_cache =
      Nx_io.Cache_dir.get_path_in_cache ~scope:[ "models"; "gpt2" ] model_id
    in
    let vocab_file = Filename.concat model_cache "vocab.json" in
    let merges_file = Filename.concat model_cache "merges.txt" in

    (* Create cache directories if they don't exist *)
    if not (Sys.file_exists model_cache) then
      Sys.command (Printf.sprintf "mkdir -p %s" model_cache) |> ignore;

    (* Download vocab.json if it doesn't exist *)
    if not (Sys.file_exists vocab_file) then (
      Printf.printf "Downloading vocab.json for %s...\n%!" model_id;
      let url =
        Printf.sprintf "https://huggingface.co/%s/resolve/main/vocab.json"
          model_id
      in
      let cmd =
        Printf.sprintf
          "curl -L -o %s %s 2>/dev/null || wget -O %s %s 2>/dev/null" vocab_file
          url vocab_file url
      in
      let exit_code = Sys.command cmd in
      if exit_code <> 0 then
        failwith
          (Printf.sprintf "Failed to download vocab.json for %s" model_id));

    (* Download merges.txt if it doesn't exist *)
    if not (Sys.file_exists merges_file) then (
      Printf.printf "Downloading merges.txt for %s...\n%!" model_id;
      let url =
        Printf.sprintf "https://huggingface.co/%s/resolve/main/merges.txt"
          model_id
      in
      let cmd =
        Printf.sprintf
          "curl -L -o %s %s 2>/dev/null || wget -O %s %s 2>/dev/null"
          merges_file url merges_file url
      in
      let exit_code = Sys.command cmd in
      if exit_code <> 0 then
        failwith
          (Printf.sprintf "Failed to download merges.txt for %s" model_id));

    (vocab_file, merges_file)

  let create ?vocab_file ?merges_file ?model_id () =
    (* Either provide vocab_file and merges_file paths, or a model_id to
       download from *)
    let vocab_file, merges_file =
      match (vocab_file, merges_file, model_id) with
      | Some vf, Some mf, _ -> (vf, mf)
      | None, None, Some id -> download_vocab_and_merges id
      | None, None, None -> download_vocab_and_merges "gpt2" (* Default *)
      | _ ->
          failwith "Either provide both vocab_file and merges_file, or model_id"
    in

    (* Create GPT-2 tokenizer with ByteLevel pre-tokenizer Use use_regex:true to
       enable GPT-2 pattern splitting *)
    let tokenizer =
      Saga.Tokenizer.from_model_file ~vocab:vocab_file ~merges:merges_file
        ~pre:
          (Saga.Pre_tokenizers.byte_level ~add_prefix_space:false
             ~use_regex:true ())
        ~decoder:(Saga.Decoders.byte_level ())
        ()
    in
    {
      tokenizer;
      vocab_size = 50257;
      (* GPT-2 vocab size *)
      bos_token_id = 50256;
      (* <|endoftext|> *)
      eos_token_id = 50256;
      (* <|endoftext|> *)
      pad_token_id = None;
      (* GPT-2 doesn't use padding by default *)
    }

  let encode_to_array t text =
    (* Use the tokenizer's encode method *)
    let encoding = Saga.Tokenizer.encode t.tokenizer text in
    Saga.Encoding.get_ids encoding

  let encode t text =
    let token_ids = encode_to_array t text in
    let seq_len = Array.length token_ids in

    (* Convert to tensors *)
    let input_ids =
      Rune.create Int32 [| 1; seq_len |] (Array.map Int32.of_int token_ids)
    in

    (* Return inputs record *)
    { input_ids; attention_mask = None; position_ids = None }

  let encode_batch t ?(max_length = 1024) ?(padding = false) texts =
    let encoded = List.map (encode_to_array t) texts in

    (* Find max length or use specified max_length *)
    let actual_max =
      if padding then max_length
      else
        List.fold_left (fun acc arr -> Int.max acc (Array.length arr)) 0 encoded
    in

    (* Ensure we don't exceed vocabulary size *)
    let vocab_size = t.vocab_size in
    let validate_tokens arr =
      Array.map
        (fun token_id ->
          if token_id >= 0 && token_id < vocab_size then token_id
          else t.eos_token_id (* Replace invalid tokens with EOS *))
        arr
    in

    (* Pad sequences if padding is enabled *)
    let padded =
      if padding then
        let pad_token_id =
          Option.value t.pad_token_id ~default:t.eos_token_id
        in
        List.map
          (fun arr ->
            let validated = validate_tokens arr in
            let len = Array.length validated in
            if len >= actual_max then Array.sub validated 0 actual_max
            else
              Array.append validated
                (Array.make (actual_max - len) pad_token_id))
          encoded
      else
        (* Truncate to max length if no padding *)
        List.map
          (fun arr ->
            let validated = validate_tokens arr in
            let len = Array.length validated in
            if len > actual_max then Array.sub validated 0 actual_max
            else validated)
          encoded
    in

    (* Convert to tensor *)
    let batch_size = List.length padded in
    let flat_data = Array.concat padded in
    let nx_tensor =
      let data = Array.map Int32.of_int flat_data in
      Nx.create Int32 [| batch_size; actual_max |] data
    in
    (* Convert to Rune tensor *)
    Rune.of_nx nx_tensor

  let decode t token_ids =
    (* Decode token IDs back to text using tokenizer *)
    Saga.Tokenizer.decode t.tokenizer token_ids

  (* Get special token IDs *)
  let get_bos_token_id t = t.bos_token_id
  let get_eos_token_id t = t.eos_token_id
  let get_pad_token_id t = t.pad_token_id
  let get_vocab_size t = t.vocab_size
end

(* GPT-2 Embeddings *)
let embeddings ~config () =
  let open Kaun.Layer in
  (* GPT-2 embeddings: token + position *)
  let token_embeddings =
    embedding ~vocab_size:config.vocab_size ~embed_dim:config.n_embd
      ~scale:false ()
  in
  let position_embeddings =
    embedding ~vocab_size:config.n_positions ~embed_dim:config.n_embd
      ~scale:false ()
  in
  let dropout = dropout ~rate:config.embd_pdrop () in

  (* Custom module that applies both embeddings and sums them *)
  {
    Kaun.init =
      (fun ~rngs ~dtype ->
        let keys = Rune.Rng.split ~n:3 rngs in
        Ptree.dict
          [
            ("token_embeddings", token_embeddings.init ~rngs:keys.(0) ~dtype);
            ( "position_embeddings",
              position_embeddings.init ~rngs:keys.(1) ~dtype );
            ("dropout", dropout.init ~rngs:keys.(2) ~dtype);
          ]);
    Kaun.apply =
      (fun params ~training ?rngs x ->
        (* x is expected to be float tensor, but we need int indices *)
        let input_ids = Rune.cast Rune.int32 x in
        let dtype = Rune.dtype x in
        let token_embeddings_table =
          Ptree.get_tensor_exn
            ~path:(Ptree.Path.of_string "token_embeddings.embedding")
            params dtype
        in
        let position_embeddings_table =
          Ptree.get_tensor_exn
            ~path:(Ptree.Path.of_string "position_embeddings.embedding")
            params dtype
        in
        let lookup_embeddings embedding_table indices =
          let shape = Rune.shape indices in
          let batch_size = shape.(0) in
          let seq_len = shape.(1) in
          let table_shape = Rune.shape embedding_table in
          if Array.length table_shape <> 2 then
            failwith
              (Printf.sprintf
                 "Embedding table has wrong shape: %d dims, expected 2"
                 (Array.length table_shape));
          let embed_dim = table_shape.(1) in
          let indices_flat = Rune.reshape [| batch_size * seq_len |] indices in
          let gathered = Rune.take ~axis:0 indices_flat embedding_table in
          Rune.reshape [| batch_size; seq_len; embed_dim |] gathered
        in
        let token_embeds = lookup_embeddings token_embeddings_table input_ids in
        let seq_len =
          (Rune.shape input_ids).(Array.length (Rune.shape input_ids) - 1)
        in
        let batch_size = (Rune.shape input_ids).(0) in
        let position_ids =
          let pos_ids = Rune.zeros Rune.int32 [| batch_size; seq_len |] in
          for b = 0 to batch_size - 1 do
            for s = 0 to seq_len - 1 do
              Rune.set [ b; s ] pos_ids
                (Rune.scalar Rune.int32 (Int32.of_int s))
            done
          done;
          pos_ids
        in
        let position_embeds =
          lookup_embeddings position_embeddings_table position_ids
        in

        (* Sum embeddings *)
        let embeddings = Rune.add token_embeds position_embeds in

        (* Apply dropout *)
        let dropout_params =
          match Ptree.get ~path:(Ptree.Path.of_string "dropout") params with
          | Some p -> p
          | None -> Ptree.List []
        in
        let embeddings =
          dropout.apply dropout_params ~training ?rngs embeddings
        in

        embeddings);
  }

(* Main Model *)

type 'a gpt2 = {
  model : Kaun.Layer.module_;
  params : Ptree.t;
  config : config;
  dtype : (float, 'a) Rune.dtype;
}

type 'a output = {
  last_hidden_state : (float, 'a) Rune.t;
  hidden_states : (float, 'a) Rune.t list option;
  attentions : (float, 'a) Rune.t list option;
}

module Gpt2_block = struct
  (* Generate causal attention mask *)
  let causal_mask ~seq_len ~dtype =
    (* Create a lower triangular matrix for causal masking *)
    let mask = ones dtype [| seq_len; seq_len |] in
    tril mask ~k:0

  (* GPT-2 uses GELU activation *)
  let gelu x =
    (* Use exact GELU with erf for numerical stability *)
    (* GELU(x) = 0.5 * x * (1 + erf(x / sqrt(2))) *)
    Rune.gelu x

  (* Multi-head attention with causal masking *)
  let causal_attention ~n_head ~hidden_size ~params x =
    let batch_size = (shape x).(0) in
    let seq_len = (shape x).(1) in
    let head_dim = hidden_size / n_head in
    let dtype = dtype x in

    (* Get Q, K, V weights and biases *)
    let get_param name = Ptree.Dict.get_tensor_exn params ~name dtype in

    let qkv_weight = get_param "qkv_weight" in
    let qkv_bias = get_param "qkv_bias" in
    let out_weight = get_param "attn_out_weight" in
    let out_bias = get_param "attn_out_bias" in

    (* Compute combined QKV *)
    let qkv =
      add (matmul x qkv_weight) (reshape [| 1; 1; 3 * hidden_size |] qkv_bias)
    in

    (* Split into Q, K, V *)
    let query = slice [ A; A; R (0, hidden_size) ] qkv in
    let key = slice [ A; A; R (hidden_size, 2 * hidden_size) ] qkv in
    let value = slice [ A; A; R (2 * hidden_size, 3 * hidden_size) ] qkv in

    (* Reshape for multi-head attention: [batch, seq, n_head, head_dim] *)
    let reshape_for_heads t =
      let t = reshape [| batch_size; seq_len; n_head; head_dim |] t in
      (* Transpose to [batch, n_head, seq, head_dim] *)
      transpose ~axes:[ 0; 2; 1; 3 ] t
    in

    let query = reshape_for_heads query in
    let key = reshape_for_heads key in
    let value = reshape_for_heads value in

    (* Scaled dot-product attention with causal mask *)
    (* scores = Q @ K^T / sqrt(head_dim) *)
    let key_t = transpose ~axes:[ 0; 1; 3; 2 ] key in
    let scores = matmul query key_t in
    let scores = div_s scores (Float.sqrt (Float.of_int head_dim)) in

    (* Apply causal mask *)
    let mask = causal_mask ~seq_len ~dtype in
    (* Expand mask for batch and heads: [1, 1, seq_len, seq_len] *)
    let mask = reshape [| 1; 1; seq_len; seq_len |] mask in
    (* Where mask is 0, set score to large negative value *)
    let neg_inf = full dtype [| 1; 1; 1; 1 |] (-1e10) in
    let scores = where (equal_s mask 0.0) neg_inf scores in

    (* Softmax over last dimension *)
    let attn_weights = softmax scores ~axes:[ 3 ] in

    (* Apply attention to values *)
    let attn_output = matmul attn_weights value in

    (* Transpose back and reshape *)
    let attn_output = transpose ~axes:[ 0; 2; 1; 3 ] attn_output in
    (* Make contiguous before reshaping - force a copy to ensure it's
       contiguous *)
    let attn_output = copy attn_output in
    let attn_output =
      reshape [| batch_size; seq_len; hidden_size |] attn_output
    in

    (* Output projection *)
    add
      (matmul attn_output out_weight)
      (reshape [| 1; 1; hidden_size |] out_bias)

  (* Feed-forward network *)
  let gelu_new x =
    let dtype = Rune.dtype x in
    let half = Rune.scalar dtype 0.5 in
    let one = Rune.scalar dtype 1.0 in
    let sqrt_2_over_pi =
      Rune.scalar dtype 0.7978845608028654
      (* sqrt(2/pi) *)
    in
    let coeff = Rune.scalar dtype 0.044715 in
    let x_cubed = Rune.mul x (Rune.mul x x) in
    let inner = Rune.mul sqrt_2_over_pi (Rune.add x (Rune.mul coeff x_cubed)) in
    let tanh_inner = Rune.tanh inner in
    Rune.mul (Rune.mul half x) (Rune.add one tanh_inner)

  let mlp ~n_inner ~hidden_size ~activation ~params x =
    let dtype = dtype x in
    let get_param name = Ptree.Dict.get_tensor_exn params ~name dtype in

    let inter_weight = get_param "inter_weight" in
    let inter_bias = get_param "inter_bias" in
    let out_weight = get_param "out_weight" in
    let out_bias = get_param "out_bias" in

    (* First linear layer *)
    let h =
      add (matmul x inter_weight) (reshape [| 1; 1; n_inner |] inter_bias)
    in
    (* GELU activation *)
    let h =
      match activation with
      | `gelu -> gelu h
      | `gelu_new -> gelu_new h
      | `relu -> Activations.relu h
      | `swish -> Activations.swish h
    in
    (* Second linear layer *)
    add (matmul h out_weight) (reshape [| 1; 1; hidden_size |] out_bias)

  (* GPT-2 transformer block with pre-layer normalization *)
  let gpt2_block ~config ~params x =
    let dtype = dtype x in
    let get_param name = Ptree.Dict.get_tensor_exn params ~name dtype in

    let ln1_weight = get_param "attn_gamma" in
    let ln1_bias = get_param "attn_beta" in
    let ln2_weight = get_param "ffn_gamma" in
    let ln2_bias = get_param "ffn_beta" in

    (* Pre-layer norm for attention *)
    let normed =
      layer_norm ~gamma:ln1_weight ~beta:ln1_bias
        ~epsilon:config.layer_norm_epsilon x
    in

    (* Self-attention with residual *)
    let attn_out =
      causal_attention ~n_head:config.n_head ~hidden_size:config.n_embd ~params
        normed
    in
    let x = add x attn_out in

    (* Pre-layer norm for FFN *)
    let normed =
      layer_norm ~gamma:ln2_weight ~beta:ln2_bias
        ~epsilon:config.layer_norm_epsilon x
    in

    (* FFN with residual *)
    let n_inner = Option.value config.n_inner ~default:(4 * config.n_embd) in
    let ffn_out =
      mlp ~n_inner ~activation:config.activation_function
        ~hidden_size:config.n_embd ~params normed
    in
    add x ffn_out

  (* Stack of GPT-2 blocks *)
  let gpt2_transformer ~config ~layer_params x =
    (* Apply each transformer block sequentially *)
    let rec apply_layers x = function
      | [] -> x
      | params :: rest ->
          let x = gpt2_block ~config ~params x in
          apply_layers x rest
    in
    apply_layers x layer_params
end

let create ?(config = default_config) () =
  (* GPT-2 uses a custom architecture with causal attention *)
  (* We'll implement it as a custom module *)
  {
    Kaun.init =
      (fun ~rngs ~dtype ->
        (* Initialize embeddings *)
        let embeddings_layer = embeddings ~config () in
        let embeddings_params = Kaun.init embeddings_layer ~rngs ~dtype in

        (* Initialize transformer blocks *)
        let layer_params = List.init config.n_layer (fun _ -> []) in

        (* Initialize final layer norm *)
        let ln_f_gamma = Rune.ones dtype [| config.n_embd |] in
        let ln_f_beta = Rune.zeros dtype [| config.n_embd |] in

        (* Return params structure *)
        Ptree.List
          [
            embeddings_params;
            Ptree.List (List.map (fun p -> Ptree.Dict p) layer_params);
            Ptree.dict
              [
                ("gamma", Ptree.tensor ln_f_gamma);
                ("beta", Ptree.tensor ln_f_beta);
              ];
          ]);
    Kaun.apply =
      (fun params ~training:_ ?rngs:_ x ->
        (* Access sub-structures via explicit paths *)
        let embeddings_params =
          Ptree.get_exn ~path:(Ptree.Path.index 0 Ptree.Path.root) params
        in
        let layer_params_list =
          Ptree.get_exn ~path:(Ptree.Path.index 1 Ptree.Path.root) params
        in
        let ln_f_params =
          Ptree.get_exn ~path:(Ptree.Path.index 2 Ptree.Path.root) params
        in

        (* Apply embeddings *)
        let embeddings_layer = embeddings ~config () in
        let x =
          Kaun.apply embeddings_layer embeddings_params ~training:false x
        in

        (* Extract layer params as dict fields *)
        let layer_params =
          let items =
            Ptree.List.items_exn ~ctx:"gpt2: layer_params" layer_params_list
          in
          List.map
            (fun item ->
              Ptree.Dict.fields_exn ~ctx:"gpt2: layer_params item" item)
            items
        in

        (* Apply GPT-2 transformer blocks *)
        let x = Gpt2_block.gpt2_transformer ~config ~layer_params x in

        (* Apply final layer norm using path lookups *)
        let gamma =
          Ptree.get_tensor_exn
            ~path:(Ptree.Path.key "gamma" Ptree.Path.root)
            ln_f_params (Rune.dtype x)
        in
        let beta =
          Ptree.get_tensor_exn
            ~path:(Ptree.Path.key "beta" Ptree.Path.root)
            ln_f_params (Rune.dtype x)
        in
        layer_norm x ~gamma ~beta ~epsilon:config.layer_norm_epsilon);
  }

let from_pretrained ?(model_id = "gpt2") ?revision ?cache_config ~dtype () =
  (* Load config and weights from HuggingFace *)
  let cache_config =
    Option.value cache_config ~default:Kaun_huggingface.Config.default
  in
  let revision = Option.value revision ~default:Kaun_huggingface.Latest in

  (* Load config JSON from HuggingFace *)
  let config_json =
    match
      Kaun_huggingface.load_config ~config:cache_config ~revision ~model_id ()
    with
    | Cached json | Downloaded (json, _) -> json
  in

  (* Parse GPT-2 specific config *)
  let gpt2_config =
    let open Yojson.Safe.Util in
    {
      vocab_size = config_json |> member "vocab_size" |> to_int;
      n_positions = config_json |> member "n_positions" |> to_int;
      n_embd = config_json |> member "n_embd" |> to_int;
      n_layer = config_json |> member "n_layer" |> to_int;
      n_head = config_json |> member "n_head" |> to_int;
      n_inner = config_json |> member "n_inner" |> to_int_option;
      activation_function =
        (match
           config_json |> member "activation_function" |> to_string_option
         with
        | Some "gelu_new" -> `gelu_new
        | Some "gelu" -> `gelu
        | Some "relu" -> `relu
        | Some "swish" | Some "silu" -> `swish
        | _ -> `gelu_new);
      resid_pdrop =
        config_json |> member "resid_pdrop" |> to_float_option
        |> Option.value ~default:0.1;
      embd_pdrop =
        config_json |> member "embd_pdrop" |> to_float_option
        |> Option.value ~default:0.1;
      attn_pdrop =
        config_json |> member "attn_pdrop" |> to_float_option
        |> Option.value ~default:0.1;
      layer_norm_epsilon =
        config_json
        |> member "layer_norm_epsilon"
        |> to_float_option |> Option.value ~default:1e-5;
      initializer_range =
        config_json |> member "initializer_range" |> to_float_option
        |> Option.value ~default:0.02;
      scale_attn_weights =
        config_json
        |> member "scale_attn_weights"
        |> to_bool_option |> Option.value ~default:true;
      use_cache =
        config_json |> member "use_cache" |> to_bool_option
        |> Option.value ~default:true;
      scale_attn_by_inverse_layer_idx =
        config_json
        |> member "scale_attn_by_inverse_layer_idx"
        |> to_bool_option
        |> Option.value ~default:false;
      reorder_and_upcast_attn =
        config_json
        |> member "reorder_and_upcast_attn"
        |> to_bool_option
        |> Option.value ~default:false;
      bos_token_id = config_json |> member "bos_token_id" |> to_int_option;
      eos_token_id = config_json |> member "eos_token_id" |> to_int_option;
      pad_token_id = config_json |> member "pad_token_id" |> to_int_option;
    }
  in

  (* Load weights using HuggingFace infrastructure *)
  let hf_params =
    Kaun_huggingface.from_pretrained ~config:cache_config ~revision ~model_id ()
  in

  (* Map HuggingFace parameter names to our expected structure *)
  let map_huggingface_to_kaun hf_params =
    (* Flatten the nested HuggingFace structure *)
    let rec flatten_ptree prefix tree =
      match tree with
      | Ptree.Tensor tensor -> [ (prefix, tensor) ]
      | Ptree.List lst ->
          List.concat
            (List.mapi
               (fun i subtree ->
                 flatten_ptree (prefix ^ "." ^ string_of_int i) subtree)
               lst)
      | Ptree.Dict fields ->
          List.fold_left
            (fun acc (name, subtree) ->
              let new_prefix =
                if prefix = "" then name else prefix ^ "." ^ name
              in
              flatten_ptree new_prefix subtree @ acc)
            [] fields
    in

    let flat_params = flatten_ptree "" hf_params in
    let embeddings_params = ref [] in
    let decoder_layers = ref [] in
    let final_layer_norm_params = ref [] in
    let set_embedding params key tensor =
      Ptree.Dict.set key
        (Ptree.dict [ ("embedding", Ptree.Tensor tensor) ])
        params
    in

    List.iter
      (fun (hf_name, tensor) ->
        match hf_name with
        (* Embeddings *)
        | s when String.starts_with ~prefix:"wte.weight" s ->
            embeddings_params :=
              set_embedding !embeddings_params "token_embeddings" tensor
        | s when String.starts_with ~prefix:"wpe.weight" s ->
            embeddings_params :=
              set_embedding !embeddings_params "position_embeddings" tensor
        (* Transformer blocks *)
        | s when String.starts_with ~prefix:"h." s -> (
            let rest = String.sub s 2 (String.length s - 2) in
            match String.split_on_char '.' rest with
            | layer_idx :: params -> (
                let layer_idx_int = int_of_string layer_idx in
                let param_name = String.concat "." params in

                (* Ensure we have enough layers *)
                while List.length !decoder_layers <= layer_idx_int do
                  decoder_layers := !decoder_layers @ [ ref [] ]
                done;

                (* Get the layer params ref *)
                let layer_params = List.nth !decoder_layers layer_idx_int in

                (* Handle different parameter types *)
                match param_name with
                | "attn.c_attn.weight" ->
                    (* Keep combined QKV weight, will split after matmul *)
                    (* GPT-2 stores as [in_features, out_features] = [768, 2304] *)
                    layer_params :=
                      Ptree.Dict.set "qkv_weight" (Ptree.Tensor tensor)
                        !layer_params
                | "attn.c_attn.bias" ->
                    (* Keep combined QKV bias *)
                    layer_params :=
                      Ptree.Dict.set "qkv_bias" (Ptree.Tensor tensor)
                        !layer_params
                | "attn.c_proj.weight" ->
                    (* GPT-2 stores as [in, out] which is [768, 768] *)
                    (* We need [768, 768] for x[*,*,768] @ W -> [*,*,768] *)
                    (* So NO transpose needed! *)
                    layer_params :=
                      Ptree.Dict.set "attn_out_weight" (Ptree.Tensor tensor)
                        !layer_params
                | "attn.c_proj.bias" ->
                    layer_params :=
                      Ptree.Dict.set "attn_out_bias" (Ptree.Tensor tensor)
                        !layer_params
                | "mlp.c_fc.weight" ->
                    (* GPT-2 stores as [in, out] which is [768, 3072] *)
                    (* We need [768, 3072] for x[*,*,768] @ W -> [*,*,3072] *)
                    (* So NO transpose needed! *)
                    layer_params :=
                      Ptree.Dict.set "inter_weight" (Ptree.Tensor tensor)
                        !layer_params
                | "mlp.c_proj.weight" ->
                    (* GPT-2 stores as [in, out] which is [3072, 768] *)
                    (* We need [3072, 768] for h[*,*,3072] @ W -> [*,*,768] *)
                    (* So NO transpose needed! *)
                    layer_params :=
                      Ptree.Dict.set "out_weight" (Ptree.Tensor tensor)
                        !layer_params
                | "mlp.c_fc.bias" ->
                    layer_params :=
                      Ptree.Dict.set "inter_bias" (Ptree.Tensor tensor)
                        !layer_params
                | "mlp.c_proj.bias" ->
                    layer_params :=
                      Ptree.Dict.set "out_bias" (Ptree.Tensor tensor)
                        !layer_params
                | "ln_1.weight" ->
                    layer_params :=
                      Ptree.Dict.set "attn_gamma" (Ptree.Tensor tensor)
                        !layer_params
                | "ln_1.bias" ->
                    layer_params :=
                      Ptree.Dict.set "attn_beta" (Ptree.Tensor tensor)
                        !layer_params
                | "ln_2.weight" ->
                    layer_params :=
                      Ptree.Dict.set "ffn_gamma" (Ptree.Tensor tensor)
                        !layer_params
                | "ln_2.bias" ->
                    layer_params :=
                      Ptree.Dict.set "ffn_beta" (Ptree.Tensor tensor)
                        !layer_params
                | _ -> () (* Ignore other parameters like attn.bias *))
            | _ -> ())
        (* Final layer norm *)
        | s when String.starts_with ~prefix:"ln_f.weight" s ->
            final_layer_norm_params :=
              Ptree.Dict.set "gamma" (Ptree.Tensor tensor)
                !final_layer_norm_params
        | s when String.starts_with ~prefix:"ln_f.bias" s ->
            final_layer_norm_params :=
              Ptree.Dict.set "beta" (Ptree.Tensor tensor)
                !final_layer_norm_params
        | _ -> () (* Ignore other parameters *))
      flat_params;

    let ensure_embedding params key =
      match List.assoc_opt key params with
      | Some (Ptree.Dict fields) ->
          if
            not
              (List.exists
                 (fun (name, _) -> String.equal name "embedding")
                 fields)
          then failwith (key ^ " missing embedding field")
      | Some _ -> failwith (key ^ " is not a dict")
      | None -> failwith (key ^ " missing")
    in
    ensure_embedding !embeddings_params "token_embeddings";
    ensure_embedding !embeddings_params "position_embeddings";

    (* Build the final sequential structure *)
    let decoder_list = List.map (fun r -> Ptree.Dict !r) !decoder_layers in

    (* Add dropout placeholder for embeddings *)
    embeddings_params :=
      Ptree.Dict.set "dropout" (Ptree.List []) !embeddings_params;

    (* Create sequential structure: embeddings, decoder layers, final layer
       norm *)
    Ptree.List
      [
        Ptree.Dict !embeddings_params;
        Ptree.List decoder_list;
        Ptree.Dict !final_layer_norm_params;
      ]
  in

  let mapped_params = map_huggingface_to_kaun hf_params in
  let model = create ~config:gpt2_config () in

  { model; params = mapped_params; config = gpt2_config; dtype }

let forward gpt2 inputs ?(training = false) ?(output_hidden_states = false)
    ?(output_attentions = false) () =
  let { model; params; dtype = target_dtype; _ } = gpt2 in
  let { input_ids; attention_mask = _; position_ids = _ } = inputs in

  (* GPT-2 forward pass using the Kaun model *)
  let open Rune in
  (* Get dtype from params to maintain polymorphism *)
  let float_input = cast target_dtype input_ids in

  (* Apply the model using Kaun *)
  let model_output = Kaun.apply model params ~training float_input in

  (* The model output is the final hidden state *)
  let last_hidden_state = model_output in

  (* For output_hidden_states and output_attentions, we would need to modify the
     model architecture to return intermediate values. For now, return minimal
     info *)
  let hidden_states =
    if output_hidden_states then Some [ last_hidden_state ] else None
  in
  let attentions = if output_attentions then None else None in

  (* Return structured output *)
  { last_hidden_state; hidden_states; attentions }

(* Language Modeling Head *)

module For_causal_lm = struct
  let create ?(config = default_config) () =
    let open Kaun.Layer in
    sequential
      [
        (* GPT-2 base model *)
        create ~config ();
        (* LM head: project to vocabulary *)
        linear ~in_features:config.n_embd ~out_features:config.vocab_size ();
      ]

  let forward ~model ~params ~compute_dtype ~input_ids ?attention_mask:_
      ?position_ids:_ ?labels ~training () =
    let open Rune in
    let dtype = compute_dtype in
    let float_input = cast dtype input_ids in
    let hidden_states = Kaun.apply model params ~training float_input in

    let find_token_embedding_weight () =
      let open Ptree in
      let try_paths =
        [
          "[0][0].token_embeddings.embedding";
          "[0].token_embeddings.embedding";
          "[0][0].token_embeddings";
          "[0].token_embeddings";
        ]
      in
      let rec loop = function
        | [] -> None
        | p :: ps -> (
            match get_tensor ~path:(Path.of_string p) params dtype with
            | Some t -> Some t
            | None -> (
                (* get_tensor may raise on dtype mismatch; treat as miss *)
                try Some (get_tensor_exn ~path:(Path.of_string p) params dtype)
                with _ -> loop ps))
      in
      loop try_paths
    in

    let logits =
      match find_token_embedding_weight () with
      | Some wte ->
          let wte_t = transpose ~axes:[ 1; 0 ] wte in
          matmul hidden_states wte_t
      | None -> hidden_states
    in

    let loss =
      match labels with
      | Some labels ->
          let batch_size = (shape labels).(0) in
          let seq_length = (shape labels).(1) in
          let vocab_size = (shape logits).(2) in
          let shift_logits = slice [ A; R (0, seq_length - 1); A ] logits in
          let shift_labels = slice [ A; R (1, seq_length); A ] labels in
          let flat_logits =
            Rune.reshape
              [| batch_size * (seq_length - 1); vocab_size |]
              shift_logits
          in
          let flat_labels =
            Rune.reshape [| batch_size * (seq_length - 1) |] shift_labels
          in
          Some
            (Kaun.Loss.softmax_cross_entropy_with_indices flat_logits
               flat_labels)
      | None -> None
    in

    (logits, loss)
end

(* Utilities *)

let parse_gpt2_config json =
  (* Parse GPT-2 specific configuration from HuggingFace JSON *)
  let open Yojson.Safe.Util in
  {
    vocab_size = json |> member "vocab_size" |> to_int;
    n_positions = json |> member "n_positions" |> to_int;
    n_embd = json |> member "n_embd" |> to_int;
    n_layer = json |> member "n_layer" |> to_int;
    n_head = json |> member "n_head" |> to_int;
    n_inner = json |> member "n_inner" |> to_int_option;
    activation_function =
      (match json |> member "activation_function" |> to_string_option with
      | Some "gelu_new" -> `gelu_new
      | Some "gelu" -> `gelu
      | Some "relu" -> `relu
      | Some "swish" | Some "silu" -> `swish
      | _ -> `gelu_new);
    resid_pdrop =
      json |> member "resid_pdrop" |> to_float_option
      |> Option.value ~default:0.1;
    embd_pdrop =
      json |> member "embd_pdrop" |> to_float_option
      |> Option.value ~default:0.1;
    attn_pdrop =
      json |> member "attn_pdrop" |> to_float_option
      |> Option.value ~default:0.1;
    layer_norm_epsilon =
      json
      |> member "layer_norm_epsilon"
      |> to_float_option |> Option.value ~default:1e-5;
    initializer_range =
      json |> member "initializer_range" |> to_float_option
      |> Option.value ~default:0.02;
    scale_attn_weights =
      json
      |> member "scale_attn_weights"
      |> to_bool_option |> Option.value ~default:true;
    use_cache =
      json |> member "use_cache" |> to_bool_option |> Option.value ~default:true;
    scale_attn_by_inverse_layer_idx =
      json
      |> member "scale_attn_by_inverse_layer_idx"
      |> to_bool_option
      |> Option.value ~default:false;
    reorder_and_upcast_attn =
      json
      |> member "reorder_and_upcast_attn"
      |> to_bool_option
      |> Option.value ~default:false;
    bos_token_id = json |> member "bos_token_id" |> to_int_option;
    eos_token_id = json |> member "eos_token_id" |> to_int_option;
    pad_token_id = json |> member "pad_token_id" |> to_int_option;
  }

let num_parameters params =
  let tensors = Ptree.flatten_with_paths params in
  List.fold_left
    (fun acc (_, t) -> acc + Array.fold_left ( * ) 1 (Ptree.Tensor.shape t))
    0 tensors

let parameter_stats params =
  let total_params = num_parameters params in
  let total_bytes = total_params * 4 in
  (* Assuming float32 *)
  Printf.sprintf "GPT-2 parameters: %d (%.2f MB)" total_params
    (float_of_int total_bytes /. 1024. /. 1024.)

(* Common GPT-2 model configurations *)
let load_gpt2_small ~dtype () = from_pretrained ~model_id:"gpt2" ~dtype ()

let load_gpt2_medium ~dtype () =
  from_pretrained ~model_id:"gpt2-medium" ~dtype ()

let load_gpt2_large ~dtype () = from_pretrained ~model_id:"gpt2-large" ~dtype ()
let load_gpt2_xl ~dtype () = from_pretrained ~model_id:"gpt2-xl" ~dtype ()