// Copyright 2016 Google Inc.

//

// Licensed under the Apache License, Version 2.0 (the "License");

// you may not use this file except in compliance with the License.

// You may obtain a copy of the License at

//

//     http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing, software

// distributed under the License is distributed on an "AS IS" BASIS,

// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

// See the License for the specific language governing permissions and

// limitations under the License.!

#include "model_interface.h"

#include <algorithm>

#include "sentencepiece_model.pb.h"

#include "third_party/absl/strings/str_format.h"

#include "util.h"

namespace sentencepiece {

ModelInterface::ModelInterface(const ModelProto &model_proto)

    : model_proto_(&model_proto), status_(util::OkStatus()) {}

ModelInterface::~ModelInterface() {}

#define RETURN_PIECE(name, default_value)                                \

  if (model_proto_->trainer_spec().name().empty()) return default_value; \

  return model_proto_->trainer_spec().name();

absl::string_view ModelInterface::unk_piece() const {

  RETURN_PIECE(unk_piece, "<unk>");

}

absl::string_view ModelInterface::bos_piece() const {

  RETURN_PIECE(bos_piece, "<s>");

}

absl::string_view ModelInterface::eos_piece() const {

  RETURN_PIECE(eos_piece, "</s>");

}

absl::string_view ModelInterface::pad_piece() const {

  RETURN_PIECE(pad_piece, "<pad>");

}

#undef RETURN_PIECE

int ModelInterface::PieceToId(absl::string_view piece) const {

  if (auto it = reserved_id_map_.find(piece); it != reserved_id_map_.end()) {

    return it->second;

  }

  if (auto it = pieces_.find(piece); it != pieces_.end()) {

    return it->second;

  }

  return unk_id_;

}

void ModelInterface::InitializePieces() {

  pieces_.clear();

  reserved_id_map_.clear();

  unk_id_ = -1;

  std::set<absl::string_view> user_defined_symbols;

  std::vector<bool> byte_found(256, false);

  int pieces_size = 0;

  int reserved_id_map_size = 0;

  for (int i = 0; i < model_proto_->pieces_size(); ++i) {

    const auto &sp = model_proto_->pieces(i);

    static constexpr size_t kMaxPieceSize = 8192;

    if (sp.piece().size() >= kMaxPieceSize) {

      status_ = util::InternalError("piece size must be less than 8k.");

      return;

    }

    const bool is_normal_piece =

        (sp.type() == ModelProto::SentencePiece::NORMAL ||

         sp.type() == ModelProto::SentencePiece::USER_DEFINED ||

         sp.type() == ModelProto::SentencePiece::UNUSED);

    if (is_normal_piece) {

      ++pieces_size;

    } else {

      ++reserved_id_map_size;

    }

  }

  pieces_.reserve(pieces_size);

  reserved_id_map_.reserve(reserved_id_map_size);

  for (int i = 0; i < model_proto_->pieces_size(); ++i) {

    const auto &sp = model_proto_->pieces(i);

    if (sp.piece().empty()) {

      status_ = util::InternalError("piece must not be empty.");

      return;

    }

    if (sp.piece().find('\0') != absl::string_view::npos) {

      status_ = util::InternalError("piece must not include null character.");

      return;

    }

    const bool is_normal_piece =

        (sp.type() == ModelProto::SentencePiece::NORMAL ||

         sp.type() == ModelProto::SentencePiece::USER_DEFINED ||

         sp.type() == ModelProto::SentencePiece::UNUSED);

    if (!port::InsertIfNotPresent(

            is_normal_piece ? &pieces_ : &reserved_id_map_, sp.piece(), i)) {

      status_ = util::InternalError(sp.piece() + " is already defined.");

      return;

    }

    if (sp.type() == ModelProto::SentencePiece::USER_DEFINED) {

      user_defined_symbols.insert(sp.piece());

    }

    if (sp.type() == ModelProto::SentencePiece::UNKNOWN) {

      if (unk_id_ >= 0) {

        status_ = util::InternalError("unk is already defined.");

        return;

      }

      unk_id_ = i;

    }

    if (sp.type() == ModelProto::SentencePiece::BYTE) {

      if (!model_proto_->trainer_spec().byte_fallback()) {

        status_ =

            util::InternalError("byte piece " + sp.piece() +

                                " is found although `byte_fallback` is false.");

        return;

      }

      const int byte = PieceToByte(sp.piece());

      if (0 <= byte && byte < 256) {

        byte_found[byte] = true;

      } else {

        status_ =

            util::InternalError("byte piece " + sp.piece() + " is invalid.");

        return;

      }

    }

  }

  if (unk_id_ == -1) {

    status_ = util::InternalError("unk is not defined.");

    return;

  }

  if (model_proto_->trainer_spec().byte_fallback()) {

    // Checks that there are 256 byte pieces.

    if (std::find(byte_found.begin(), byte_found.end(), false) !=

        byte_found.end()) {

      status_ = util::InternalError(

          "there are not 256 byte pieces although `byte_fallback` is true.");

      return;

    }

  }

  matcher_ = std::make_unique<normalizer::PrefixMatcher>(user_defined_symbols);

}

std::vector<absl::string_view> SplitIntoWords(absl::string_view text,

                                              bool treat_ws_as_suffix,

                                              bool allow_ws_only_pieces) {

  const char *begin = text.data();

  const char *end = text.data() + text.size();

  // Space symbol (U+2581)

  constexpr absl::string_view kSpaceSymbol = "\xe2\x96\x81";

  bool in_ws_sequence = false;

  std::vector<absl::string_view> result;

  if (treat_ws_as_suffix) {  // put ws tokens at the end of non-ws sequences.

    if (begin < end) result.emplace_back(begin, 0);

    while (begin < end) {

      const int mblen =

          std::min<int>(string_util::OneCharLen(begin), end - begin);

      const bool is_ws = absl::string_view(begin, mblen) == kSpaceSymbol;

      if (is_ws) {  // keep track of sequences consecutive ws tokens.

        in_ws_sequence = true;

      } else if (in_ws_sequence) {

        if (allow_ws_only_pieces) result.emplace_back(begin, 0);

        in_ws_sequence = false;

      }

      result.back() =

          absl::string_view(result.back().data(), result.back().size() + mblen);

      begin += mblen;

      if (begin < end && is_ws && !allow_ws_only_pieces)

        result.emplace_back(begin, 0);

    }

  } else {

    while (begin < end) {

      const int mblen =

          std::min<int>(string_util::OneCharLen(begin), end - begin);

      bool is_ws = absl::string_view(begin, mblen) == kSpaceSymbol;

      // if is whitespace (and not in sequence if allow_ws_only_pieces is True)

      if (begin == text.data() ||

          (is_ws && (!in_ws_sequence || !allow_ws_only_pieces))) {

        result.emplace_back(begin, 0);  // add empty string piece.

        in_ws_sequence = true;

      }

      if (in_ws_sequence && !is_ws) in_ws_sequence = false;

      result.back() =

          absl::string_view(result.back().data(), result.back().size() + mblen);

      begin += mblen;

    }

  }

  return result;

}

const std::string &ByteToPiece(unsigned char c) {

  static const std::vector<std::string> *const kBytePieces = []() {

    auto *v = new std::vector<std::string>(256);

    for (int i = 0; i < 256; ++i) {

      (*v)[i] = absl::StrFormat("<0x%02X>", i);

    }

    return v;

  }();

  return (*kBytePieces)[c];

}

int PieceToByte(absl::string_view piece) {

  using PieceToByteMap = absl::flat_hash_map<absl::string_view, unsigned char>;

  static const auto *const kMap = []() -> PieceToByteMap * {

    auto *m = new PieceToByteMap();

    for (int i = 0; i < 256; ++i) {

      (*m)[ByteToPiece(i)] = i;

    }

    return m;

  }();

  if (const auto it = kMap->find(piece); it != kMap->end()) {

    return it->second;

  }

  return -1;

}

}  // namespace sentencepiece