fix w2ner.pieces2word nil slice issues

sunhailin-Leo · Jul 8, 2024 · 2dba54d · 2dba54d
1 parent 6e384dd
commit 2dba54d
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/models/transformers/bert_w2ner.go b/models/transformers/bert_w2ner.go
@@ -111,15 +111,11 @@ func (w *W2NerModelService) getBertInputFeature(batchInferData [][]string) []*W2
 		for j := 0; j < padTokenLen; j++ {
 			gridMask2d[j] = generateAllTrueSlice(padTokenLen)
 
-			if j >= len(inferTokens) {
+			if j >= len(inferTokens) || len(inferTokens[j]) == 0 {
 				pieces2word[j] = make([]bool, len(batchInputFeatures[i].TokenIDs))
 				continue
 			}
 
-			if len(inferTokens[j]) == 0 {
-				continue
-			}
-
 			idx := utils.GenerateRange[int](start, start+len(inferTokens[j]))
 			pieces2word[j] = make([]bool, len(batchInputFeatures[i].TokenIDs))
 			for k := 0; k < len(pieces2word[j]); k++ {