cleaner inheritence

marian-nmt · Jan 24, 2018 · 6e223c2 · 6e223c2
1 parent a1e22b1
commit 6e223c2
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 11 deletions.
diff --git a/src/data/corpus.cpp b/src/data/corpus.cpp
@@ -134,7 +134,7 @@ Corpus::Corpus(std::vector<std::string> paths,
                std::vector<Ptr<Vocab>> vocabs,
                Ptr<Config> options,
                size_t maxLength)
-    : DatasetBase(paths),
+    : CorpusBase(paths),
       options_(options),
       vocabs_(vocabs),
       maxLength_(maxLength ? maxLength : options_->get<size_t>("max-length")),

diff --git a/src/data/corpus.h b/src/data/corpus.h
@@ -206,7 +206,14 @@ class CorpusBatch : public Batch {
 
 class CorpusIterator;
 
-typedef DatasetBase<SentenceTuple, CorpusIterator, CorpusBatch> CorpusBase;
+class CorpusBase : public DatasetBase<SentenceTuple, CorpusIterator, CorpusBatch> {
+public:
+
+  CorpusBase() : DatasetBase() {}
+  CorpusBase(std::vector<std::string> paths) : DatasetBase(paths) {}
+
+  virtual std::vector<Ptr<Vocab>>& getVocabs() = 0;
+};
 
 class CorpusIterator
     : public boost::iterator_facade<CorpusIterator,

diff --git a/src/data/corpus_sqlite.cpp b/src/data/corpus_sqlite.cpp
@@ -113,7 +113,7 @@ CorpusSQLite::CorpusSQLite(std::vector<std::string> paths,
                            std::vector<Ptr<Vocab>> vocabs,
                            Ptr<Config> options,
                            size_t maxLength)
-    : DatasetBase(paths),
+    : CorpusBase(paths),
       options_(options),
       vocabs_(vocabs),
       maxLength_(maxLength ? maxLength : options_->get<size_t>("max-length")),

diff --git a/src/training/training.h b/src/training/training.h
@@ -43,14 +43,7 @@ class Train : public ModelTask {
     if((options_->has("valid-sets") || options_->has("valid-script-path"))
        && options_->get<size_t>("valid-freq") > 0) {
 
-      // @TODO: solve this with better polymorphism
-      std::vector<Ptr<Vocab>> vocabs;
-      if(options_->get<bool>("sqlite"))
-        vocabs = std::static_pointer_cast<CorpusSQLite>(dataset)->getVocabs();
-      else
-        vocabs = std::static_pointer_cast<Corpus>(dataset)->getVocabs();
-
-      for(auto validator : Validators(vocabs, options_))
+      for(auto validator : Validators(dataset->getVocabs(), options_))
         scheduler->addValidator(validator);
     }