fix local normalizer code for summary features

3 files changed, 15 insertions, 8 deletions

diff --git a/decoder/decoder.cc b/decoder/decoder.cc
index a16a9b5a..89425198 100644
--- a/decoder/decoder.cc
+++ b/decoder/decoder.cc
@@ -806,15 +806,21 @@ bool DecoderImpl::Decode(const string& input, DecoderObserver* o) {
     }
 
     if (rp.fid_summary) {
-      Hypergraph::EdgeProbs posteriors;
-      const prob_t z = forest.ComputeEdgePosteriors(1.0, &posteriors);
+      const prob_t z = forest.PushWeightsToGoal(1.0);
       if (!SILENT) { cerr << "  " << passtr << " adding summary feature " << FD::Convert(rp.fid_summary) << " log(Z)=" << log(z) << endl; }
-      assert(forest.edges_.size() == posteriors.size());
       if (!isfinite(log(z)) || isnan(log(z))) {
         cerr << "  " << passtr << " !!! Invalid partition detected, abandoning.\n";
       } else {
-        for (int i = 0; i < posteriors.size(); ++i)
-          forest.edges_[i].feature_values_.set_value(rp.fid_summary, log(posteriors[i] / z));
+        for (int i = 0; i < forest.edges_.size(); ++i) {
+          const double log_prob_transition = log(forest.edges_[i].edge_prob_); // locally normalized by the edge
+                                                                            // head node by forest.PushWeightsToGoal
+          if (!isfinite(log_prob_transition) || isnan(log_prob_transition)) {
+            cerr << "Edge: i=" << i << " got bad inside prob: " << *forest.edges_[i].rule_ << endl;
+            abort();
+          }
+
+          forest.edges_[i].feature_values_.set_value(rp.fid_summary, log_prob_transition);
+        }
       }
     }
 
diff --git a/decoder/hg.cc b/decoder/hg.cc
index 39ac5132..a4028b0e 100644
--- a/decoder/hg.cc
+++ b/decoder/hg.cc
@@ -226,9 +226,9 @@ prob_t Hypergraph::PushViterbiWeightsToGoal(int fid) {
 }
 
 
-void Hypergraph::PushWeightsToGoal(double scale) {
+prob_t Hypergraph::PushWeightsToGoal(double scale) {
   vector<prob_t> posts;
-  ComputeEdgePosteriors(scale, &posts);
+  const prob_t inside_z = ComputeEdgePosteriors(scale, &posts);
   for (int i = 0; i < nodes_.size(); ++i) {
     const Hypergraph::Node& node = nodes_[i];
     prob_t z = prob_t::Zero();
@@ -238,6 +238,7 @@ void Hypergraph::PushWeightsToGoal(double scale) {
       edges_[node.in_edges_[j]].edge_prob_ = posts[node.in_edges_[j]] / z;
     }
   }
+  return inside_z;
 }
 
 struct EdgeExistsWeightFunction {
diff --git a/decoder/hg.h b/decoder/hg.h
index aa1202b1..e5ef05f8 100644
--- a/decoder/hg.h
+++ b/decoder/hg.h
@@ -449,7 +449,7 @@ public:
   void PushWeightsToSource(double scale = 1.0);
   // same, except weights are pushed to the goal, works for HGs,
   // not just lattices
-  void PushWeightsToGoal(double scale = 1.0);
+  prob_t PushWeightsToGoal(double scale = 1.0);
 
   // contrary to PushWeightsToGoal, use viterbi semiring; store log(p) to fid.  note that p_viterbi becomes 1; k*p_viterbi becomes k.  also modifies edge_prob_ (note that the fid stored log(p) will stick around even if you reweight)
   // afterwards, product of edge_prob_ for a derivation will equal 1 for the viterbi (p_v before, 1 after), and in general (k*p_v before, k after).  returns inside(goal)