test_tokens.py 1.39 KB
Newer Older
1
from text_generation_server.utils.tokens import (
2
3
    StopSequenceCriteria,
    StoppingCriteria,
4
    FinishReason,
5
6
7
)


8
def test_stop_sequence_criteria():
9
    criteria = StopSequenceCriteria("/test;")
10

11
12
13
14
    assert not criteria("/")
    assert not criteria("/test")
    assert criteria("/test;")
    assert not criteria("/test; ")
15
16


17
18
19
20
21
22
23
24
25
def test_stop_sequence_criteria_escape():
    criteria = StopSequenceCriteria("<|stop|>")

    assert not criteria("<")
    assert not criteria("<|stop")
    assert criteria("<|stop|>")
    assert not criteria("<|stop|> ")


26
27
28
def test_stopping_criteria():
    criteria = StoppingCriteria(0, [StopSequenceCriteria("/test;")], max_new_tokens=5)
    assert criteria(65827, "/test") == (False, None)
29
    assert criteria(30, ";") == (True, FinishReason.FINISH_REASON_STOP_SEQUENCE)
30
31


32
33
34
def test_stopping_criteria_eos():
    criteria = StoppingCriteria(0, [StopSequenceCriteria("/test;")], max_new_tokens=5)
    assert criteria(1, "") == (False, None)
35
    assert criteria(0, "") == (True, FinishReason.FINISH_REASON_EOS_TOKEN)
36
37
38


def test_stopping_criteria_max():
39
40
41
42
43
    criteria = StoppingCriteria(0, [StopSequenceCriteria("/test;")], max_new_tokens=5)
    assert criteria(1, "") == (False, None)
    assert criteria(1, "") == (False, None)
    assert criteria(1, "") == (False, None)
    assert criteria(1, "") == (False, None)
44
    assert criteria(1, "") == (True, FinishReason.FINISH_REASON_LENGTH)