| 739 | } |
| 740 | |
| 741 | NBestEncodeResult Model::SampleEncodeAndScore(absl::string_view normalized, |
| 742 | float inv_theta, int samples, |
| 743 | bool wor, |
| 744 | bool include_best) const { |
| 745 | if (!status().ok() || normalized.empty()) { |
| 746 | return {}; |
| 747 | } |
| 748 | NBestEncodeResult results; |
| 749 | Lattice lattice; |
| 750 | lattice.SetSentence(normalized); |
| 751 | PopulateNodes(&lattice); |
| 752 | |
| 753 | const std::vector<float> alpha = lattice.ForwardAlgorithm(inv_theta); |
| 754 | const float marginal = alpha[lattice.eos_node()->node_id]; |
| 755 | |
| 756 | if (include_best) { |
| 757 | if (!wor) { |
| 758 | LOG(ERROR) << "include_best not supported for wor false"; |
| 759 | return {}; |
| 760 | } |
| 761 | EncodeResult result; |
| 762 | const auto best_path = lattice.Viterbi(); |
| 763 | for (const auto *node : best_path.first) { |
| 764 | result.emplace_back(node->piece, node->id); |
| 765 | } |
| 766 | |
| 767 | // Inclusion probability if we always include the best is 1. |
| 768 | results.emplace_back(result, 0.0); |
| 769 | } |
| 770 | |
| 771 | if (wor) { |
| 772 | // Draw k+1 samples as we need perturbed score of k+1th element |
| 773 | auto nbest_samples = lattice.NBest(samples + 1, true, inv_theta); |
| 774 | |
| 775 | if (include_best) { |
| 776 | std::vector<std::vector<Lattice::Node *>> nbest_paths( |
| 777 | nbest_samples.size()); |
| 778 | for (int i = 0; i < nbest_samples.size(); i++) { |
| 779 | nbest_paths[i] = nbest_samples[i].first; |
| 780 | } |
| 781 | // Remove the best result from the samples if necessary |
| 782 | const auto best_path = lattice.Viterbi(); |
| 783 | |
| 784 | const int index_of_best = |
| 785 | (std::find(nbest_paths.begin(), nbest_paths.end(), best_path.first) - |
| 786 | nbest_paths.begin()); |
| 787 | |
| 788 | if (index_of_best != nbest_samples.size()) { |
| 789 | nbest_samples.erase(nbest_samples.begin() + index_of_best); |
| 790 | } else { |
| 791 | nbest_samples.pop_back(); |
| 792 | } |
| 793 | } |
| 794 | // We use the perturbed score of the k+1th element to calculate the |
| 795 | // inclusion probability. |
| 796 | const double kappa = static_cast<double>(nbest_samples.back().second); |
| 797 | // Discard the last sample |
| 798 | nbest_samples.pop_back(); |