rulm

1
from transformers import AutoTokenizer, AutoConfig
2

3
from src.util.dl import fix_tokenizer
4

5

6
def test_fix_tokenizer():
7
    model_name = "TheBloke/Llama-2-7B-fp16"
8
    tokenizer = AutoTokenizer.from_pretrained(model_name)
9
    config = AutoConfig.from_pretrained(model_name)
10
    tokenizer = fix_tokenizer(tokenizer, config)
11
    assert tokenizer.bos_token_id == 1
12
    assert tokenizer.eos_token_id == 2
13
    assert tokenizer.pad_token_id == 0
14
    assert tokenizer.unk_token_id == 0
15
    assert tokenizer.model_max_length == 4096
16

17
    model_name = "ai-forever/ruGPT-3.5-13B"
18
    tokenizer = AutoTokenizer.from_pretrained(model_name)
19
    config = AutoConfig.from_pretrained(model_name)
20
    tokenizer = fix_tokenizer(tokenizer, config)
21
    assert tokenizer.pad_token_id == 0
22
    assert tokenizer.bos_token_id == 2
23
    assert tokenizer.eos_token_id == 3
24
    assert tokenizer.unk_token_id == 1
25
    assert tokenizer.model_max_length == 2048
26