Spaces:
Sleeping
Sleeping
File size: 3,780 Bytes
f6d689c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
import unittest
from unittest.mock import MagicMock, patch
import torch
from src.data.tokenization import Tokenizer, TokenizerConfig
class TestTokenizer(unittest.TestCase):
@patch("src.data.tokenization.AutoTokenizer")
def test_tokenizer_initialization(self, mock_auto_tokenizer):
mock_hf_tokenizer = MagicMock()
mock_hf_tokenizer.pad_token_id = 0
mock_hf_tokenizer.bos_token_id = 1
mock_hf_tokenizer.eos_token_id = 2
mock_hf_tokenizer.vocab_size = 1000
mock_auto_tokenizer.from_pretrained.return_value = mock_hf_tokenizer
config = TokenizerConfig(pretrained_model_name="test-model")
tokenizer = Tokenizer(config)
self.assertEqual(tokenizer.pad_token_id, 0)
self.assertEqual(tokenizer.bos_token_id, 1)
self.assertEqual(tokenizer.eos_token_id, 2)
self.assertEqual(tokenizer.vocab_size, 1000)
mock_auto_tokenizer.from_pretrained.assert_called_with("test-model")
@patch("src.data.tokenization.AutoTokenizer")
def test_encode(self, mock_auto_tokenizer):
mock_hf_tokenizer = MagicMock()
mock_hf_tokenizer.pad_token_id = 0
mock_hf_tokenizer.bos_token_id = 1
mock_hf_tokenizer.eos_token_id = 2
mock_hf_tokenizer.encode.return_value = [10, 11, 12]
mock_auto_tokenizer.from_pretrained.return_value = mock_hf_tokenizer
tokenizer = Tokenizer()
ids = tokenizer.encode("hello world")
self.assertEqual(ids, [10, 11, 12])
mock_hf_tokenizer.encode.assert_called()
@patch("src.data.tokenization.AutoTokenizer")
def test_batch_encode(self, mock_auto_tokenizer):
mock_hf_tokenizer = MagicMock()
mock_hf_tokenizer.pad_token_id = 0
mock_hf_tokenizer.bos_token_id = 1
mock_hf_tokenizer.eos_token_id = 2
# Mock return value for __call__
mock_hf_tokenizer.return_value = {
"input_ids": torch.tensor([[10, 11], [12, 13]]),
"attention_mask": torch.tensor([[1, 1], [1, 1]]),
}
mock_auto_tokenizer.from_pretrained.return_value = mock_hf_tokenizer
tokenizer = Tokenizer()
output = tokenizer.batch_encode(["hello", "world"])
self.assertIn("input_ids", output)
self.assertIn("attention_mask", output)
self.assertIsInstance(output["input_ids"], torch.Tensor)
self.assertIsInstance(output["attention_mask"], torch.Tensor)
@patch("src.data.tokenization.AutoTokenizer")
def test_decode(self, mock_auto_tokenizer):
mock_hf_tokenizer = MagicMock()
mock_hf_tokenizer.pad_token_id = 0
mock_hf_tokenizer.bos_token_id = 1
mock_hf_tokenizer.eos_token_id = 2
mock_hf_tokenizer.decode.return_value = "hello world"
mock_auto_tokenizer.from_pretrained.return_value = mock_hf_tokenizer
tokenizer = Tokenizer()
text = tokenizer.decode([10, 11, 12])
self.assertEqual(text, "hello world")
mock_hf_tokenizer.decode.assert_called()
@patch("src.data.tokenization.AutoTokenizer")
def test_prepare_decoder_inputs(self, mock_auto_tokenizer):
mock_hf_tokenizer = MagicMock()
mock_hf_tokenizer.pad_token_id = 0
mock_hf_tokenizer.bos_token_id = 1
mock_hf_tokenizer.eos_token_id = 2
mock_auto_tokenizer.from_pretrained.return_value = mock_hf_tokenizer
tokenizer = Tokenizer()
labels = torch.tensor([[10, 11, 2], [12, 2, 0]]) # 0 is pad
decoder_inputs = tokenizer.prepare_decoder_inputs(labels)
# Should shift right and prepend BOS (1)
expected = torch.tensor([[1, 10, 11], [1, 12, 2]])
self.assertTrue(torch.equal(decoder_inputs, expected))
if __name__ == "__main__":
unittest.main()
|