/// <summary>
/// Append tokens in beam search results to the source tokens
/// </summary>
/// <param name="srcTokens">(batch_size, seq_len)</param>
/// <param name="beamBatchResults">(beam_search_size, batch_size)</param>
/// <returns>output shape: (beam_search_size, batch_size)</returns>
public static List <List <BeamSearchStatus> > AppendOutputTokens(List <List <int> > srcTokens, List <List <BeamSearchStatus> > beamBatchResults)
{
int beamSize = beamBatchResults.Count;
int batchSize = beamBatchResults[0].Count;
List <List <BeamSearchStatus> > output = new List <List <BeamSearchStatus> >();
for (int i = 0; i < beamSize; i++)
{
List <BeamSearchStatus> batch = new List <BeamSearchStatus>();
output.Add(batch);
for (int j = 0; j < batchSize; j++)
{
BeamSearchStatus item = new BeamSearchStatus();
item.OutputIds.AddRange(srcTokens[j]);
item.OutputIds.AddRange(beamBatchResults[i][j].OutputIds);
item.Score = beamBatchResults[i][j].Score;
batch.Add(item);
}
}
return(output);
}
/// <summary>
/// Initialize beam search status for all beam search and batch results
/// </summary>
/// <param name="beamSearchSize"></param>
/// <param name="batchSize"></param>
/// <param name="tgtTokensList"></param>
/// <returns>shape: (beam_search_size, batch_size)</returns>
public static List <List <BeamSearchStatus> > InitBeamSearchStatusListList(int batchSize, List <List <int> > tgtTokensList)
{
List <List <BeamSearchStatus> > beam2batchStatus = new List <List <BeamSearchStatus> >();
List <BeamSearchStatus> bssList = new List <BeamSearchStatus>();
for (int i = 0; i < batchSize; i++)
{
BeamSearchStatus bss = new BeamSearchStatus();
bss.OutputIds.AddRange(tgtTokensList[i]);
bssList.Add(bss);
}
beam2batchStatus.Add(bssList);
return(beam2batchStatus);
}
