|
""" |
|
HF Chat Templates prompt strategy |
|
""" |
|
|
|
import logging |
|
from typing import Any, Dict, List, Optional |
|
|
|
from axolotl.prompt_tokenizers import PromptTokenizingStrategy |
|
from axolotl.prompters import Prompter |
|
from axolotl.utils.chat_templates import chat_templates |
|
|
|
LOG = logging.getLogger("axolotl") |
|
|
|
|
|
class ChatTemplatePrompter(Prompter): |
|
"""prompter for HF chat templates""" |
|
|
|
def __init__( |
|
self, |
|
tokenizer, |
|
chat_template=None, |
|
max_length=2048, |
|
message_field_role: str = "from", |
|
message_field_content: str = "value", |
|
roles: Optional[Dict[str, List[str]]] = None, |
|
): |
|
if roles: |
|
self.roles = {s: t for t, sources in roles.items() for s in sources} |
|
else: |
|
self.roles = { |
|
"human": "user", |
|
"user": "user", |
|
"assistant": "assistant", |
|
"gpt": "assistant", |
|
"system": "system", |
|
} |
|
self.message_field_role = message_field_role |
|
self.message_field_content = message_field_content |
|
self.tokenizer = tokenizer |
|
self.chat_template = chat_template |
|
self.max_length = max_length |
|
|
|
def build_prompt(self, conversation, add_generation_prompt=False): |
|
turns = [ |
|
{ |
|
"role": self.roles[t[self.message_field_role]], |
|
"content": t[self.message_field_content], |
|
} |
|
for t in conversation |
|
] |
|
|
|
return self.tokenizer.apply_chat_template( |
|
turns, |
|
truncation=True, |
|
max_length=self.max_length, |
|
add_generation_prompt=add_generation_prompt, |
|
chat_template=self.chat_template, |
|
) |
|
|
|
|
|
class ChatTemplateStrategy(PromptTokenizingStrategy): |
|
""" |
|
Tokenizing strategy for instruction-based prompts. |
|
""" |
|
|
|
_messages = "conversations" |
|
|
|
@property |
|
def messages(self): |
|
return self._messages |
|
|
|
@messages.setter |
|
def messages(self, messages): |
|
self._messages = messages |
|
|
|
def tokenize_prompt(self, prompt): |
|
turns = self.get_conversation_thread(prompt) |
|
prompt_ids = self.prompter.build_prompt(turns[:-1], add_generation_prompt=True) |
|
input_ids = self.prompter.build_prompt(turns) |
|
|
|
if not self.train_on_inputs: |
|
user_prompt_len = len(prompt_ids) |
|
labels = [-100] * user_prompt_len + input_ids[user_prompt_len:] |
|
else: |
|
labels = input_ids |
|
|
|
tokenized_prompt = { |
|
"input_ids": input_ids, |
|
"labels": labels, |
|
"attention_mask": [1] * len(input_ids), |
|
} |
|
|
|
return tokenized_prompt |
|
|
|
def get_conversation_thread(self, prompt): |
|
return prompt[self.messages] |
|
|
|
|
|
def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None): |
|
chat_template = ( |
|
ds_cfg["chat_template"] if ds_cfg and "chat_template" in ds_cfg else "chatml" |
|
) |
|
message_field_role = ( |
|
ds_cfg["message_field_role"] |
|
if ds_cfg and "message_field_role" in ds_cfg |
|
else "from" |
|
) |
|
message_field_content = ( |
|
ds_cfg["message_field_content"] |
|
if ds_cfg and "message_field_content" in ds_cfg |
|
else "value" |
|
) |
|
roles = ds_cfg["roles"] if ds_cfg and "roles" in ds_cfg else None |
|
|
|
strategy = ChatTemplateStrategy( |
|
ChatTemplatePrompter( |
|
tokenizer, |
|
chat_templates(chat_template), |
|
message_field_role=message_field_role, |
|
message_field_content=message_field_content, |
|
roles=roles, |
|
), |
|
tokenizer, |
|
cfg.train_on_inputs, |
|
cfg.sequence_len, |
|
) |
|
if ds_cfg and "field_messages" in ds_cfg and hasattr(strategy, "messages"): |
|
strategy.messages = ds_cfg["field_messages"] |
|
return strategy |
|
|
