src/include/fst/extensions/far/extract.h

   1 // See www.openfst.org for extensive documentation on this weighted
   2 // finite-state transducer library.
   3 //
   4 // Extracts component FSTs from an finite-state archive.
   5
   6 #ifndef FST_EXTENSIONS_FAR_EXTRACT_H_
   7 #define FST_EXTENSIONS_FAR_EXTRACT_H_
   8
   9 #include <memory>
  10 #include <string>
  11 #include <vector>
  12
  13 #include <fst/extensions/far/far.h>
  14
  15 namespace fst {
  16
  17 template <class Arc>
  18 inline void FarWriteFst(const Fst<Arc> *fst, string key, string *okey,
  19                         int *nrep, int32 generate_filenames, int i,
  20                         const string &filename_prefix,
  21                         const string &filename_suffix) {
  22   if (key == *okey) {
  23     ++*nrep;
  24   } else {
  25     *nrep = 0;
  26   }
  27   *okey = key;
  28   string ofilename;
  29   if (generate_filenames) {
  30     std::ostringstream tmp;
  31     tmp.width(generate_filenames);
  32     tmp.fill('0');
  33     tmp << i;
  34     ofilename = tmp.str();
  35   } else {
  36     if (*nrep > 0) {
  37       std::ostringstream tmp;
  38       tmp << '.' << nrep;
  39       key.append(tmp.str().data(), tmp.str().size());
  40     }
  41     ofilename = key;
  42   }
  43   fst->Write(filename_prefix + ofilename + filename_suffix);
  44 }
  45
  46 template <class Arc>
  47 void FarExtract(const std::vector<string> &ifilenames, int32 generate_filenames,
  48                 const string &keys, const string &key_separator,
  49                 const string &range_delimiter, const string &filename_prefix,
  50                 const string &filename_suffix) {
  51   std::unique_ptr<FarReader<Arc>> far_reader(
  52       FarReader<Arc>::Open(ifilenames));
  53   if (!far_reader) return;
  54   string okey;
  55   int nrep = 0;
  56   std::vector<char *> key_vector;
  57   // User has specified a set of FSTs to extract, where some of these may in
  58   // fact be ranges.
  59   if (!keys.empty()) {
  60     auto *keys_cstr = new char[keys.size() + 1];
  61     strcpy(keys_cstr, keys.c_str());
  62     SplitString(keys_cstr, key_separator.c_str(), &key_vector, true);
  63     int i = 0;
  64     for (size_t k = 0; k < key_vector.size(); ++k, ++i) {
  65       string key = key_vector[k];
  66       auto *key_cstr = new char[key.size() + 1];
  67       strcpy(key_cstr, key.c_str());
  68       std::vector<char *> range_vector;
  69       SplitString(key_cstr, range_delimiter.c_str(), &range_vector, false);
  70       if (range_vector.size() == 1) {  // Not a range
  71         if (!far_reader->Find(key)) {
  72           LOG(ERROR) << "FarExtract: Cannot find key " << key;
  73           return;
  74         }
  75         const auto *fst = far_reader->GetFst();
  76         FarWriteFst(fst, key, &okey, &nrep, generate_filenames, i,
  77                     filename_prefix, filename_suffix);
  78       } else if (range_vector.size() == 2) {  // A legal range
  79         string begin_key = range_vector[0];
  80         string end_key = range_vector[1];
  81         if (begin_key.empty() || end_key.empty()) {
  82           LOG(ERROR) << "FarExtract: Illegal range specification " << key;
  83           return;
  84         }
  85         if (!far_reader->Find(begin_key)) {
  86           LOG(ERROR) << "FarExtract: Cannot find key " << begin_key;
  87           return;
  88         }
  89         for (; !far_reader->Done(); far_reader->Next(), ++i) {
  90           const auto &ikey = far_reader->GetKey();
  91           if (end_key < ikey) break;
  92           const auto *fst = far_reader->GetFst();
  93           FarWriteFst(fst, ikey, &okey, &nrep, generate_filenames, i,
  94                       filename_prefix, filename_suffix);
  95         }
  96       } else {
  97         LOG(ERROR) << "FarExtract: Illegal range specification " << key;
  98         return;
  99       }
 100       delete[] key_cstr;
 101     }
 102     delete[] keys_cstr;
 103     return;
 104   }
 105   // Nothing specified, so just extracts everything.
 106   for (size_t i = 1; !far_reader->Done(); far_reader->Next(), ++i) {
 107     const auto &key = far_reader->GetKey();
 108     const auto *fst = far_reader->GetFst();
 109     FarWriteFst(fst, key, &okey, &nrep, generate_filenames, i, filename_prefix,
 110                 filename_suffix);
 111   }
 112   return;
 113 }
 114
 115 }  // namespace fst
 116
 117 #endif  // FST_EXTENSIONS_FAR_EXTRACT_H_