Chris4K/A.I.StateMachine

1# prompt_builder.py2from typing import Protocol, List, Tuple3from transformers import AutoTokenizer4 5 6class PromptTemplate(Protocol):7    """Protocol for prompt templates."""8    def format(self, context: str, user_input: str, chat_history: List[Tuple[str, str]], **kwargs) -> str:9        pass10 11 12class LlamaPromptTemplate:13    def format(self, context: str, user_input: str, chat_history: List[Tuple[str, str]], max_history_turns: int = 1) -> str:14        system_message = f"Please assist based on the following context: {context}"15        prompt = f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{system_message}<|eot_id|>"16        17        for user_msg, assistant_msg in chat_history[-max_history_turns:]:18            prompt += f"<|start_header_id|>user<|end_header_id|>\n\n{user_msg}<|eot_id|>"19            prompt += f"<|start_header_id|>assistant<|end_header_id|>\n\n{assistant_msg}<|eot_id|>"20            21        prompt += f"<|start_header_id|>user<|end_header_id|>\n\n{user_input}<|eot_id|>"22        prompt += "<|start_header_id|>assistant<|end_header_id|>\n\n"23        return prompt24        25 26class TransformersPromptTemplate:27    def __init__(self, model_path: str):28        self.tokenizer = AutoTokenizer.from_pretrained(model_path)29        30    def format(self, context: str, user_input: str, chat_history: List[Tuple[str, str]], **kwargs) -> str:31        messages = [32            {33                "role": "system",34                "content": f"Please assist based on the following context: {context}",35            }36        ]37        38        for user_msg, assistant_msg in chat_history:39            messages.extend([40                {"role": "user", "content": user_msg},41                {"role": "assistant", "content": assistant_msg}42            ])43            44        messages.append({"role": "user", "content": user_input})45        46        tokenized_chat = self.tokenizer.apply_chat_template(47            messages,48            tokenize=False,49            add_generation_prompt=True50        )51        return tokenized_chat52