generate index of previous states in inference

source/sample/transformer/T2TPredictor.h

@@ -53,6 +53,9 @@ public:
     /* predictions */
     XTensor prediction;
     
+    /* id of the previous state that generates the current one  */
+    XTensor preID;
+
     /* score of every prediction (last state of the path) */
     XTensor score;
 

source/sample/transformer/T2TSearch.cpp

@@ -95,6 +95,7 @@ void T2TSearch::Generate(T2TStateBundle * beam)
     XTensor scoreTopK;
     XTensor &score = beam->score;
     XTensor &index = beam->prediction;
+    XTensor &preID = beam->preID;
     int order = score.order;
 
     CheckNTErrors(order >= 2, "The tensor must be of order 2 or larger.");
@@ -115,12 +116,22 @@ void T2TSearch::Generate(T2TStateBundle * beam)
                1.0F, score.devID, score.mem);
     InitTensor(&index, order, dimsTopK, X_INT,
                1.0F, score.devID, score.mem);
+    InitTensor(&preID, order, dimsTopK, X_INT,
+               1.0F, -1);
     
     score.Reshape(order, dimsBeam);
     
     /* keep the most promissing candidates in the beam */
     TopK(score, scoreTopK, index, 0, beamSize);
     
+    CopyValues(scoreTopK, preID);
+    
+    int sizePredict = score.GetDim(-1);
+    
+    /* pre id !!! */
+    
+    /* mod !!! */
+
     score.Reshape(order, dims);
 }
 
@@ -134,6 +145,7 @@ void T2TSearch::Expand(T2TStateBundle * prev, T2TStateBundle * beam)
 
     T2TState * states = beam->states;
     XTensor &predict = beam->prediction;
+    
     XTensor index = *NewTensorBuf(predict.order - 1, predict.dimSize, X_FLOAT, 1.0F, 
                                   predict.devID, predict.mem);