return pair

Author: js1010

cpp/include/ioutils.hpp

@@ -32,7 +32,7 @@ class IoUtils {
   ~IoUtils();
   bool Init(std::string opt_path);
   int LoadStreamFile(std::string filepath);
-  int ReadStreamForVocab(int num_lines);
+  std::pair<int, bool> ReadStreamForVocab(int num_lines);
   void GetWordVocab(int min_count);
  private:
   void ParseLine(std::string line, std::vector<std::string>& line_vec);

cpp/src/ioutils.cc

@@ -59,20 +59,25 @@ int IoUtils::LoadStreamFile(std::string filepath) {
   return count;
 }
 
-int IoUtils::ReadStreamForVocab(int num_lines) {
+std::pair<int, bool> IoUtils::ReadStreamForVocab(int num_lines) {
   int read_cnt = 0;
   std::string line;
   std::vector<std::string> line_vec;
-  while (getline(stream_fin_, line) and read_cnt < num_lines) {
+  while (not stream_fin_.eof() and read_cnt < num_lines) {
+    getline(stream_fin_, line);
     ParseLine(line, line_vec);
     for (auto& word: line_vec) {
       if (not word_count_.count(word)) word_count_[word] = 0;
       word_count_[word]++;
     }
     read_cnt++;
   }
-  if (read_cnt < num_lines) stream_fin_.close();
-  return read_cnt;
+  bool finished = false;
+  if (stream_fin_.eof()) {
+    stream_fin_.close();
+    finished = true;
+  }
+  return {read_cnt, finished};
 }
 
 void IoUtils::GetWordVocab(int min_count) {

cusim/ioutils/bindings.cc

@@ -27,7 +27,7 @@ class IoUtilsBind {
     return obj_.LoadStreamFile(filepath);
   }
 
-  int ReadStreamForVocab(int num_lines) {
+  std::pair<int, bool> ReadStreamForVocab(int num_lines) {
     return obj_.ReadStreamForVocab(num_lines);
   }
 

cusim/ioutils/pyioutils.py

@@ -33,9 +33,9 @@ def load_stream_vocab(self, filepath, min_count, chunk_lines=100000):
     full_num_lines = self.obj.load_stream_file(filepath)
     pbar = tqdm.trange(full_num_lines)
     while True:
-      num_lines = self.obj.read_stream_for_vocab(chunk_lines)
+      num_lines, finished = self.obj.read_stream_for_vocab(chunk_lines)
       pbar.update(num_lines)
-      if num_lines < chunk_lines:
-        pbar.close()
+      if finished:
         break
+    pbar.close()
     self.obj.get_word_vocab(min_count)

examples/example1.py

@@ -33,7 +33,7 @@ def download():
 def run():
   download()
   iou = IoUtils()
-  iou.load_stream_vocab(DATA_PATH, 5)
+  iou.load_stream_vocab(DATA_PATH, 5, 10000)
 
 
 if __name__ == "__main__":