Source file wordpiece.ml

(** WordPiece tokenization implementation *)

exception Error of string

type vocab = (string, int) Hashtbl.t
type vocab_r = (int, string) Hashtbl.t
type token = { id : int; value : string; offsets : int * int }

type config = {
  vocab : vocab;
  unk_token : string;
  continuing_subword_prefix : string;
  max_input_chars_per_word : int;
}

type t = {
  mutable vocab : vocab;
  mutable vocab_r : vocab_r;
  unk_token : string;
  mutable continuing_subword_prefix : string;
  max_input_chars_per_word : int;
}

let create_internal vocab unk_token continuing_subword_prefix
    max_input_chars_per_word =
  let vocab_r = Hashtbl.create (Hashtbl.length vocab) in
  Hashtbl.iter (fun k v -> Hashtbl.add vocab_r v k) vocab;
  (* Only raise error if vocabulary is non-empty but missing UNK token *)
  if Hashtbl.length vocab > 0 && not (Hashtbl.mem vocab unk_token) then
    raise (Error "WordPiece error: Missing [UNK] token from the vocabulary");
  {
    vocab;
    vocab_r;
    unk_token;
    continuing_subword_prefix;
    max_input_chars_per_word;
  }

let create (cfg : config) =
  create_internal cfg.vocab cfg.unk_token cfg.continuing_subword_prefix
    cfg.max_input_chars_per_word

let default () = create_internal (Hashtbl.create 0) "[UNK]" "##" 100

let read_file ~vocab_file =
  let vocab = Hashtbl.create 10000 in
  let ic = open_in vocab_file in
  let index = ref 0 in
  (try
     while true do
       let line = input_line ic in
       let token = String.trim line in
       if token <> "" then (
         Hashtbl.add vocab token !index;
         incr index)
     done
   with End_of_file -> ());
  close_in ic;
  vocab

let read_bytes bytes =
  let vocab = Hashtbl.create 10000 in
  let str = Bytes.to_string bytes in
  let lines = String.split_on_char '\n' str in
  List.iteri
    (fun index line ->
      let token = String.trim line in
      if token <> "" then Hashtbl.add vocab token index)
    lines;
  vocab

let from_file ~vocab_file =
  let vocab = read_file ~vocab_file in
  (* Use default values for BERT-style WordPiece *)
  create_internal vocab "[UNK]" "##" 100

let from_file_with_config ~vocab_file ~unk_token ~continuing_subword_prefix
    ~max_input_chars_per_word =
  let vocab = read_file ~vocab_file in
  create_internal vocab unk_token continuing_subword_prefix
    max_input_chars_per_word

let tokenize model sequence =
  (* Handle empty vocabulary case *)
  if Hashtbl.length model.vocab = 0 then []
  else
    let char_count =
      let decoder = Uutf.decoder (`String sequence) in
      let count = ref 0 in
      let rec loop () =
        match Uutf.decode decoder with
        | `Uchar _ ->
            incr count;
            loop ()
        | `End -> !count
        | `Malformed _ ->
            incr count;
            loop ()
        | `Await -> assert false
      in
      loop ()
    in
    if char_count > model.max_input_chars_per_word then
      let id = Hashtbl.find model.vocab model.unk_token in
      [ { id; value = model.unk_token; offsets = (0, String.length sequence) } ]
    else
      let rec tokenize_greedy start acc =
        if start >= String.length sequence then List.rev acc
        else
          let rec find_longest_match end_pos =
            if end_pos <= start then None
            else
              let substr = String.sub sequence start (end_pos - start) in
              let token_str =
                if start > 0 then model.continuing_subword_prefix ^ substr
                else substr
              in
              match Hashtbl.find_opt model.vocab token_str with
              | Some id ->
                  Some { id; value = token_str; offsets = (start, end_pos) }
              | None ->
                  let new_end =
                    let rec find_char_start pos =
                      if pos <= start then start
                      else if Char.code sequence.[pos - 1] land 0xC0 <> 0x80
                      then pos - 1
                      else find_char_start (pos - 1)
                    in
                    find_char_start end_pos
                  in
                  if new_end <= start then None else find_longest_match new_end
          in
          match find_longest_match (String.length sequence) with
          | Some token -> tokenize_greedy (snd token.offsets) (token :: acc)
          | None ->
              let id = Hashtbl.find model.vocab model.unk_token in
              [
                {
                  id;
                  value = model.unk_token;
                  offsets = (0, String.length sequence);
                };
              ]
      in
      tokenize_greedy 0 []

let token_to_id model token = Hashtbl.find_opt model.vocab token
let id_to_token model id = Hashtbl.find_opt model.vocab_r id
let get_vocab model = Hashtbl.fold (fun k v acc -> (k, v) :: acc) model.vocab []
let get_vocab_size model = Hashtbl.length model.vocab
let get_unk_token model = model.unk_token
let get_continuing_subword_prefix model = model.continuing_subword_prefix
let get_max_input_chars_per_word model = model.max_input_chars_per_word

let save model ~path ?name () =
  let vocab_file =
    match name with
    | Some n -> Filename.concat path (n ^ "-vocab.txt")
    | None -> Filename.concat path "vocab.txt"
  in
  let vocab_list =
    Hashtbl.fold (fun k v acc -> (v, k) :: acc) model.vocab []
    |> List.sort compare
    |> List.map (fun (_, k) -> k)
  in
  let oc = open_out vocab_file in
  List.iter
    (fun token ->
      output_string oc token;
      output_char oc '\n')
    vocab_list;
  close_out oc;
  vocab_file

let from_bpe bpe =
  let vocab = Hashtbl.create (Bpe.get_vocab_size bpe) in
  List.iter (fun (k, id) -> Hashtbl.add vocab k id) (Bpe.get_vocab bpe);
  let unk_token =
    match Bpe.get_unk_token bpe with Some u -> u | None -> "[UNK]"
  in
  let continuing_subword_prefix =
    match Bpe.get_continuing_subword_prefix bpe with
    | Some p -> p
    | None -> "##"
  in
  create_internal vocab unk_token continuing_subword_prefix 100

(* Save reference to internal create before Builder shadows it *)
let create_wordpiece = create_internal

(** Builder module *)
module Builder = struct
  type builder = {
    mutable files : string option;
    mutable vocab : vocab;
    mutable unk_token : string;
    mutable continuing_subword_prefix : string;
    mutable max_input_chars_per_word : int;
  }

  let create () =
    {
      files = None;
      vocab = Hashtbl.create 0;
      unk_token = "[UNK]";
      continuing_subword_prefix = "##";
      max_input_chars_per_word = 100;
    }

  let files builder f =
    builder.files <- Some f;
    builder

  let vocab builder v =
    builder.vocab <- v;
    builder

  let unk_token builder token =
    builder.unk_token <- token;
    builder

  let continuing_subword_prefix builder prefix =
    builder.continuing_subword_prefix <- prefix;
    builder

  let max_input_chars_per_word builder max =
    builder.max_input_chars_per_word <- max;
    builder

  let build builder =
    let vocab =
      match builder.files with
      | None -> builder.vocab
      | Some f -> read_file ~vocab_file:f
    in
    create_wordpiece vocab builder.unk_token builder.continuing_subword_prefix
      builder.max_input_chars_per_word
end

(** Serialization *)
let to_yojson model =
  let vocab_list =
    Hashtbl.fold (fun k v acc -> (v, k) :: acc) model.vocab []
    |> List.sort compare
    |> List.map (fun (_, k) -> (k, `Int (Hashtbl.find model.vocab k)))
  in
  `Assoc
    [
      ("type", `String "WordPiece");
      ("unk_token", `String model.unk_token);
      ("continuing_subword_prefix", `String model.continuing_subword_prefix);
      ("max_input_chars_per_word", `Int model.max_input_chars_per_word);
      ("vocab", `Assoc vocab_list);
    ]

let of_yojson json =
  match json with
  | `Assoc fields ->
      let get_field name =
        List.assoc_opt name fields |> function
        | Some v -> v
        | None -> raise (Error ("Missing field: " ^ name))
      in
      let () =
        match get_field "type" with
        | `String "WordPiece" -> ()
        | _ -> raise (Error "Invalid type")
        | exception _ -> ()
      in
      let unk_token =
        match get_field "unk_token" with
        | `String s -> s
        | _ -> raise (Error "Invalid unk_token")
      in
      let continuing_subword_prefix =
        match get_field "continuing_subword_prefix" with
        | `String s -> s
        | _ -> raise (Error "Invalid continuing_subword_prefix")
      in
      let max_input_chars_per_word =
        match get_field "max_input_chars_per_word" with
        | `Int i -> i
        | _ -> raise (Error "Invalid max_input_chars_per_word")
      in
      let vocab_json = get_field "vocab" in
      let vocab =
        match vocab_json with
        | `Assoc pairs ->
            let h = Hashtbl.create (List.length pairs) in
            List.iter
              (fun (k, v) ->
                match v with
                | `Int id -> Hashtbl.add h k id
                | _ -> raise (Error "Invalid vocab entry"))
              pairs;
            h
        | _ -> raise (Error "Invalid vocab")
      in
      create_internal vocab unk_token continuing_subword_prefix
        max_input_chars_per_word
  | _ -> raise (Error "Invalid JSON structure")

let from_bytes bytes =
  let str = Bytes.to_string bytes in
  of_yojson (Yojson.Basic.from_string str)

(** Trainer module *)
module Trainer = struct
  type trainer_config = {
    min_frequency : int;
    vocab_size : int;
    show_progress : bool;
    special_tokens : string list;
    limit_alphabet : int option;
    initial_alphabet : char list;
    continuing_subword_prefix : string;
    end_of_word_suffix : string option;
  }

  type trainer = {
    config : trainer_config; [@warning "-69"]
    bpe_config : Bpe.Trainer.trainer_config; [@warning "-69"]
    bpe_trainer : Bpe.Trainer.trainer;
  }

  let default_config =
    {
      min_frequency = 0;
      vocab_size = 30000;
      show_progress = true;
      special_tokens = [ "[UNK]"; "[CLS]"; "[SEP]"; "[PAD]"; "[MASK]" ];
      limit_alphabet = None;
      initial_alphabet = [];
      continuing_subword_prefix = "##";
      end_of_word_suffix = None;
    }

  let create config =
    let bpe_config =
      {
        Bpe.Trainer.min_frequency = config.min_frequency;
        vocab_size = config.vocab_size;
        show_progress = config.show_progress;
        special_tokens = config.special_tokens;
        limit_alphabet = config.limit_alphabet;
        initial_alphabet = config.initial_alphabet;
        continuing_subword_prefix = Some config.continuing_subword_prefix;
        end_of_word_suffix = config.end_of_word_suffix;
        max_token_length = None;
      }
    in
    { config; bpe_config; bpe_trainer = Bpe.Trainer.create bpe_config }

  let feed trainer texts = Bpe.Trainer.feed trainer.bpe_trainer texts

  let train trainer model =
    let bpe_model = Bpe.default () in
    let special_tokens = Bpe.Trainer.train trainer.bpe_trainer bpe_model in
    let new_model = from_bpe bpe_model in
    model.vocab <- new_model.vocab;
    model.vocab_r <- new_model.vocab_r;
    model.continuing_subword_prefix <- new_model.continuing_subword_prefix;
    special_tokens
end