File size: 2,060 Bytes
03f35b0 9d1523a 03f35b0 9d1523a 03f35b0 9d1523a 03f35b0 |
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 |
# prompt_builder.py
from typing import Protocol, List, Tuple
from transformers import AutoTokenizer
class PromptTemplate(Protocol):
"""Protocol for prompt templates."""
def format(self, context: str, user_input: str, chat_history: List[Tuple[str, str]], **kwargs) -> str:
pass
class LlamaPromptTemplate:
def format(self, context: str, user_input: str, chat_history: List[Tuple[str, str]], max_history_turns: int = 1) -> str:
system_message = f"Please assist based on the following context: {context}"
prompt = f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{system_message}<|eot_id|>"
for user_msg, assistant_msg in chat_history[-max_history_turns:]:
prompt += f"<|start_header_id|>user<|end_header_id|>\n\n{user_msg}<|eot_id|>"
prompt += f"<|start_header_id|>assistant<|end_header_id|>\n\n{assistant_msg}<|eot_id|>"
prompt += f"<|start_header_id|>user<|end_header_id|>\n\n{user_input}<|eot_id|>"
prompt += "<|start_header_id|>assistant<|end_header_id|>\n\n"
return prompt
class TransformersPromptTemplate:
def __init__(self, model_path: str):
self.tokenizer = AutoTokenizer.from_pretrained(model_path)
def format(self, context: str, user_input: str, chat_history: List[Tuple[str, str]], **kwargs) -> str:
messages = [
{
"role": "system",
"content": f"Please assist based on the following context: {context}",
}
]
for user_msg, assistant_msg in chat_history:
messages.extend([
{"role": "user", "content": user_msg},
{"role": "assistant", "content": assistant_msg}
])
messages.append({"role": "user", "content": user_input})
tokenized_chat = self.tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
return tokenized_chat
|