path: root/contrib/moses2/defer/CubePruningPerMiniStack/Search.cpp

/*
 * Search.cpp
 *
 *  Created on: 16 Nov 2015
 *      Author: hieu
 */
#include <boost/foreach.hpp>
#include "Search.h"
#include "../Manager.h"
#include "../Hypothesis.h"
#include "../../InputPaths.h"
#include "../../InputPath.h"
#include "../../System.h"
#include "../../Sentence.h"
#include "../../TranslationTask.h"
#include "../../legacy/Util2.h"

using namespace std;

namespace Moses2
{

namespace NSCubePruningPerMiniStack
{

////////////////////////////////////////////////////////////////////////
Search::Search(Manager &mgr)
:Moses2::Search(mgr)
,m_stacks(mgr)

,m_queue(QueueItemOrderer(),
		std::vector<QueueItem*>() )

,m_seenPositions()
{
}

Search::~Search()
{
}

void Search::Decode()
{
	// init stacks
	m_stacks.Init(mgr.GetInput().GetSize() + 1);

	const Bitmap &initBitmap = mgr.GetBitmaps().GetInitialBitmap();
	Hypothesis *initHypo = Hypothesis::Create(mgr.GetSystemPool(), mgr);
	initHypo->Init(mgr, mgr.GetInputPaths().GetBlank(), mgr.GetInitPhrase(), initBitmap);
	initHypo->EmptyHypothesisState(mgr.GetInput());

	m_stacks.Add(initHypo, mgr.GetHypoRecycle());

	for (size_t stackInd = 0; stackInd < m_stacks.GetSize() - 1; ++stackInd) {
		CreateSearchGraph(stackInd);
	}

	for (size_t stackInd = 1; stackInd < m_stacks.GetSize(); ++stackInd) {
		//cerr << "stackInd=" << stackInd << endl;
		Decode(stackInd);

		//cerr << m_stacks << endl;
	}

	//DebugCounts();
}

void Search::Decode(size_t stackInd)
{
  NSCubePruningMiniStack::Stack &stack = m_stacks[stackInd];
  BOOST_FOREACH(NSCubePruningMiniStack::Stack::Coll::value_type &val, stack.GetColl()) {
	  NSCubePruningMiniStack::MiniStack &miniStack = *val.second;
	  Decode(miniStack);
  }

}

void Search::Decode(NSCubePruningMiniStack::MiniStack &miniStack)
{
	Recycler<Hypothesis*> &hypoRecycler  = mgr.GetHypoRecycle();

	// reuse queue from previous stack. Clear it first
	std::vector<QueueItem*> &container = Container(m_queue);
	//cerr << "container=" << container.size() << endl;
	BOOST_FOREACH(QueueItem *item, container) {
		// recycle unused hypos from queue
		Hypothesis *hypo = item->hypo;
		hypoRecycler.Recycle(hypo);

		// recycle queue item
		m_queueItemRecycler.push_back(item);
	}
	container.clear();

	m_seenPositions.clear();

	// add top hypo from every edge into queue
	CubeEdges &edges = *m_cubeEdges[&miniStack];

	BOOST_FOREACH(CubeEdge *edge, edges) {
		//cerr << "edge=" << *edge << endl;
		edge->CreateFirst(mgr, m_queue, m_seenPositions, m_queueItemRecycler);
	}

	size_t pops = 0;
	while (!m_queue.empty() && pops < mgr.system.popLimit) {
		// get best hypo from queue, add to stack
		//cerr << "queue=" << queue.size() << endl;
		QueueItem *item = m_queue.top();
		m_queue.pop();

		CubeEdge *edge = item->edge;

		// add hypo to stack
		Hypothesis *hypo = item->hypo;
		//cerr << "hypo=" << *hypo << " " << hypo->GetBitmap() << endl;
		m_stacks.Add(hypo, hypoRecycler);

		edge->CreateNext(mgr, item, m_queue, m_seenPositions, m_queueItemRecycler);

		++pops;
	}

	/*
	// create hypo from every edge. Increase diversity
	while (!m_queue.empty()) {
		QueueItem *item = m_queue.top();
		m_queue.pop();

		if (item->hypoIndex == 0 && item->tpIndex == 0) {
			CubeEdge &edge = item->edge;

			// add hypo to stack
			Hypothesis *hypo = item->hypo;
			//cerr << "hypo=" << *hypo << " " << hypo->GetBitmap() << endl;
			m_stacks.Add(hypo, mgr.GetHypoRecycle());
		}
	}
	*/
}


void Search::CreateSearchGraph(size_t stackInd)
{
  NSCubePruningMiniStack::Stack &stack = m_stacks[stackInd];
  MemPool &pool = mgr.GetPool();

  BOOST_FOREACH(const NSCubePruningMiniStack::Stack::Coll::value_type &val, stack.GetColl()) {
	  const Bitmap &hypoBitmap = *val.first.first;
	  size_t hypoEndPos = val.first.second;
	  //cerr << "key=" << hypoBitmap << " " << hypoEndPos << endl;

	  // create edges to next hypos from existing hypos
	  const InputPaths &paths = mgr.GetInputPaths();

	  BOOST_FOREACH(const InputPath *path, paths) {
		const Range &pathRange = path->range;
		//cerr << "pathRange=" << pathRange << endl;

		if (!path->IsUsed()) {
			continue;
		}
		if (!CanExtend(hypoBitmap, hypoEndPos, pathRange)) {
			continue;
		}

		const Bitmap &newBitmap = mgr.GetBitmaps().GetBitmap(hypoBitmap, pathRange);

		// sort hypo for a particular bitmap and hypoEndPos
		const NSCubePruningMiniStack::MiniStack &miniStack = *val.second;


		// add cube edge
  	    size_t numPt = mgr.system.mappings.size();
  	    for (size_t i = 0; i < numPt; ++i) {
  	    	const TargetPhrases *tps = path->targetPhrases[i];
  			if (tps && tps->GetSize()) {
  				// create next mini stack
  				NSCubePruningMiniStack::MiniStack &nextMiniStack = m_stacks.GetMiniStack(newBitmap, pathRange);

  				CubeEdge *edge = new (pool.Allocate<CubeEdge>()) CubeEdge(mgr, miniStack, *path, *tps, newBitmap);

				CubeEdges *edges;
				boost::unordered_map<NSCubePruningMiniStack::MiniStack*, CubeEdges*>::iterator iter = m_cubeEdges.find(&nextMiniStack);
				if (iter == m_cubeEdges.end()) {
					edges = new (pool.Allocate<CubeEdges>()) CubeEdges();
					m_cubeEdges[&nextMiniStack] = edges;
				}
				else {
					edges = iter->second;
				}

				edges->push_back(edge);
  			}
		}
	  }
  }

}


const Hypothesis *Search::GetBestHypo() const
{
	const NSCubePruningMiniStack::Stack &lastStack = m_stacks.Back();
	std::vector<const Hypothesis*> sortedHypos = lastStack.GetBestHypos(1);

	const Hypothesis *best = NULL;
	if (sortedHypos.size()) {
		best = sortedHypos[0];
	}
	return best;
}

void Search::DebugCounts()
{
	std::map<size_t, size_t> counts;

	for (size_t stackInd = 0; stackInd < m_stacks.GetSize(); ++stackInd) {
		//cerr << "stackInd=" << stackInd << endl;
		const NSCubePruningMiniStack::Stack &stack = m_stacks[stackInd];
		BOOST_FOREACH(const NSCubePruningMiniStack::Stack::Coll::value_type &val, stack.GetColl()) {
			const NSCubePruningMiniStack::MiniStack &miniStack = *val.second;
			size_t count = miniStack.GetColl().size();

			if (counts.find(count) == counts.end()) {
				counts[count] = 0;
			}
			else {
				++counts[count];
			}
		}
		//cerr << m_stacks << endl;
	}

	std::map<size_t, size_t>::const_iterator iter;
	for (iter = counts.begin(); iter != counts.end(); ++iter) {
		cerr << iter->first << "=" << iter->second << " ";
	}
	cerr << endl;
}



}

}