diff --git a/litgpt/prompts.py b/litgpt/prompts.py index e513e6102e..79062adba7 100644 --- a/litgpt/prompts.py +++ b/litgpt/prompts.py @@ -133,10 +133,7 @@ def stop_tokens(self, tokenizer: "Tokenizer") -> Tuple[List[int], ...]: class Falcon(PromptStyle): def apply(self, prompt: str, **kwargs: str) -> str: - # First line could be modified. AFAIK Falcon doesn't impose a specific system prompt - # The instruction to not prefix its replies doesn't work always, but better than nothing - # I've also tried just "{prompt}\n" but the model seems to ramble more often - return f"Do not prefix your replies with 'Bot: '\nUser: {prompt}\n" + return f"{prompt}\nAnswer:" def stop_tokens(self, tokenizer: "Tokenizer") -> Tuple[List[int], ...]: return ( @@ -218,14 +215,14 @@ def apply(self, prompt: Union[str, List[Dict[str, str]]], **kwargs: str) -> str: def encode_header(role: str) -> List[str]: return [f"<|start_header_id|>{role}<|end_header_id|>\n\n"] - + def encode_message(message: Dict[str, str]) -> List[str]: tokens = encode_header(message["role"]) # NOTE: Meta stripped this. I'm not sure I agree, but who am I to argue? tokens.append(message["content"].strip()) tokens.append("<|eot_id|>") return tokens - + def has_system_prompt(messages: List[Dict[str, str]]) -> bool: return messages[0].get("role", "") == "system" if len(messages) else False @@ -309,7 +306,6 @@ def apply(self, prompt: str, **kwargs: str) -> str: return f'<|system|>\nYou are a helpful assistant.<|end|>\n<|user|>\n{prompt}<|end|>\n<|assistant|>\n' - class TinyLlama(PromptStyle): def apply(self, prompt: str, **kwargs: str) -> str: return (