use progressbar in reading stream

Author: js1010

cpp/include/ioutils.hpp

@@ -30,12 +30,17 @@ class IoUtils {
  public:
   IoUtils();
   ~IoUtils();
-  void LoadGensimVocab(std::string filepath, int min_count);
+  bool Init(std::string opt_path);
+  int LoadStreamFile(std::string filepath);
+  int ReadStreamForVocab(int num_lines);
+  void GetWordVocab(int min_count);
  private:
-  std::vector<std::string> parse_line(std::string line);
+  void ParseLine(std::string line, std::vector<std::string>& line_vec);
 
+  std::ifstream stream_fin_;
+  json11::Json opt_;
   std::shared_ptr<spdlog::logger> logger_;
-  std::unordered_map<std::string, int> word_idmap_;
+  std::unordered_map<std::string, int> word_idmap_, word_count_;
   std::vector<std::string> word_list_;
 };  // class IoUtils
 

cpp/src/ioutils.cc

@@ -13,9 +13,23 @@ IoUtils::IoUtils() {
 
 IoUtils::~IoUtils() {}
 
-std::vector<std::string> IoUtils::parse_line(std::string line) {
+bool IoUtils::Init(std::string opt_path) {
+  std::ifstream in(opt_path.c_str());
+  if (not in.is_open()) return false;
+
+  std::string str((std::istreambuf_iterator<char>(in)),
+      std::istreambuf_iterator<char>());
+  std::string err_cmt;
+  auto _opt = json11::Json::parse(str, err_cmt);
+  if (not err_cmt.empty()) return false;
+  opt_ = _opt;
+  CuSimLogger().set_log_level(opt_["c_log_level"].int_value());
+  return true;
+}
+
+void IoUtils::ParseLine(std::string line, std::vector<std::string>& ret) {
+  ret.clear();
   int n = line.size();
-  std::vector<std::string> ret;
   std::string element;
   for (int i = 0; i < n; ++i) {
     if (line[i] == ' ') {
@@ -28,26 +42,42 @@ std::vector<std::string> IoUtils::parse_line(std::string line) {
   if (element.size() > 0) {
     ret.push_back(element);
   }
-  return ret;
 }
 
-void IoUtils::LoadGensimVocab(std::string filepath, int min_count) {
-  INFO("read gensim file to generate vocabulary: {}, min_count: {}", filepath, min_count);
-  std::ifstream fin(filepath.c_str());
-  std::unordered_map<std::string, int> word_count;
-  while (not fin.eof()) {
-    std::string line;
-    getline(fin, line);
-    std::vector<std::string> line_vec = parse_line(line);
+int IoUtils::LoadStreamFile(std::string filepath) {
+  INFO("read gensim file to generate vocabulary: {}", filepath);
+  stream_fin_.open(filepath.c_str());
+  int count = 0;
+  std::string line;
+  while (getline(stream_fin_, line))
+    count++;
+  stream_fin_.close();
+  stream_fin_.open(filepath.c_str());
+  word_idmap_.clear();
+  word_list_.clear();
+  word_count_.clear();
+  return count;
+}
+
+int 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) {
+    ParseLine(line, line_vec);
     for (auto& word: line_vec) {
-      if (not word_count.count(word)) word_count[word] = 0;
-      word_count[word]++;
+      if (not word_count_.count(word)) word_count_[word] = 0;
+      word_count_[word]++;
     }
+    read_cnt++;
   }
-  INFO("number of raw words: {}", word_count.size());
-  word_idmap_.clear();
-  word_list_.clear();
-  for (auto& it: word_count) {
+  if (read_cnt < num_lines) stream_fin_.close();
+  return read_cnt;
+}
+
+void IoUtils::GetWordVocab(int min_count) {
+  INFO("number of raw words: {}", word_count_.size());
+  for (auto& it: word_count_) {
     if (it.second >= min_count) {
       word_idmap_[it.first] = word_idmap_.size();
       word_list_.push_back(it.first);

cusim/aux.py

@@ -10,7 +10,6 @@
 
 import jsmin
 from google.protobuf.json_format import Parse, MessageToDict
-from cusim.config_pb2 import ConfigProto
 
 # get_logger and Option refer to
 # https://github.com/kakao/buffalo/blob/
@@ -57,7 +56,8 @@ def load_json_file(fname):
   return ret
 
 # use protobuf to restrict field and types
-def get_opt_as_proto(raw, proto_type=ConfigProto):
+def get_opt_as_proto(raw, proto_type=None):
+  assert proto_type is not None
   proto = proto_type()
   # convert raw to proto
   Parse(json.dumps(Option(raw)), proto)

cusim/ioutils/bindings.cc

@@ -18,9 +18,23 @@ typedef py::array_t<int, py::array::c_style | py::array::forcecast> int_array;
 class IoUtilsBind {
  public:
   IoUtilsBind() {}
-  void LoadGensimVocab(std::string filepath, int min_count) {
-    obj_.LoadGensimVocab(filepath, min_count);
+
+  bool Init(std::string opt_path) {
+    return obj_.Init(opt_path);
+  }
+
+  int LoadStreamFile(std::string filepath) {
+    return obj_.LoadStreamFile(filepath);
+  }
+
+  int ReadStreamForVocab(int num_lines) {
+    return obj_.ReadStreamForVocab(num_lines);
+  }
+
+  void GetWordVocab(int min_count) {
+    return obj_.GetWordVocab(min_count);
   }
+
  private:
   cusim::IoUtils obj_;
 };
@@ -30,8 +44,10 @@ PYBIND11_PLUGIN(ioutils_bind) {
 
   py::class_<IoUtilsBind>(m, "IoUtilsBind")
   .def(py::init())
-  .def("load_gensim_vocab", &IoUtilsBind::LoadGensimVocab,
-      py::arg("filepath"), py::arg("min_count"))
+  .def("init", &IoUtilsBind::Init, py::arg("opt_path"))
+  .def("load_stream_file", &IoUtilsBind::LoadStreamFile, py::arg("filepath"))
+  .def("read_stream_for_vocab", &IoUtilsBind::ReadStreamForVocab, py::arg("num_lines"))
+  .def("get_word_vocab", &IoUtilsBind::GetWordVocab, py::arg("min_count"))
   .def("__repr__",
   [](const IoUtilsBind &a) {
     return "<IoUtilsBind>";

cusim/ioutils/pyioutils.py

@@ -5,14 +5,37 @@
 # LICENSE file in the root directory of this source tree.
 
 # pylint: disable=no-name-in-module,too-few-public-methods,no-member
+import os
+import json
+import tempfile
+import tqdm
+
 from cusim import aux
 from cusim.ioutils.ioutils_bind import IoUtilsBind
-
+from cusim.config_pb2 import IoUtilsConfigProto
 
 class IoUtils:
-  def __init__(self, log_level=2):
-    self.logger = aux.get_logger("ioutils", level=log_level)
+  def __init__(self, opt=None):
+    self.opt = aux.get_opt_as_proto(opt or {}, IoUtilsConfigProto)
+    self.logger = aux.get_logger("ioutils", level=self.opt.py_log_level)
+
+    tmp = tempfile.NamedTemporaryFile(mode='w', delete=False)
+    opt_content = json.dumps(aux.proto_to_dict(self.opt), indent=2)
+    tmp.write(opt_content)
+    tmp.close()
+
+    self.logger.info("opt: %s", opt_content)
     self.obj = IoUtilsBind()
+    assert self.obj.init(bytes(tmp.name, "utf8")), f"failed to load {tmp.name}"
+    os.remove(tmp.name)
 
-  def load_gensim_vocab(self, filepath, min_count):
-    self.obj.load_gensim_vocab(filepath, min_count)
+  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)
+      pbar.update(num_lines)
+      if num_lines < chunk_lines:
+        pbar.close()
+        break
+    self.obj.get_word_vocab(min_count)

cusim/proto/config.proto

@@ -6,22 +6,7 @@
 
 syntax = "proto2";
 
-message ConfigProto {
-  optional int32 seed = 1 [default = 777];
-  optional int32 c_log_level = 3 [default = 2];
-  optional int32 py_log_level = 4 [default = 2];
-  optional int32 max_m = 5 [default = 12];
-  optional int32 max_m0 = 6 [default = 24];
-  optional int32 ef_construction = 7 [default = 150];
-  // optional int32 ef_search = 8 [default = 50];
-  optional double level_mult = 9;
-  optional bool save_remains = 10;
-  optional double hyper_threads = 11 [default = 10];
-  optional int32 block_dim = 12 [default = 32];
-  optional string dist_type = 13 [default = "dot"];
-  optional int32 visited_table_size = 17;
-  optional int32 visited_list_size = 14 [default = 8192];
-  optional bool nrz = 15 [default = true];
-  optional bool reverse_cand = 16;
-  optional double heuristic_coef = 18 [default = 0.25];
+message IoUtilsConfigProto {
+  optional int32 py_log_level = 1 [default = 2];
+  optional int32 c_log_level = 2 [default = 2];
 }

examples/example1.py

@@ -5,16 +5,35 @@
 # LICENSE file in the root directory of this source tree.
 
 # pylint: disable=no-name-in-module,logging-format-truncated
+import os
+import subprocess
 import fire
+
+from gensim import downloader as api
 from cusim import aux, IoUtils
 
 LOGGER = aux.get_logger()
-CORPORA_PATH = "res/corpora.txt"
+DOWNLOAD_PATH = "./res"
+DATASET = "wiki-english-20171001"
+DATA_PATH = f"./res/{DATASET}.stream.txt"
 MIN_COUNT = 5
 
+def download():
+  if os.path.exists(DATA_PATH):
+    LOGGER.info("%s already exists", DATA_PATH)
+    return
+  api.BASE_DIR = DOWNLOAD_PATH
+  filepath = api.load(DATASET, return_path=True)
+  LOGGER.info("filepath: %s", filepath)
+  cmd = ["gunzip", "-c", filepath, ">", DATA_PATH]
+  cmd = " ".join(cmd)
+  LOGGER.info("cmd: %s", cmd)
+  subprocess.call(cmd, shell=True)
+
 def run():
+  download()
   iou = IoUtils()
-  iou.load_gensim_vocab(CORPORA_PATH, MIN_COUNT)
+  iou.load_stream_vocab(DATA_PATH, 5)
 
 
 if __name__ == "__main__":

examples/requirements.txt

@@ -1 +1,2 @@
 fire
+gensim