|
|
|
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 |
|
|
