Implement hypergraph output for chart moses

1 files changed, 24 insertions, 7 deletions

diff --git a/moses/ChartManager.cpp b/moses/ChartManager.cpp
index 8033c3280..97d05a60e 100644
--- a/moses/ChartManager.cpp
+++ b/moses/ChartManager.cpp
@@ -25,6 +25,7 @@
 #include "ChartHypothesis.h"
 #include "ChartKBestExtractor.h"
 #include "ChartTranslationOptions.h"
+#include "HypergraphOutput.h"
 #include "StaticData.h"
 #include "DecodeStep.h"
 #include "TreeInput.h"
@@ -223,8 +224,9 @@ void ChartManager::CalcNBest(
   }
 }
 
-void ChartManager::GetSearchGraph(long translationId, std::ostream &outputSearchGraphStream) const
+void ChartManager::WriteSearchGraph(const ChartSearchGraphWriter& writer) const
 {
+
   size_t size = m_source.GetSize();
 
   // which hypotheses are reachable?
@@ -237,7 +239,11 @@ void ChartManager::GetSearchGraph(long translationId, std::ostream &outputSearch
     // no hypothesis
     return;
   }
-  FindReachableHypotheses( hypo, reachable);
+  size_t winners = 0;
+  size_t losers = 0;
+
+  FindReachableHypotheses( hypo, reachable, &winners, &losers);
+  writer.WriteHeader(winners, losers);
 
   for (size_t width = 1; width <= size; ++width) {
     for (size_t startPos = 0; startPos <= size-width; ++startPos) {
@@ -246,12 +252,13 @@ void ChartManager::GetSearchGraph(long translationId, std::ostream &outputSearch
       TRACE_ERR(" " << range << "=");
 
       const ChartCell &cell = m_hypoStackColl.Get(range);
-      cell.GetSearchGraph(translationId, outputSearchGraphStream, reachable);
+      cell.WriteSearchGraph(writer, reachable);
     }
   }
 }
 
-void ChartManager::FindReachableHypotheses( const ChartHypothesis *hypo, std::map<unsigned,bool> &reachable ) const
+void ChartManager::FindReachableHypotheses(
+  const ChartHypothesis *hypo, std::map<unsigned,bool> &reachable, size_t* winners, size_t* losers) const
 {
   // do not recurse, if already visited
   if (reachable.find(hypo->GetId()) != reachable.end()) {
@@ -260,9 +267,14 @@ void ChartManager::FindReachableHypotheses( const ChartHypothesis *hypo, std::ma
 
   // recurse
   reachable[ hypo->GetId() ] = true;
+  if (hypo->GetWinningHypothesis() == hypo) {
+    (*winners)++;
+  } else {
+    (*losers)++;
+  }
   const std::vector<const ChartHypothesis*> &previous = hypo->GetPrevHypos();
   for(std::vector<const ChartHypothesis*>::const_iterator i = previous.begin(); i != previous.end(); ++i) {
-    FindReachableHypotheses( *i, reachable );
+    FindReachableHypotheses( *i, reachable, winners, losers );
   }
 
   // also loop over recombined hypotheses (arcs)
@@ -271,14 +283,19 @@ void ChartManager::FindReachableHypotheses( const ChartHypothesis *hypo, std::ma
     ChartArcList::const_iterator iterArc;
     for (iterArc = arcList->begin(); iterArc != arcList->end(); ++iterArc) {
       const ChartHypothesis &arc = **iterArc;
-      FindReachableHypotheses( &arc, reachable );
+      FindReachableHypotheses( &arc, reachable, winners, losers );
     }
   }
 }
 
 void ChartManager::OutputSearchGraphAsHypergraph(std::ostream &outputSearchGraphStream) const {
-  //TODO
+  ChartSearchGraphWriterHypergraph writer(&outputSearchGraphStream);
+  WriteSearchGraph(writer);
+}
 
+void ChartManager::OutputSearchGraphMoses(std::ostream &outputSearchGraphStream) const {
+  ChartSearchGraphWriterMoses writer(&outputSearchGraphStream, m_lineNumber);
+  WriteSearchGraph(writer);
 }
 
 } // namespace Moses