diff options
Diffstat (limited to 'src')
-rw-r--r-- | src/sentencepiece_processor.cc | 12 |
1 files changed, 10 insertions, 2 deletions
diff --git a/src/sentencepiece_processor.cc b/src/sentencepiece_processor.cc index 817219f..7abb848 100644 --- a/src/sentencepiece_processor.cc +++ b/src/sentencepiece_processor.cc @@ -42,9 +42,13 @@ bool SentencePieceProcessor::Load(const std::string &filename) { return false; } + return Load(ifs); +} + +bool SentencePieceProcessor::Load(std::istream &istream) { model_proto_ = port::MakeUnique<ModelProto>(); - if (!model_proto_->ParseFromIstream(&ifs)) { - LOG(WARNING) << "Model file is broken: " << filename; + if (!model_proto_->ParseFromIstream(&istream)) { + LOG(WARNING) << "Model file is broken"; return false; } @@ -59,6 +63,10 @@ void SentencePieceProcessor::LoadOrDie(const std::string &filename) { CHECK(Load(filename)) << "failed to load model: " << filename; } +void SentencePieceProcessor::LoadOrDie(std::istream &istream) { + CHECK(Load(istream)) << "failed to load model"; +} + void SentencePieceProcessor::SetEncodeExtraOptions( const std::string &extra_options) { encode_extra_options_ = ParseExtraOptions(extra_options); |