Upload folder using huggingface_hub

41f6dd8 verified almost 2 years ago

9.95 kB

// vim:tabstop=2
#include "ProbingPT.h"
#include "moses/StaticData.h"
#include "moses/FactorCollection.h"
#include "moses/TargetPhraseCollection.h"
#include "moses/InputFileStream.h"
#include "probingpt/querying.h"
#include "probingpt/probing_hash_utils.h"

using namespace std;

namespace Moses
{
ProbingPT::ProbingPT(const std::string &line)
  : PhraseDictionary(line,true)
  ,m_engine(NULL)
  ,load_method(util::POPULATE_OR_READ)
{
  ReadParameters();

  assert(m_input.size() == 1);
  assert(m_output.size() == 1);
}

ProbingPT::~ProbingPT()
{
  delete m_engine;
}

void ProbingPT::Load(AllOptions::ptr const& opts)
{
  m_options = opts;
  SetFeaturesToApply();

  m_engine = new probingpt::QueryEngine(m_filePath.c_str(), load_method);

  m_unkId = 456456546456;

  FactorCollection &vocab = FactorCollection::Instance();

  // source vocab
  const std::map<uint64_t, std::string> &sourceVocab =
    m_engine->getSourceVocab();
  std::map<uint64_t, std::string>::const_iterator iterSource;
  for (iterSource = sourceVocab.begin(); iterSource != sourceVocab.end();
       ++iterSource) {
    string wordStr = iterSource->second;
    //cerr << "wordStr=" << wordStr << endl;

    const Factor *factor = vocab.AddFactor(wordStr);

    uint64_t probingId = iterSource->first;
    size_t factorId = factor->GetId();

    if (factorId >= m_sourceVocab.size()) {
      m_sourceVocab.resize(factorId + 1, m_unkId);
    }
    m_sourceVocab[factorId] = probingId;
  }

  // target vocab
  InputFileStream targetVocabStrme(m_filePath + "/TargetVocab.dat");
  string line;
  while (getline(targetVocabStrme, line)) {
    vector<string> toks = Tokenize(line, "\t");
    UTIL_THROW_IF2(toks.size() != 2, string("Incorrect format:") + line + "\n");

    //cerr << "wordStr=" << toks[0] << endl;

    const Factor *factor = vocab.AddFactor(toks[0]);
    uint32_t probingId = Scan<uint32_t>(toks[1]);

    if (probingId >= m_targetVocab.size()) {
      m_targetVocab.resize(probingId + 1);
    }

    m_targetVocab[probingId] = factor;
  }

  // alignments
  CreateAlignmentMap(m_filePath + "/Alignments.dat");

  // memory mapped file to tps
  string filePath = m_filePath + "/TargetColl.dat";
  file.open(filePath.c_str());
  if (!file.is_open()) {
    throw "Couldn't open file ";
  }

  data = file.data();
  //size_t size = file.size();

  // cache
  //CreateCache(system);

}

void ProbingPT::CreateAlignmentMap(const std::string path)
{
  const std::vector< std::vector<unsigned char> > &probingAlignColl = m_engine->getAlignments();
  m_aligns.resize(probingAlignColl.size(), NULL);

  for (size_t i = 0; i < probingAlignColl.size(); ++i) {
    AlignmentInfo::CollType aligns;

    const std::vector<unsigned char> &probingAligns = probingAlignColl[i];
    for (size_t j = 0; j < probingAligns.size(); j += 2) {
      size_t startPos = probingAligns[j];
      size_t endPos = probingAligns[j+1];
      //cerr << "startPos=" << startPos << " " << endPos << endl;
      aligns.insert(std::pair<size_t,size_t>(startPos, endPos));
    }

    const AlignmentInfo *align = AlignmentInfoCollection::Instance().Add(aligns);
    m_aligns[i] = align;
    //cerr << "align=" << align->Debug(system) << endl;
  }
}

void ProbingPT::SetParameter(const std::string& key, const std::string& value)
{
  if (key == "load") {
    if (value == "lazy") {
      load_method = util::LAZY;
    } else if (value == "populate_or_lazy") {
      load_method = util::POPULATE_OR_LAZY;
    } else if (value == "populate_or_read" || value == "populate") {
      load_method = util::POPULATE_OR_READ;
    } else if (value == "read") {
      load_method = util::READ;
    } else if (value == "parallel_read") {
      load_method = util::PARALLEL_READ;
    } else {
      UTIL_THROW2("load method not supported" << value);
    }
  } else {
    PhraseDictionary::SetParameter(key, value);
  }

}

void ProbingPT::InitializeForInput(ttasksptr const& ttask)
{

}

void ProbingPT::GetTargetPhraseCollectionBatch(const InputPathList &inputPathQueue) const
{
  InputPathList::const_iterator iter;
  for (iter = inputPathQueue.begin(); iter != inputPathQueue.end(); ++iter) {
    InputPath &inputPath = **iter;
    const Phrase &sourcePhrase = inputPath.GetPhrase();

    if (sourcePhrase.GetSize() > m_options->search.max_phrase_length) {
      continue;
    }

    TargetPhraseCollection::shared_ptr tpColl = CreateTargetPhrase(sourcePhrase);
    inputPath.SetTargetPhrases(*this, tpColl, NULL);
  }
}

TargetPhraseCollection::shared_ptr ProbingPT::CreateTargetPhrase(const Phrase &sourcePhrase) const
{
  // create a target phrase from the 1st word of the source, prefix with 'ProbingPT:'
  assert(sourcePhrase.GetSize());

  std::pair<bool, uint64_t> keyStruct = GetKey(sourcePhrase);
  if (!keyStruct.first) {
    return TargetPhraseCollection::shared_ptr();
  }

  // check in cache
  CachePb::const_iterator iter = m_cachePb.find(keyStruct.second);
  if (iter != m_cachePb.end()) {
    //cerr << "FOUND IN CACHE " << keyStruct.second << " " << sourcePhrase.Debug(mgr.system) << endl;
    TargetPhraseCollection *tps = iter->second;
    return TargetPhraseCollection::shared_ptr(tps);
  }

  // query pt
  TargetPhraseCollection *tps = CreateTargetPhrases(sourcePhrase,
                                keyStruct.second);
  return TargetPhraseCollection::shared_ptr(tps);
}

std::pair<bool, uint64_t> ProbingPT::GetKey(const Phrase &sourcePhrase) const
{
  std::pair<bool, uint64_t> ret;

  // create a target phrase from the 1st word of the source, prefix with 'ProbingPT:'
  size_t sourceSize = sourcePhrase.GetSize();
  assert(sourceSize);

  uint64_t probingSource[sourceSize];
  GetSourceProbingIds(sourcePhrase, ret.first, probingSource);
  if (!ret.first) {
    // source phrase contains a word unknown in the pt.
    // We know immediately there's no translation for it
  } else {
    ret.second = m_engine->getKey(probingSource, sourceSize);
  }

  return ret;

}

void ProbingPT::GetSourceProbingIds(const Phrase &sourcePhrase,
                                    bool &ok, uint64_t probingSource[]) const
{

  size_t size = sourcePhrase.GetSize();
  for (size_t i = 0; i < size; ++i) {
    const Word &word = sourcePhrase.GetWord(i);
    uint64_t probingId = GetSourceProbingId(word);
    if (probingId == m_unkId) {
      ok = false;
      return;
    } else {
      probingSource[i] = probingId;
    }
  }

  ok = true;
}

uint64_t ProbingPT::GetSourceProbingId(const Word &word) const
{
  uint64_t ret = 0;

  for (size_t i = 0; i < m_input.size(); ++i) {
    FactorType factorType = m_input[i];
    const Factor *factor = word[factorType];

    size_t factorId = factor->GetId();
    if (factorId >= m_sourceVocab.size()) {
      return m_unkId;
    }
    ret += m_sourceVocab[factorId];
  }

  return ret;
}

TargetPhraseCollection *ProbingPT::CreateTargetPhrases(
  const Phrase &sourcePhrase, uint64_t key) const
{
  TargetPhraseCollection *tps = NULL;

  //Actual lookup
  std::pair<bool, uint64_t> query_result; // 1st=found, 2nd=target file offset
  query_result = m_engine->query(key);
  //cerr << "key2=" << query_result.second << endl;

  if (query_result.first) {
    const char *offset = data + query_result.second;
    uint64_t *numTP = (uint64_t*) offset;

    tps = new TargetPhraseCollection();

    offset += sizeof(uint64_t);
    for (size_t i = 0; i < *numTP; ++i) {
      TargetPhrase *tp = CreateTargetPhrase(offset);
      assert(tp);
      tp->EvaluateInIsolation(sourcePhrase, GetFeaturesToApply());

      tps->Add(tp);

    }

    tps->Prune(true, m_tableLimit);
    //cerr << *tps << endl;
  }

  return tps;

}

TargetPhrase *ProbingPT::CreateTargetPhrase(
  const char *&offset) const
{
  probingpt::TargetPhraseInfo *tpInfo = (probingpt::TargetPhraseInfo*) offset;
  size_t numRealWords = tpInfo->numWords / m_output.size();

  TargetPhrase *tp = new TargetPhrase(this);

  offset += sizeof(probingpt::TargetPhraseInfo);

  // scores
  float *scores = (float*) offset;

  size_t totalNumScores = m_engine->num_scores + m_engine->num_lex_scores;

  if (m_engine->logProb) {
    // set pt score for rule
    tp->GetScoreBreakdown().PlusEquals(this, scores);

    // save scores for other FF, eg. lex RO. Just give the offset
    /*
    if (m_engine->num_lex_scores) {
      tp->scoreProperties = scores + m_engine->num_scores;
    }
    */
  } else {
    // log score 1st
    float logScores[totalNumScores];
    for (size_t i = 0; i < totalNumScores; ++i) {
      logScores[i] = FloorScore(TransformScore(scores[i]));
    }

    // set pt score for rule
    tp->GetScoreBreakdown().PlusEquals(this, logScores);

    // save scores for other FF, eg. lex RO.
    /*
    tp->scoreProperties = pool.Allocate<SCORE>(m_engine->num_lex_scores);
    for (size_t i = 0; i < m_engine->num_lex_scores; ++i) {
      tp->scoreProperties[i] = logScores[i + m_engine->num_scores];
    }
    */
  }

  offset += sizeof(float) * totalNumScores;

  // words
  for (size_t targetPos = 0; targetPos < numRealWords; ++targetPos) {
    Word &word = tp->AddWord();
    for (size_t i = 0; i < m_output.size(); ++i) {
      FactorType factorType = m_output[i];

      uint32_t *probingId = (uint32_t*) offset;

      const Factor *factor = GetTargetFactor(*probingId);
      assert(factor);

      word[factorType] = factor;

      offset += sizeof(uint32_t);
    }
  }

  // align
  uint32_t alignTerm = tpInfo->alignTerm;
  //cerr << "alignTerm=" << alignTerm << endl;
  UTIL_THROW_IF2(alignTerm >= m_aligns.size(), "Unknown alignInd");
  tp->SetAlignTerm(m_aligns[alignTerm]);

  // properties TODO

  return tp;
}

//////////////////////////////////////////////////////////////////


ChartRuleLookupManager *ProbingPT::CreateRuleLookupManager(
  const ChartParser &,
  const ChartCellCollectionBase &,
  std::size_t)
{
  abort();
  return NULL;
}

TO_STRING_BODY(ProbingPT);

// friend
ostream& operator<<(ostream& out, const ProbingPT& phraseDict)
{
  return out;
}

}

	// vim:tabstop=2
	#include "ProbingPT.h"
	#include "moses/StaticData.h"
	#include "moses/FactorCollection.h"
	#include "moses/TargetPhraseCollection.h"
	#include "moses/InputFileStream.h"
	#include "probingpt/querying.h"
	#include "probingpt/probing_hash_utils.h"

	using namespace std;

	namespace Moses
	{
	ProbingPT::ProbingPT(const std::string &line)
	: PhraseDictionary(line,true)
	,m_engine(NULL)
	,load_method(util::POPULATE_OR_READ)
	{
	ReadParameters();

	assert(m_input.size() == 1);
	assert(m_output.size() == 1);
	}

	ProbingPT::~ProbingPT()
	{
	delete m_engine;
	}

	void ProbingPT::Load(AllOptions::ptr const& opts)
	{
	m_options = opts;
	SetFeaturesToApply();

	m_engine = new probingpt::QueryEngine(m_filePath.c_str(), load_method);

	m_unkId = 456456546456;

	FactorCollection &vocab = FactorCollection::Instance();

	// source vocab
	const std::map<uint64_t, std::string> &sourceVocab =
	m_engine->getSourceVocab();
	std::map<uint64_t, std::string>::const_iterator iterSource;
	for (iterSource = sourceVocab.begin(); iterSource != sourceVocab.end();
	++iterSource) {
	string wordStr = iterSource->second;
	//cerr << "wordStr=" << wordStr << endl;

	const Factor *factor = vocab.AddFactor(wordStr);

	uint64_t probingId = iterSource->first;
	size_t factorId = factor->GetId();

	if (factorId >= m_sourceVocab.size()) {
	m_sourceVocab.resize(factorId + 1, m_unkId);
	}
	m_sourceVocab[factorId] = probingId;
	}

	// target vocab
	InputFileStream targetVocabStrme(m_filePath + "/TargetVocab.dat");
	string line;
	while (getline(targetVocabStrme, line)) {
	vector<string> toks = Tokenize(line, "\t");
	UTIL_THROW_IF2(toks.size() != 2, string("Incorrect format:") + line + "\n");

	//cerr << "wordStr=" << toks[0] << endl;

	const Factor *factor = vocab.AddFactor(toks[0]);
	uint32_t probingId = Scan<uint32_t>(toks[1]);

	if (probingId >= m_targetVocab.size()) {
	m_targetVocab.resize(probingId + 1);
	}

	m_targetVocab[probingId] = factor;
	}

	// alignments
	CreateAlignmentMap(m_filePath + "/Alignments.dat");

	// memory mapped file to tps
	string filePath = m_filePath + "/TargetColl.dat";
	file.open(filePath.c_str());
	if (!file.is_open()) {
	throw "Couldn't open file ";
	}

	data = file.data();
	//size_t size = file.size();

	// cache
	//CreateCache(system);

	}

	void ProbingPT::CreateAlignmentMap(const std::string path)
	{
	const std::vector< std::vector<unsigned char> > &probingAlignColl = m_engine->getAlignments();
	m_aligns.resize(probingAlignColl.size(), NULL);

	for (size_t i = 0; i < probingAlignColl.size(); ++i) {
	AlignmentInfo::CollType aligns;

	const std::vector<unsigned char> &probingAligns = probingAlignColl[i];
	for (size_t j = 0; j < probingAligns.size(); j += 2) {
	size_t startPos = probingAligns[j];
	size_t endPos = probingAligns[j+1];
	//cerr << "startPos=" << startPos << " " << endPos << endl;
	aligns.insert(std::pair<size_t,size_t>(startPos, endPos));
	}

	const AlignmentInfo *align = AlignmentInfoCollection::Instance().Add(aligns);
	m_aligns[i] = align;
	//cerr << "align=" << align->Debug(system) << endl;
	}
	}

	void ProbingPT::SetParameter(const std::string& key, const std::string& value)
	{
	if (key == "load") {
	if (value == "lazy") {
	load_method = util::LAZY;
	} else if (value == "populate_or_lazy") {
	load_method = util::POPULATE_OR_LAZY;
	} else if (value == "populate_or_read" \|\| value == "populate") {
	load_method = util::POPULATE_OR_READ;
	} else if (value == "read") {
	load_method = util::READ;
	} else if (value == "parallel_read") {
	load_method = util::PARALLEL_READ;
	} else {
	UTIL_THROW2("load method not supported" << value);
	}
	} else {
	PhraseDictionary::SetParameter(key, value);
	}

	}

	void ProbingPT::InitializeForInput(ttasksptr const& ttask)
	{

	}

	void ProbingPT::GetTargetPhraseCollectionBatch(const InputPathList &inputPathQueue) const
	{
	InputPathList::const_iterator iter;
	for (iter = inputPathQueue.begin(); iter != inputPathQueue.end(); ++iter) {
	InputPath &inputPath = **iter;
	const Phrase &sourcePhrase = inputPath.GetPhrase();

	if (sourcePhrase.GetSize() > m_options->search.max_phrase_length) {
	continue;
	}

	TargetPhraseCollection::shared_ptr tpColl = CreateTargetPhrase(sourcePhrase);
	inputPath.SetTargetPhrases(*this, tpColl, NULL);
	}
	}

	TargetPhraseCollection::shared_ptr ProbingPT::CreateTargetPhrase(const Phrase &sourcePhrase) const
	{
	// create a target phrase from the 1st word of the source, prefix with 'ProbingPT:'
	assert(sourcePhrase.GetSize());

	std::pair<bool, uint64_t> keyStruct = GetKey(sourcePhrase);
	if (!keyStruct.first) {
	return TargetPhraseCollection::shared_ptr();
	}

	// check in cache
	CachePb::const_iterator iter = m_cachePb.find(keyStruct.second);
	if (iter != m_cachePb.end()) {
	//cerr << "FOUND IN CACHE " << keyStruct.second << " " << sourcePhrase.Debug(mgr.system) << endl;
	TargetPhraseCollection *tps = iter->second;
	return TargetPhraseCollection::shared_ptr(tps);
	}

	// query pt
	TargetPhraseCollection *tps = CreateTargetPhrases(sourcePhrase,
	keyStruct.second);
	return TargetPhraseCollection::shared_ptr(tps);
	}

	std::pair<bool, uint64_t> ProbingPT::GetKey(const Phrase &sourcePhrase) const
	{
	std::pair<bool, uint64_t> ret;

	// create a target phrase from the 1st word of the source, prefix with 'ProbingPT:'
	size_t sourceSize = sourcePhrase.GetSize();
	assert(sourceSize);

	uint64_t probingSource[sourceSize];
	GetSourceProbingIds(sourcePhrase, ret.first, probingSource);
	if (!ret.first) {
	// source phrase contains a word unknown in the pt.
	// We know immediately there's no translation for it
	} else {
	ret.second = m_engine->getKey(probingSource, sourceSize);
	}

	return ret;

	}

	void ProbingPT::GetSourceProbingIds(const Phrase &sourcePhrase,
	bool &ok, uint64_t probingSource[]) const
	{

	size_t size = sourcePhrase.GetSize();
	for (size_t i = 0; i < size; ++i) {
	const Word &word = sourcePhrase.GetWord(i);
	uint64_t probingId = GetSourceProbingId(word);
	if (probingId == m_unkId) {
	ok = false;
	return;
	} else {
	probingSource[i] = probingId;
	}
	}

	ok = true;
	}

	uint64_t ProbingPT::GetSourceProbingId(const Word &word) const
	{
	uint64_t ret = 0;

	for (size_t i = 0; i < m_input.size(); ++i) {
	FactorType factorType = m_input[i];
	const Factor *factor = word[factorType];

	size_t factorId = factor->GetId();
	if (factorId >= m_sourceVocab.size()) {
	return m_unkId;
	}
	ret += m_sourceVocab[factorId];
	}

	return ret;
	}

	TargetPhraseCollection *ProbingPT::CreateTargetPhrases(
	const Phrase &sourcePhrase, uint64_t key) const
	{
	TargetPhraseCollection *tps = NULL;

	//Actual lookup
	std::pair<bool, uint64_t> query_result; // 1st=found, 2nd=target file offset
	query_result = m_engine->query(key);
	//cerr << "key2=" << query_result.second << endl;

	if (query_result.first) {
	const char *offset = data + query_result.second;
	uint64_t numTP = (uint64_t) offset;

	tps = new TargetPhraseCollection();

	offset += sizeof(uint64_t);
	for (size_t i = 0; i < *numTP; ++i) {
	TargetPhrase *tp = CreateTargetPhrase(offset);
	assert(tp);
	tp->EvaluateInIsolation(sourcePhrase, GetFeaturesToApply());

	tps->Add(tp);

	}

	tps->Prune(true, m_tableLimit);
	//cerr << *tps << endl;
	}

	return tps;

	}

	TargetPhrase *ProbingPT::CreateTargetPhrase(
	const char *&offset) const
	{
	probingpt::TargetPhraseInfo tpInfo = (probingpt::TargetPhraseInfo) offset;
	size_t numRealWords = tpInfo->numWords / m_output.size();

	TargetPhrase *tp = new TargetPhrase(this);

	offset += sizeof(probingpt::TargetPhraseInfo);

	// scores
	float scores = (float) offset;

	size_t totalNumScores = m_engine->num_scores + m_engine->num_lex_scores;

	if (m_engine->logProb) {
	// set pt score for rule
	tp->GetScoreBreakdown().PlusEquals(this, scores);

	// save scores for other FF, eg. lex RO. Just give the offset
	/*
	if (m_engine->num_lex_scores) {
	tp->scoreProperties = scores + m_engine->num_scores;
	}
	*/
	} else {
	// log score 1st
	float logScores[totalNumScores];
	for (size_t i = 0; i < totalNumScores; ++i) {
	logScores[i] = FloorScore(TransformScore(scores[i]));
	}

	// set pt score for rule
	tp->GetScoreBreakdown().PlusEquals(this, logScores);

	// save scores for other FF, eg. lex RO.
	/*
	tp->scoreProperties = pool.Allocate<SCORE>(m_engine->num_lex_scores);
	for (size_t i = 0; i < m_engine->num_lex_scores; ++i) {
	tp->scoreProperties[i] = logScores[i + m_engine->num_scores];
	}
	*/
	}

	offset += sizeof(float) * totalNumScores;

	// words
	for (size_t targetPos = 0; targetPos < numRealWords; ++targetPos) {
	Word &word = tp->AddWord();
	for (size_t i = 0; i < m_output.size(); ++i) {
	FactorType factorType = m_output[i];

	uint32_t probingId = (uint32_t) offset;

	const Factor factor = GetTargetFactor(probingId);
	assert(factor);

	word[factorType] = factor;

	offset += sizeof(uint32_t);
	}
	}

	// align
	uint32_t alignTerm = tpInfo->alignTerm;
	//cerr << "alignTerm=" << alignTerm << endl;
	UTIL_THROW_IF2(alignTerm >= m_aligns.size(), "Unknown alignInd");
	tp->SetAlignTerm(m_aligns[alignTerm]);

	// properties TODO

	return tp;
	}

	//////////////////////////////////////////////////////////////////


	ChartRuleLookupManager *ProbingPT::CreateRuleLookupManager(
	const ChartParser &,
	const ChartCellCollectionBase &,
	std::size_t)
	{
	abort();
	return NULL;
	}

	TO_STRING_BODY(ProbingPT);

	// friend
	ostream& operator<<(ostream& out, const ProbingPT& phraseDict)
	{
	return out;
	}

	}