criteo2hugectr.cpp

/*
 * Copyright (c) 2020, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <sys/stat.h>

#include <fstream>
#include <ios>
#include <iostream>
#include <sstream>
#include <unordered_set>
#include <vector>

#include "HugeCTR/include/data_generator.hpp"

using namespace HugeCTR;

static std::string usage_str =
    "usage: ./criteo2hugectr in.txt dir/prefix file_list.txt [option:#keys for wide model,default "
    "is 0] [option: Number of files in each file_list.txt,default is 0(all in one file)]";
static const int N = 40960;  // number of samples per data file
static int KEYS_WIDE_MODEL = 0;
static const int KEYS_DENSE_MODEL = 26;
static const int dense_dim = 13;
typedef unsigned int T;
static const long long label_dim = 1;
static int SLOT_NUM = 26;
static int FILELIST_LENGTH = 0;  // number of files in each file_list.txt
const int RANGE[] = {0,       1460,    2018,    337396,  549106,  549411,  549431,
                     561567,  562200,  562203,  613501,  618803,  951403,  954582,
                     954609,  966800,  1268011, 1268021, 1272862, 1274948, 1274952,
                     1599225, 1599242, 1599257, 1678991, 1679087, 1737709};
std::unordered_set<unsigned int> keyset;

std::vector<std::string> &split(const std::string &s, char delim, std::vector<std::string> &elems) {
  std::stringstream ss(s);
  std::string item;
  while (std::getline(ss, item, delim)) {
    elems.push_back(item);
  }
  return elems;
}

int main(int argc, char *argv[]) {
  const std::string tmp_file_list_name("file_list.tmp");

  if (argc != 4 && argc != 5 && argc != 6) {
    std::cout << usage_str << std::endl;
    exit(-1);
  }

  if (argc == 5) {
    if (atoi(argv[4]) != 0) {
      KEYS_WIDE_MODEL = atoi(argv[4]);
      SLOT_NUM += 1;
    } else {
      std::cout << "Checking the range of each key..." << std::endl;
    }
  }

  if (argc == 6) {
    if (atoi(argv[4]) != 0) {
      KEYS_WIDE_MODEL = atoi(argv[4]);
      SLOT_NUM += 1;
      std::cout << "slot_num for w&D is:" << SLOT_NUM << std::endl;
    }
    if (atoi(argv[5]) > 0) {
      FILELIST_LENGTH = atoi(argv[5]);
      std::cout << FILELIST_LENGTH << std::endl;
    } else {
      std::cerr << "The number of files in file_list should greater than 0 (default is 0)..." << std::endl;
    }
  }

  // open txt file
  std::ifstream txt_file(argv[1], std::ifstream::binary);
  if (!txt_file.is_open()) {
    std::cerr << "Cannot open argv[1]" << std::endl;
  }
  // create a data file under prefix
  std::string data_prefix(argv[2]);
  std::string directory;
  const size_t last_slash_idx = data_prefix.rfind('/');
  if (std::string::npos != last_slash_idx) {
    directory = data_prefix.substr(0, last_slash_idx);
  }
  check_make_dir(directory);
  // check file_list.txt prefix
  std::string file_name(argv[3]);
  if (file_name.find(".") == std::string::npos) {
    std::cerr << "Please provide aviable file_list with extension(.txt) " << std::endl;
    exit(-1);
  }
  const size_t last_point_idx = file_name.rfind('.');
  std::string file_name_prefix = file_name.substr(0, last_point_idx);
  std::string file_name_postfix = file_name.substr(last_point_idx);

  int file_counter = 0;
  std::ofstream file_list_tmp(tmp_file_list_name, std::ofstream::out);

  if (!file_list_tmp.is_open()) {
    std::cerr << "Cannot open file_list.tmp" << std::endl;
  }

  std::string keyset_name(file_name_prefix + ".keyset");
  if (FILELIST_LENGTH > 0) {
    keyset_name = file_name_prefix + "." + std::to_string(file_counter) + ".keyset";
  }
  std::ofstream keyset_file(keyset_name, std::ofstream::binary);
  DataWriter<Check_t::None> keyset_writer(keyset_file);

  do {
    std::string data_file_name(data_prefix + std::to_string(file_counter) + ".data");
    std::ofstream data_file(data_file_name, std::ofstream::binary);
    file_list_tmp << (data_file_name + "\n");
    std::cout << data_file_name << std::endl;
    if (!data_file.is_open()) {
      std::cerr << "Cannot open data_file" << std::endl;
    }

    DataWriter<Check_t::Sum> data_writer(data_file);
    DataSetHeader header = {
        1, static_cast<long long>(N), label_dim, dense_dim, static_cast<long long>(SLOT_NUM), 0, 0,
        0};
    data_writer.append(reinterpret_cast<char *>(&header), sizeof(DataSetHeader));
    data_writer.write();
    // read N lines
    int i = 0;
    for (; i < N; i++) {
      // read a line
      std::string line;
      std::getline(txt_file, line);
      // if end
      if (txt_file.eof()) {
        txt_file.close();
        if (i == 0 && keyset.size() == 0) {
          return 0;
        }
        data_file.seekp(std::ios_base::beg);
        DataSetHeader last_header = {1,
                                     static_cast<long long>(i),
                                     label_dim,
                                     dense_dim,
                                     static_cast<long long>(SLOT_NUM),
                                     0,
                                     0,
                                     0};
        data_writer.append(reinterpret_cast<char *>(&last_header), sizeof(DataSetHeader));
        data_writer.write();
        data_file.close();
        std::cout << "last keyset size is:" << keyset.size() << std::endl;
        keyset_writer.write();
        keyset_file.close();

        file_list_tmp.close();
        // redirect
        {
          std::ifstream tmp(tmp_file_list_name);
          if (!tmp.is_open()) {
            std::cerr << "Cannot open " << tmp_file_list_name << std::endl;
          }

          std::cout << "Opening " << argv[3] << std::endl;
          // std::ofstream file_list(argv[3], std::ofstream::out);

          std::string file_list_name = file_name;

          if (FILELIST_LENGTH > 0) {
            file_list_name = file_name_prefix + "." +
                             std::to_string((file_counter / FILELIST_LENGTH)) + file_name_postfix;
          }

          std::ofstream file_list(file_list_name, std::ofstream::out);

          if (!file_list.is_open()) {
            std::cerr << "Cannot open " << argv[3] << std::endl;
          }

          // file_list << (std::to_string(file_counter + 1) + "\n");
          if (i != 0) {
            file_counter++;
          }
          if (FILELIST_LENGTH == 0) {
            file_list << (std::to_string(file_counter) + "\n");
          } else if (file_counter % FILELIST_LENGTH != 0 && FILELIST_LENGTH != 1) {
            file_list << (std::to_string(file_counter % FILELIST_LENGTH) + "\n");
          } else {
            file_list << (std::to_string(FILELIST_LENGTH) + "\n");
          }

          file_list << tmp.rdbuf();

          tmp.close();
          file_list.close();
        }

        return 0;
      }
      std::vector<std::string> vec_string;
      split(line, ' ', vec_string);
      if (vec_string.size() != (unsigned int)(KEYS_WIDE_MODEL + KEYS_DENSE_MODEL + dense_dim +
                                              label_dim))  // first one is label
      {
        std::cerr << "vec_string.size() != KEYS_WIDE_MODEL+KEYS_DENSE_MODEL+dense_dim+label_dim"
                  << std::endl;
        std::cerr << line << std::endl;
        exit(-1);
      }
#ifndef NDEBUG
      std::cout << std::endl;
#endif
      for (int j = 0; j < dense_dim + label_dim; j++) {
        float label_dense = std::stod(vec_string[j]);
        data_writer.append(reinterpret_cast<char *>(&label_dense), sizeof(float));
#ifndef NDEBUG
        std::cout << label_dense << ' ';
#endif
      }

      if (KEYS_WIDE_MODEL != 0) {
        data_writer.append(reinterpret_cast<char *>(&KEYS_WIDE_MODEL), sizeof(int));
        for (int j = dense_dim + label_dim; j < KEYS_WIDE_MODEL + dense_dim + label_dim; j++) {
          T key = static_cast<T>(std::stoll(vec_string[j]));
          data_writer.append(reinterpret_cast<char *>(&key), sizeof(T));
          if (keyset.find(key) == keyset.end()) {
            keyset.insert(key);
            keyset_writer.append(reinterpret_cast<char *>(&key), sizeof(T));
          }

#ifndef NDEBUG
          std::cout << key << ',';
#endif
        }
      }
      int nnz = 1;
      for (int j = KEYS_WIDE_MODEL + dense_dim + label_dim;
           j < KEYS_WIDE_MODEL + KEYS_DENSE_MODEL + dense_dim + label_dim; j++) {
        T key = static_cast<T>(std::stoll(vec_string[j]));
        data_writer.append(reinterpret_cast<char *>(&nnz), sizeof(int));
        data_writer.append(reinterpret_cast<char *>(&key), sizeof(T));
        if (keyset.find(key) == keyset.end()) {
          keyset.insert(key);
          keyset_writer.append(reinterpret_cast<char *>(&key), sizeof(T));
        }
        if (KEYS_WIDE_MODEL == 0) {
          if (key < static_cast<T>(RANGE[j - dense_dim - label_dim]) ||
              key > static_cast<T>(RANGE[j + 1 - dense_dim - label_dim])) {
            std::cout << key << " in feature:" << j << " our of range:" << RANGE[j] << ","
                      << RANGE[j + 1] << std::endl;
          }
        }
#ifndef NDEBUG
        std::cout << key << ',';
#endif
      }
      data_writer.write();
    }
    data_file.close();

    if (FILELIST_LENGTH > 0 && (file_counter + 1) % FILELIST_LENGTH == 0) {
      int NOFile = (file_counter / FILELIST_LENGTH);
      file_list_tmp.close();
      std::ifstream tmp(tmp_file_list_name);
      if (!tmp.is_open()) {
        std::cerr << "Cannot open " << tmp_file_list_name << std::endl;
      }

      std::cout << "Opening " << argv[3] << std::endl;

      std::string file_list_name =
          file_name_prefix + "." + std::to_string(NOFile) + file_name_postfix;

      std::ofstream file_list(file_list_name, std::ofstream::out);
      if (!file_list.is_open()) {
        std::cerr << "Cannot open " << argv[3] << std::endl;
      }

      file_list << (std::to_string(FILELIST_LENGTH) + "\n");
      file_list << tmp.rdbuf();

      tmp.close();
      file_list.close();

      std::cout << std::to_string(NOFile) << " keyset size is: " << keyset.size() << std::endl;
      keyset.clear();
      keyset_writer.write();
      keyset_file.close();
      keyset_name = file_name_prefix + "." + std::to_string(NOFile + 1) + ".keyset";
      keyset_file.open(keyset_name, std::ofstream::binary);
      DataWriter<Check_t::None> keyset_writer(keyset_file);

      std::cout << "reopen file_list_tmp" << std::endl;
      file_list_tmp.open(tmp_file_list_name);
      if (!file_list_tmp.is_open()) {
        std::cerr << "Cannot open file_list.tmp during iteration" << std::endl;
      }
    }
    file_counter++;

#ifndef NDEBUG
    std::cout << std::endl;
    if (file_counter > 2) {
      break;
    }
#endif
  } while (1);

  return 0;
}