| 177 | } |
| 178 | |
| 179 | Status Init(Env* env, const string& filename) { |
| 180 | string data; |
| 181 | TF_RETURN_IF_ERROR(ReadFileToString(env, filename, &data)); |
| 182 | StringPiece input = data; |
| 183 | string w; |
| 184 | corpus_size_ = 0; |
| 185 | std::unordered_map<string, int32> word_freq; |
| 186 | while (ScanWord(&input, &w)) { |
| 187 | ++(word_freq[w]); |
| 188 | ++corpus_size_; |
| 189 | } |
| 190 | if (corpus_size_ < window_size_ * 10) { |
| 191 | return errors::InvalidArgument( |
| 192 | "The text file ", filename, |
| 193 | " contains too little data: ", corpus_size_, " words"); |
| 194 | } |
| 195 | typedef std::pair<string, int32> WordFreq; |
| 196 | std::vector<WordFreq> ordered; |
| 197 | for (const auto& p : word_freq) { |
| 198 | if (p.second >= min_count_) ordered.push_back(p); |
| 199 | } |
| 200 | LOG(INFO) << "Data file: " << filename << " contains " << data.size() |
| 201 | << " bytes, " << corpus_size_ << " words, " << word_freq.size() |
| 202 | << " unique words, " << ordered.size() |
| 203 | << " unique frequent words."; |
| 204 | word_freq.clear(); |
| 205 | std::sort(ordered.begin(), ordered.end(), |
| 206 | [](const WordFreq& x, const WordFreq& y) { |
| 207 | return x.second > y.second; |
| 208 | }); |
| 209 | vocab_size_ = static_cast<int32>(1 + ordered.size()); |
| 210 | Tensor word(DT_STRING, TensorShape({vocab_size_})); |
| 211 | Tensor freq(DT_INT32, TensorShape({vocab_size_})); |
| 212 | word.flat<tstring>()(0) = "UNK"; |
| 213 | static const int32 kUnkId = 0; |
| 214 | std::unordered_map<string, int32> word_id; |
| 215 | int64 total_counted = 0; |
| 216 | for (std::size_t i = 0; i < ordered.size(); ++i) { |
| 217 | const auto& w = ordered[i].first; |
| 218 | auto id = i + 1; |
| 219 | word.flat<tstring>()(id) = w; |
| 220 | auto word_count = ordered[i].second; |
| 221 | freq.flat<int32>()(id) = word_count; |
| 222 | total_counted += word_count; |
| 223 | word_id[w] = id; |
| 224 | } |
| 225 | freq.flat<int32>()(kUnkId) = corpus_size_ - total_counted; |
| 226 | word_ = word; |
| 227 | freq_ = freq; |
| 228 | corpus_.reserve(corpus_size_); |
| 229 | input = data; |
| 230 | while (ScanWord(&input, &w)) { |
| 231 | corpus_.push_back(gtl::FindWithDefault(word_id, w, kUnkId)); |
| 232 | } |
| 233 | precalc_examples_.resize(kPrecalc); |
| 234 | sentence_.resize(kSentenceSize); |
| 235 | return Status::OK(); |
| 236 | } |
no test coverage detected