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"));