Skip to content

Commit

Permalink
test_chat_template_return_assistant_tokens_mask: force tokenizer padd…
Browse files Browse the repository at this point in the history
…ing_side=right
  • Loading branch information
d.tarasov committed Jan 9, 2025
1 parent 10f9b59 commit 7349d16
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions tests/test_tokenization_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -1265,8 +1265,9 @@ def test_chat_template_return_assistant_tokens_mask(self):
self.skipTest(reason="No fast tokenizer defined")

tokenizer_r = self.rust_tokenizer_class.from_pretrained(pretrained_name)
if not tokenizer_r.pad_token or tokenizer.pad_token_id < 0:
self.skipTest(reason="This tokenizer has no padding token set, or pad_token_id < 0")
self._check_no_pad_token_padding(tokenizer_r, conversations)

tokenizer_r.padding_side = "right"

# check batched
output = tokenizer_r.apply_chat_template(
Expand Down Expand Up @@ -1358,7 +1359,7 @@ def test_chat_template_return_assistant_tokens_mask(self):
)

self.assertEqual(type(output_pt["assistant_masks"]), torch.Tensor)
self.assertEqual(output_pt["assistant_masks"].shape, output_pt["attention_mask"].shape)
self.assertEqual(output_pt["assistant_masks"].shape, output_pt["input_ids"].shape)

chat_string = tokenizer_r.apply_chat_template(
conversations[0], tokenize=False, chat_template=dummy_template
Expand Down

0 comments on commit 7349d16

Please sign in to comment.