Fix bug introduced in WordpieceBatch change
@@ -256,8 +256,7 @@ class FullTransformerBatch:
for doc_spans in self.spans:
if len(doc_spans) == 0 or len(doc_spans[0]) == 0:
outputs.append(TransformerData.empty())
- token_count = 0
- else:
+ continue
start_i = token_positions[doc_spans[0][0]]
end_i = token_positions[doc_spans[-1][-1]] + 1
end = start + len(doc_spans)
@@ -271,7 +270,6 @@ class FullTransformerBatch:
align=doc_align,
)
)
- token_count = sum(doc_tokens.lengths)
- prev_tokens += token_count
+ prev_tokens += doc_tokens.input_ids.size
start += len(doc_spans)
return outputs
