Fix : input_ids sequence length problem

helper_functions.py

@@ -112,7 +112,7 @@ def transform_single_text(
         tokens = tokenize_whole_text(text, tokenizer)
     input_id_chunks, mask_chunks = split_tokens_into_smaller_chunks(tokens, chunk_size, stride, minimal_chunk_length)
     add_special_tokens_at_beginning_and_end(input_id_chunks, mask_chunks)
-    input_ids, attention_mask = stack_tokens_from_all_chunks(input_id_chunks, mask_chunks)
+    add_padding_tokens(input_id_chunks, mask_chunks , chunk_size)
     input_ids, attention_mask = stack_tokens_from_all_chunks(input_id_chunks, mask_chunks)
     return input_ids, attention_mask
 
@@ -158,19 +158,16 @@ def add_special_tokens_at_beginning_and_end(input_id_chunks: list[Tensor], mask_
         mask_chunks[i] = torch.cat([Tensor([1]), mask_chunks[i], Tensor([1])])
 
 
-
-def add_padding_tokens(input_id_chunks: list[Tensor], mask_chunks: list[Tensor], chunk_size: int) -> None:
-    """Adds padding tokens at the end to make sure that all chunks have exactly chunk_size tokens."""
-    pad_token_id = 0  # Assuming this is defined somewhere in your code
+def add_padding_tokens(input_id_chunks: list[Tensor], mask_chunks: list[Tensor] , chunk_size) -> None:
+    """Adds padding tokens (token id = 0) at the end to make sure that all chunks have exactly 512 tokens."""
     for i in range(len(input_id_chunks)):
         # get required padding length
-        pad_len = chunk_size +2 - input_id_chunks[i].shape[0]
+        pad_len = chunk_size + 2 - input_id_chunks[i].shape[0]
         # check if tensor length satisfies required chunk size
         if pad_len > 0:
             # if padding length is more than 0, we must add padding
-            input_id_chunks[i] = torch.cat([input_id_chunks[i], torch.tensor([pad_token_id] * pad_len)])
-            mask_chunks[i] = torch.cat([mask_chunks[i], torch.tensor([0] * pad_len)])
-
+            input_id_chunks[i] = torch.cat([input_id_chunks[i], Tensor([0] * pad_len)])
+            mask_chunks[i] = torch.cat([mask_chunks[i], Tensor([0] * pad_len)])
 
 
 def stack_tokens_from_all_chunks(input_id_chunks: list[Tensor], mask_chunks: list[Tensor]) -> tuple[Tensor, Tensor]:
@@ -191,6 +188,13 @@ def split_overlapping(tensor: Tensor, chunk_size: int, stride: int, minimal_chun
 
 ## Voice part
 
+def stack_tokens_from_all_chunks_for_inference(input_id_chunks: list[Tensor], mask_chunks: list[Tensor]) -> tuple[Tensor, Tensor]:
+    """Reshapes data to a form compatible with BERT model input."""
+    input_ids = torch.stack(input_id_chunks)
+    attention_mask = torch.stack(mask_chunks)
+
+    return input_ids.long(), attention_mask.int()
+
 def transform_for_inference_text(text: str,
     tokenizer: PreTrainedTokenizerBase,
     chunk_size: int,
@@ -204,7 +208,7 @@ def transform_for_inference_text(text: str,
     input_id_chunks, mask_chunks = split_tokens_into_smaller_chunks(tokens, chunk_size, stride, minimal_chunk_length)
     add_special_tokens_at_beginning_and_end_inference(input_id_chunks, mask_chunks)
     add_padding_tokens_inference(input_id_chunks, mask_chunks, chunk_size)
-    input_ids, attention_mask = stack_tokens_from_all_chunks(input_id_chunks, mask_chunks)
+    input_ids, attention_mask = stack_tokens_from_all_chunks_for_inference(input_id_chunks, mask_chunks)
     return {"input_ids": input_ids, "attention_mask": attention_mask}
 
 def add_special_tokens_at_beginning_and_end_inference(input_id_chunks: list[Tensor], mask_chunks: list[Tensor]) -> None: