diff options
Diffstat (limited to 'src/trainer_interface.h')
-rw-r--r-- | src/trainer_interface.h | 14 |
1 files changed, 7 insertions, 7 deletions
diff --git a/src/trainer_interface.h b/src/trainer_interface.h index 6cd2469..552b206 100644 --- a/src/trainer_interface.h +++ b/src/trainer_interface.h @@ -88,13 +88,13 @@ class TrainerInterface { virtual ~TrainerInterface(); - virtual void SetSentenceIterator(SentenceIterator *sentence_iterator) { + // Loads sentence from `sentence_iterator` and stores the model + // to `output_model_proto`. + virtual util::Status Train(SentenceIterator *sentence_iterator, + ModelProto *output_model_proto) { sentence_iterator_ = sentence_iterator; - } - - virtual void SetOutputSerializedModelProto( - std::string *serialized_model_proto) { - serialized_model_proto_ = serialized_model_proto; + output_model_proto_ = output_model_proto; + return Train(); } virtual util::Status Train() { return status(); } @@ -158,7 +158,7 @@ class TrainerInterface { SentenceIterator *sentence_iterator_ = nullptr; // Emits model to this proto instead of file. - std::string *serialized_model_proto_ = nullptr; + ModelProto *output_model_proto_ = nullptr; private: // Serialize final_pieces_ to |model_proto|. |