vllm_fairseq / mosesdecoder /moses2 /legacy /OutputCollector.cpp
sleepyhead111's picture
Add files using upload-large-folder tool
99f07cf verified
#include "OutputCollector.h"
namespace Moses2
{
OutputCollector::OutputCollector(std::ostream* outStream,
std::ostream* debugStream) :
m_nextOutput(0), m_outStream(outStream), m_debugStream(debugStream), m_isHoldingOutputStream(
false), m_isHoldingDebugStream(false) {
}
OutputCollector::OutputCollector(std::string xout, std::string xerr) :
m_nextOutput(0) {
// TO DO open magic streams instead of regular ofstreams! [UG]
if (xout == "/dev/stderr") {
m_outStream = &std::cerr;
m_isHoldingOutputStream = false;
}
else if (xout.size() && xout != "/dev/stdout" && xout != "-") {
m_outStream = new std::ofstream(xout.c_str());
UTIL_THROW_IF2(!m_outStream->good(),
"Failed to open output file" << xout);
m_isHoldingOutputStream = true;
}
else {
m_outStream = &std::cout;
m_isHoldingOutputStream = false;
}
if (xerr == "/dev/stdout") {
m_debugStream = &std::cout;
m_isHoldingDebugStream = false;
}
else if (xerr.size() && xerr != "/dev/stderr") {
m_debugStream = new std::ofstream(xerr.c_str());
UTIL_THROW_IF2(!m_debugStream->good(),
"Failed to open debug stream" << xerr);
m_isHoldingDebugStream = true;
}
else {
m_debugStream = &std::cerr;
m_isHoldingDebugStream = false;
}
}
OutputCollector::~OutputCollector() {
if (m_isHoldingOutputStream) delete m_outStream;
if (m_isHoldingDebugStream) delete m_debugStream;
}
void OutputCollector::Write(int sourceId, const std::string& output, const std::string& debug) {
#ifdef WITH_THREADS
boost::mutex::scoped_lock lock(m_mutex);
#endif
if (sourceId == m_nextOutput) {
//This is the one we were expecting
*m_outStream << output << std::flush;
*m_debugStream << debug << std::flush;
++m_nextOutput;
//see if there's any more
std::unordered_map<int, std::string>::iterator iter;
while ((iter = m_outputs.find(m_nextOutput)) != m_outputs.end()) {
*m_outStream << iter->second << std::flush;
++m_nextOutput;
std::unordered_map<int, std::string>::iterator debugIter = m_debugs.find(
iter->first);
m_outputs.erase(iter);
if (debugIter != m_debugs.end()) {
*m_debugStream << debugIter->second << std::flush;
m_debugs.erase(debugIter);
}
}
}
else {
//save for later
m_outputs[sourceId] = output;
m_debugs[sourceId] = debug;
}
}
}