diff options
Diffstat (limited to 'src/trainer_interface.cc')
-rw-r--r-- | src/trainer_interface.cc | 18 |
1 files changed, 7 insertions, 11 deletions
diff --git a/src/trainer_interface.cc b/src/trainer_interface.cc index 37f7003..5cdb300 100644 --- a/src/trainer_interface.cc +++ b/src/trainer_interface.cc @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License.! -#include "trainer_interface.h" - #include <cstdlib> #include <memory> #include <set> @@ -34,6 +32,7 @@ #include "third_party/absl/strings/str_format.h" #include "third_party/absl/strings/str_join.h" #include "third_party/absl/strings/str_split.h" +#include "trainer_interface.h" #include "unicode_script.h" #include "util.h" @@ -50,7 +49,6 @@ const char TrainerInterface::kUPPBoundaryStr[] = "\t"; namespace { util::Status VerifySpec(const TrainerSpec &trainer_spec) { - // CHECK_OR_RETURN(!trainer_spec.model_prefix().empty()); CHECK_GT_OR_RETURN(trainer_spec.vocab_size(), 0); if (trainer_spec.model_type() == TrainerSpec::UNIGRAM || @@ -313,10 +311,10 @@ util::Status TrainerInterface::LoadSentences() { (sentence_iterator_ == nullptr && !trainer_spec_.input().empty())) << "SentenceIterator and trainer_spec.input() must be exclusive."; - CHECK_OR_RETURN((serialized_model_proto_ != nullptr && - trainer_spec_.model_prefix().empty()) || - (serialized_model_proto_ == nullptr && - !trainer_spec_.model_prefix().empty())) + CHECK_OR_RETURN( + (output_model_proto_ != nullptr && + trainer_spec_.model_prefix().empty()) || + (output_model_proto_ == nullptr && !trainer_spec_.model_prefix().empty())) << "ModelProto and trainer_spec.model_prefix() must be exclusive."; const bool is_tsv = trainer_spec_.input_format() == "tsv"; @@ -647,10 +645,8 @@ util::Status TrainerInterface::SaveVocab(absl::string_view filename) const { } util::Status TrainerInterface::Save() const { - if (serialized_model_proto_) { - ModelProto model_proto; - RETURN_IF_ERROR(Serialize(&model_proto)); - *serialized_model_proto_ = model_proto.SerializeAsString(); + if (output_model_proto_) { + RETURN_IF_ERROR(Serialize(output_model_proto_)); } else { RETURN_IF_ERROR(SaveModel(trainer_spec_.model_prefix() + ".model")); RETURN_IF_ERROR(SaveVocab(trainer_spec_.model_prefix() + ".vocab")); |