gemma.cpp/gemma/tokenizer.cc

99 lines
3.1 KiB
C++

// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// 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
//
// https://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 "gemma/tokenizer.h"
#include <stdio.h>
#include <memory>
#include <string>
#include <vector>
#include "compression/io.h" // Path
#include "hwy/base.h"
#include "hwy/profiler.h"
// copybara:import_next_line:sentencepiece
#include "src/sentencepiece_processor.h"
namespace gcpp {
// Set this to true to debug tokenizer tokens.
constexpr bool kShowTokenization = false;
class GemmaTokenizer::Impl {
public:
Impl() = default;
explicit Impl(const Path& tokenizer_path) {
PROFILER_ZONE("Startup.tokenizer");
spp_ = std::make_unique<sentencepiece::SentencePieceProcessor>();
if (!spp_->Load(tokenizer_path.path).ok()) {
HWY_ABORT("Failed to load the tokenizer file.");
}
}
bool Encode(const std::string& input,
std::vector<std::string>* pieces) const {
return spp_ && spp_->Encode(input, pieces).ok();
}
bool Encode(const std::string& input, std::vector<int>* ids) const {
if constexpr (kShowTokenization) {
bool is_ok = spp_ && spp_->Encode(input, ids).ok();
for (int i = 0; i < static_cast<int>(ids->size()); i++) {
fprintf(stderr, "%3d: %d\n", i, (*ids)[i]);
}
return is_ok;
} else {
return spp_ && spp_->Encode(input, ids).ok();
}
}
// Given a sequence of ids, decodes it into a detokenized output.
bool Decode(const std::vector<int>& ids, std::string* detokenized) const {
return spp_ && spp_->Decode(ids, detokenized).ok();
}
private:
std::unique_ptr<sentencepiece::SentencePieceProcessor> spp_;
};
GemmaTokenizer::GemmaTokenizer(const Path& tokenizer_path) {
impl_ = std::make_unique<Impl>(tokenizer_path);
}
// Default suffices, but they must be defined after GemmaTokenizer::Impl.
GemmaTokenizer::GemmaTokenizer() = default;
GemmaTokenizer::~GemmaTokenizer() = default;
GemmaTokenizer::GemmaTokenizer(GemmaTokenizer&& other) = default;
GemmaTokenizer& GemmaTokenizer::operator=(GemmaTokenizer&& other) = default;
bool GemmaTokenizer::Encode(const std::string& input,
std::vector<std::string>* pieces) const {
return impl_->Encode(input, pieces);
}
bool GemmaTokenizer::Encode(const std::string& input,
std::vector<int>* ids) const {
return impl_->Encode(input, ids);
}
// Given a sequence of ids, decodes it into a detokenized output.
bool GemmaTokenizer::Decode(const std::vector<int>& ids,
std::string* detokenized) const {
return impl_->Decode(ids, detokenized);
}
} // namespace gcpp