FilesExpand file tree

 main

/

test_llama_chat_format.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

88 lines (79 loc) · 3.3 KB

 main

/

test_llama_chat_format.py

Copy path

Top

File metadata and controls

• Code
• Blame

88 lines (79 loc) · 3.3 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

import json

import jinja2

from llama_cpp import (

ChatCompletionRequestUserMessage,

)

import llama_cpp.llama_types as llama_types

import llama_cpp.llama_chat_format as llama_chat_format

from llama_cpp.llama_chat_format import hf_tokenizer_config_to_chat_formatter

def test_mistral_instruct():

chat_template = "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}"

chat_formatter = jinja2.Template(chat_template)

messages = [

llama_types.ChatCompletionRequestUserMessage(role="user", content="Instruction"),

llama_types.ChatCompletionRequestAssistantMessage(role="assistant", content="Model answer"),

llama_types.ChatCompletionRequestUserMessage(role="user", content="Follow-up instruction"),

]

response = llama_chat_format.format_mistral_instruct(

messages=messages,

)

reference = chat_formatter.render(

messages=messages,

bos_token="<s>",

eos_token="</s>",

)

assert response.prompt == reference

mistral_7b_tokenizer_config = """{

"add_bos_token": true,

"add_eos_token": false,

"added_tokens_decoder": {

"0": {

"content": "<unk>",

"lstrip": false,

"normalized": false,

"rstrip": false,

"single_word": false,

"special": true

},

"1": {

"content": "<s>",

"lstrip": false,

"normalized": false,

"rstrip": false,

"single_word": false,

"special": true

},

"2": {

"content": "</s>",

"lstrip": false,

"normalized": false,

"rstrip": false,

"single_word": false,

"special": true

}

},

"additional_special_tokens": [],

"bos_token": "<s>",

"clean_up_tokenization_spaces": false,

"eos_token": "</s>",

"legacy": true,

"model_max_length": 1000000000000000019884624838656,

"pad_token": null,

"sp_model_kwargs": {},

"spaces_between_special_tokens": false,

"tokenizer_class": "LlamaTokenizer",

"unk_token": "<unk>",

"use_default_system_prompt": false,

"chat_template": "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}"

}"""

def test_hf_tokenizer_config_str_to_chat_formatter():

tokenizer_config = json.loads(mistral_7b_tokenizer_config)

chat_formatter = hf_tokenizer_config_to_chat_formatter(

tokenizer_config

)

chat_formatter_respoonse = chat_formatter(

messages=[

ChatCompletionRequestUserMessage(role="user", content="Hello, world!"),

]

)

assert chat_formatter_respoonse.prompt == ("<s>[INST] Hello, world! [/INST]</s>" "")