[QEff. Finetune]: Added fix for pad_to_max_length in tokenization. (#599)

Signed-off-by: meetkuma <[email protected]>

Author: quic-meetkuma

QEfficient/finetune/dataset/alpaca_dataset.py

@@ -58,10 +58,15 @@ def __getitem__(self, index):
         else:
             prompt = PROMPT_DICT["prompt_input"].format_map(ann)
         example = prompt + ann["output"]
+
+        if self.context_length is not None:
+            padding_type = "max_length"
+        else:
+            padding_type = True
         prompt = torch.tensor(
-            self.tokenizer.encode(prompt, max_length=self.context_length, pad_to_max_length=True), dtype=torch.int64
+            self.tokenizer.encode(prompt, max_length=self.context_length, padding=padding_type), dtype=torch.int64
         )
-        example = self.tokenizer.encode(example, max_length=self.context_length, pad_to_max_length=True)
+        example = self.tokenizer.encode(example, max_length=self.context_length, padding=padding_type)
         example.append(self.tokenizer.eos_token_id)
         example = torch.tensor(example, dtype=torch.int64)
         labels = copy.deepcopy(example)

QEfficient/finetune/dataset/custom_dataset/sample_dataset_preproc.py

@@ -61,17 +61,22 @@ def apply_prompt_template(sample):
     dataset = dataset.map(apply_prompt_template, remove_columns=list(dataset.features))
 
     def tokenize_add_label(sample):
+        if context_length is not None:
+            padding_type = "max_length"
+        else:
+            padding_type = True
+
         input = tokenizer.encode(
             tokenizer.bos_token + sample["input"],
             add_special_tokens=False,
             max_length=context_length,
-            pad_to_max_length=True,
+            padding=padding_type,
         )
         label = tokenizer.encode(
             sample["label"] + tokenizer.pad_token + tokenizer.eos_token,
             add_special_tokens=False,
             max_length=context_length,
-            pad_to_max_length=True,
+            padding=padding_type,
         )
 
         sample = {

QEfficient/finetune/dataset/grammar_dataset.py

@@ -44,17 +44,23 @@ def convert_to_features(self, example_batch):
         target_ = example_batch["target"]
 
         prompt = f"Correct this to standard English: {input_}\n---\nCorrected: "
+
+        if self.context_length is not None:
+            padding_type = "max_length"
+        else:
+            padding_type = True
+
         prompt_ids = self.tokenizer.encode(
             self.tokenizer.bos_token + prompt,
             add_special_tokens=False,
             max_length=self.context_length,
-            pad_to_max_length=True,
+            padding=padding_type,
         )
         label_ids = self.tokenizer.encode(
             target_ + self.tokenizer.eos_token,
             add_special_tokens=False,
             max_length=self.context_length,
-            pad_to_max_length=True,
+            padding=padding_type,
         )
 
         sample = {