[WIP] Remove Instruct/Chat versions of models & introduce a new ChatTemplate API, fix Anthropic API (#820)

This PR should significantly reduce the number of user-facing classes we
have in Guidance, and reduce subtle bugs introduced by using a
mis-specified Chat Template (models currently silently default to the
ChatML syntax, which many of the latest models don't adhere to). It
should also make it easier for users to add new models to guidance,
either via PR or in their own codebases.

Before:

```python
from guidance.models.transformers import Transformers, TransformersChat

class Llama(Transformers):
    pass

# Users have to do this for most chat models in guidance
class LlamaChat(TransformersChat, Llama):
    def get_role_start(self, role_name, **kwargs):
        if role_name == "system":
            return self._system_prefex + "<<SYS>>\n"
        elif role_name == "user":
            if str(self).endswith("\n<</SYS>>\n\n"):
                return ""  
            else:
                return "[INST] "
        else:
            return " "

    def get_role_end(self, role_name=None):
        if role_name == "system":
            return "\n<</SYS>>\n\n"
        elif role_name == "user":
            return " [/INST]"
        else:
            return " "

lm = LlamaChat(path_to_llama)
```

After:
```python
from guidance.models import Transformers

lm = Transformers(path_to_llama) # automagically works
```

If you're using a rare model and the auto import doesn't automatically
work...

After pt2:
```python
# users can copy paste from https://huggingface.co/meta-llama/Llama-2-7b-chat-hf/blob/main/tokenizer_config.json#L12 
llama2_template = "{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% set system_message = false %}{% endif %}{% for message in loop_messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if loop.index0 == 0 and system_message != false %}{% set content = '<<SYS>>\\n' + system_message + '\\n<</SYS>>\\n\\n' + message['content'] %}{% else %}{% set content = message['content'] %}{% endif %}{% if message['role'] == 'user' %}{{ bos_token + '[INST] ' + content.strip() + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ ' '  + content.strip() + ' ' + eos_token }}{% endif %}{% endfor %}"

lm = Transformers(path_to_llama, chat_template=llama2_template)
```

or, in the worst case for maximal robustness and customizability:
```python
from guidance._chat import ChatTemplate, UnsupportedRoleException

class Llama2ChatTemplate(ChatTemplate):
    template_str = llama2_template

    def get_role_start(self, role_name):
        if role_name == "system":
            return "[INST] <<SYS>>\n"
        elif role_name == "user":
            return "<s>[INST]"
        elif role_name == "assistant":
            return " "
        else:
            raise UnsupportedRoleException(role_name, self)
        
    def get_role_end(self, role_name=None):
        if role_name == "system":
            return "\n<</SYS>"
        elif role_name == "user":
            return " [/INST]"
        elif role_name == "assistant":
            return "</s>"
        else:
            raise UnsupportedRoleException(role_name, self)

lm = Transformers(path_to_llama, chat_template=Llama2ChatTemplate)
```

The first big change is the removal of the `Chat` and `Instruct` mixins,
and an introduction of a new `guidance._chat.ChatTemplate` class, which
handles the same responsibilities as those mixins used to.

Users can construct a subclass of `ChatTemplate` and pass it to models
with a new `chat_template` argument (defaulted to None).

The way this works for local models is to leverage the `chat_template`
property in `huggingface transformers` and `llamacpp`'s GGUF files. When
a user tries to load a model, guidance now follows the following order
of operations:

1. See if the user passed in a `ChatTemplate` -- if so, use that
directly.
2. If the user passed string "chat_template", set that as a
"template_str". If the user did not pass in anything, set the
"template_str" based on metadata from the huggingface.AutoTokenizer or
gguf metadata fields.
3. Check the `template_str` against a local cache in guidance which
maintains template converters for the most popular models on
huggingface/in llama.cpp. We index this cache based on the actual
chat_template string, so any model that uses one of these chat templates
-- even if it isn't explicitly listed in the documentation -- will
automatically load the right guidance class.
4. If we don't have anything in the cache, try to automatically convert
the jinja2 template into the new guidance._chat.ChatTemplate syntax.
Warn the user if we attempt this. [NOTE: This is not yet implemented and
may come in a future PR.]
5. Default to the `ChatML` syntax, with a warning to the user.

Currently this PR updates the following user facing `guidance.Models`
classes:

- Transformers (removing TransformersChat and the Llama/Mistral
subclasses)
- LlamaCpp (removing LlamaCppChat and Mistral subclasses)
- Anthropic

For now, `Anthropic` should be representative of how grammarless classes
will work. I wanted to start with OpenAI, but many other guidance.models
classes inherit from OpenAI, so I'm saving that for later. Also while I
was at it I upgraded the `Anthropic` class to use their latest SDK, so
`guidance.models.Anthropic` should now work with the latest Claude3
models.

# TODO

A decent amount left to do here. In no particular order...

1. Unrelated to this change, but guidance cannot properly handle
`llama3` or `phi-3`'s true tokenizers/chat templates. We need to fix
that independently.
2. Add better warnings to users when we fall back to using the ChatML
template.
3. Extend support through the rest of the guidance.models classes
(mostly just remote ones left, starting with OpenAI).
4. Write out more templates for popular models in the ChatTemplateCache,
and also add an alias system so that we can look up models in the cache
by common names (e.g. "llama3").
5. Add a deprecation warning to people trying to use very old models on
`Anthropic`.
6. Much more testing and documentation. We should, for example, add
documentation on how to import/initialize a new ChatTemplate and use it
for your own models.
7. Write the auto-converter from huggingface `jinja2` to guidance
ChatTemplate. A battery of unit tests here that compare against the
original `transformers.apply_chat_template` method would make this more
robust. Can be in a future PR as this is complex logic. A start to this
was attempted in #791 by
@ilmarinen, and we could eventually pull this in and expand its
coverage.
8. Probably get rid of the folders in `guidance.models.llama_cpp` and
`guidance.models.transformers` because we don't need to maintain a bunch
of subclasses for them anymore.

Would appreciate any and all feedback, particularly on the logical flow
and new user facing (simpler) API. @marcotcr @paulbkoch @slundberg
@riedgar-ms @hudson-ai

---------

Co-authored-by: Richard Edgar (Microsoft) <[email protected]>
Co-authored-by: Paul Koch <[email protected]>

Author: 3 people

.github/workflows/workflow-pr-gate.yml

@@ -134,6 +134,7 @@ jobs:
     needs:
       - unit-tests-linux-python-other
       - unit-tests-gpu-python-latest
+      - server-tests
     name: End Stage 2
     runs-on: ubuntu-latest
     steps:

.gitignore

@@ -22,4 +22,6 @@ guidance/_rust/Cargo.lock
 
 notebooks/**/*.papermill_out.ipynb
 
-.mypy_cache/*
+.mypy_cache/*
+
+**/scratch.*

README.md

@@ -55,8 +55,8 @@ else:
 ```python
 from guidance import user, assistant
 
-# load a chat model
-chat_lm = models.LlamaCppChat(path)
+# load a model
+chat_lm = models.LlamaCpp(path)
 
 # wrap with chat block contexts
 with user():

guidance/_chat.py

@@ -0,0 +1,210 @@
+import warnings
+import uuid
+import inspect
+
+class ChatTemplate:
+    """Contains template for all chat and instruct tuned models."""
+
+    def get_role_start(self, role_name, **kwargs):
+        raise NotImplementedError(
+            "You need to use a ChatTemplate subclass that overrides the get_role_start method"
+        )
+
+    def get_role_end(self, role_name=None):
+        raise NotImplementedError(
+            "You need to use a ChatTemplate subclass that overrides the get_role_start method"
+        )
+    
+class ChatTemplateCache: 
+    def __init__(self):
+        self._cache = {}
+
+    def __getitem__(self, key):
+        key_compact = key.replace(" ", "")
+        return self._cache[key_compact]
+
+
+    def __setitem__(self, key, value):
+        key_compact = key.replace(" ", "")
+        self._cache[key_compact] = value
+
+    def __contains__(self, key):
+        key_compact = key.replace(" ", "")
+        return key_compact in self._cache
+    
+# Feels weird having to instantiate this, but it's a singleton for all purposes
+# TODO [HN]: Add an alias system so we can instantiate with other simple keys (e.g. "llama2" instead of the full template string)
+CHAT_TEMPLATE_CACHE = ChatTemplateCache() 
+
+class UnsupportedRoleException(Exception):
+    def __init__(self, role_name, instance):
+        self.role_name = role_name
+        self.instance = instance
+        super().__init__(self._format_message())
+
+    def _format_message(self):
+        return (f"Role {self.role_name} is not supported by the {self.instance.__class__.__name__} chat template. ")
+
+def load_template_class(chat_template=None):
+    """Utility method to find the best chat template.
+    
+    Order of precedence:
+    - If it's a chat template class, use it directly
+    - If it's a string, check the cache of popular model templates
+    - If it's a string and not in the cache, try to create a class dynamically
+    - [TODO] If it's a string and can't be created, default to ChatML and raise a warning
+    - If it's None, default to ChatML and raise a warning
+    """
+    if inspect.isclass(chat_template) and issubclass(chat_template, ChatTemplate):
+        if chat_template is ChatTemplate:
+            raise Exception("You can't use the base ChatTemplate class directly. Create or use a subclass instead.")
+        return chat_template
+    
+    elif isinstance(chat_template, str):
+        # First check the cache of popular model types
+        # TODO: Expand keys of cache to include aliases for popular model types (e.g. "llama2, phi3")
+        # Can possibly accomplish this with an "aliases" dictionary that maps all aliases to the canonical key in cache
+        if chat_template in CHAT_TEMPLATE_CACHE:
+            return CHAT_TEMPLATE_CACHE[chat_template]
+        # TODO: Add logic here to try to auto-create class dynamically via _template_class_from_string method
+    
+    # Only warn when a user provided a chat template that we couldn't load
+    if chat_template is not None:
+        warnings.warn(f"""Chat template {chat_template} was unable to be loaded directly into guidance.
+                        Defaulting to the ChatML format which may not be optimal for the selected model. 
+                        For best results, create and pass in a `guidance.ChatTemplate` subclass for your model.""")
+    
+    # By default, use the ChatML Template. Warnings to user will happen downstream only if they use chat roles.
+    return ChatMLTemplate
+        
+
+def _template_class_from_string(template_str):
+    """Utility method to try to create a chat template class from a string."""
+    # TODO: Try to build this, perhaps based on passing unit tests we create?
+    pass
+
+
+# CACHE IMPLEMENTATIONS:
+
+# --------------------------------------------------
+# @@@@ ChatML @@@@
+# --------------------------------------------------
+# Note that all grammarless models will default to this syntax, since we typically send chat formatted messages.
+chatml_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}"
+class ChatMLTemplate(ChatTemplate):
+    template_str = chatml_template
+
+    def get_role_start(self, role_name):
+        return f"<|im_start|>{role_name}\n"
+        
+    def get_role_end(self, role_name=None):
+        return "<|im_end|>\n"
+    
+CHAT_TEMPLATE_CACHE[chatml_template] = ChatMLTemplate
+
+
+# --------------------------------------------------
+# @@@@ Llama-2 @@@@
+# --------------------------------------------------
+# [05/08/24] https://huggingface.co/meta-llama/Llama-2-7b-chat-hf/blob/main/tokenizer_config.json#L12
+llama2_template = "{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% set system_message = false %}{% endif %}{% for message in loop_messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if loop.index0 == 0 and system_message != false %}{% set content = '<<SYS>>\\n' + system_message + '\\n<</SYS>>\\n\\n' + message['content'] %}{% else %}{% set content = message['content'] %}{% endif %}{% if message['role'] == 'user' %}{{ bos_token + '[INST] ' + content.strip() + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ ' '  + content.strip() + ' ' + eos_token }}{% endif %}{% endfor %}"
+class Llama2ChatTemplate(ChatTemplate):
+    # available_roles = ["system", "user", "assistant"]
+    template_str = llama2_template
+
+    def get_role_start(self, role_name):
+        if role_name == "system":
+            return "[INST] <<SYS>>\n"
+        elif role_name == "user":
+            return "<s>[INST]"
+        elif role_name == "assistant":
+            return " "
+        else:
+            raise UnsupportedRoleException(role_name, self)
+        
+    def get_role_end(self, role_name=None):
+        if role_name == "system":
+            return "\n<</SYS>"
+        elif role_name == "user":
+            return " [/INST]"
+        elif role_name == "assistant":
+            return "</s>"
+        else:
+            raise UnsupportedRoleException(role_name, self)
+
+CHAT_TEMPLATE_CACHE[llama2_template] = Llama2ChatTemplate
+
+
+# --------------------------------------------------
+# @@@@ Llama-3 @@@@
+# --------------------------------------------------
+# [05/08/24] https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct/blob/main/tokenizer_config.json#L2053
+llama3_template = "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}"
+class Llama3ChatTemplate(ChatTemplate):
+    # available_roles = ["system", "user", "assistant"]
+    template_str = llama3_template
+
+    def get_role_start(self, role_name):
+        if role_name == "system":
+            return "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n"
+        elif role_name == "user":
+            return "<|start_header_id|>user<|end_header_id>\n\n"
+        elif role_name == "assistant":
+            return "<|start_header_id|>assistant<|end_header_id>\n\n"
+        else:
+            raise UnsupportedRoleException(role_name, self)
+        
+    def get_role_end(self, role_name=None):
+        return "<|eot_id|>"
+    
+CHAT_TEMPLATE_CACHE[llama3_template] = Llama3ChatTemplate
+
+# --------------------------------------------------
+# @@@@ Phi-3 @@@@
+# --------------------------------------------------
+# [05/08/24] https://huggingface.co/microsoft/Phi-3-mini-4k-instruct/blob/main/tokenizer_config.json#L119
+phi3_template = "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') %}{{'<|user|>' + '\n' + message['content'] + '<|end|>' + '\n' + '<|assistant|>' + '\n'}}{% elif (message['role'] == 'assistant') %}{{message['content'] + '<|end|>' + '\n'}}{% endif %}{% endfor %}"
+class Phi3ChatTemplate(ChatTemplate):
+    # available_roles = ["user", "assistant"]
+    template_str = phi3_template
+
+    def get_role_start(self, role_name):
+        if role_name == "user":
+            return "<|user|>"
+        elif role_name == "assistant":
+            return "<|assistant|>"
+        else:
+            raise UnsupportedRoleException(role_name, self)
+        
+    def get_role_end(self, role_name=None):
+        return "<|end|>"
+    
+CHAT_TEMPLATE_CACHE[phi3_template] = Phi3ChatTemplate
+
+
+# --------------------------------------------------
+# @@@@ Mistral-7B-Instruct-v0.2 @@@@
+# --------------------------------------------------
+# [05/08/24] https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2/blob/main/tokenizer_config.json#L42
+mistral_7b_instruct_template = "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}"
+class Mistral7BInstructChatTemplate(ChatTemplate):
+    # available_roles = ["user", "assistant"]
+    template_str = mistral_7b_instruct_template
+
+    def get_role_start(self, role_name):
+        if role_name == "user":
+            return "[INST] "
+        elif role_name == "assistant":
+            return ""
+        else:
+            raise UnsupportedRoleException(role_name, self)
+        
+    def get_role_end(self, role_name=None):
+        if role_name == "user":
+            return " [/INST]"
+        elif role_name == "assistant":
+            return "</s>"
+        else:
+            raise UnsupportedRoleException(role_name, self)
+        
+CHAT_TEMPLATE_CACHE[mistral_7b_instruct_template] = Mistral7BInstructChatTemplate

guidance/library/_role.py

@@ -7,14 +7,10 @@
 span_start = "<||_html:<span style='background-color: rgba(255, 180, 0, 0.3); border-radius: 3px;'>_||>"
 span_end = "<||_html:</span>_||>"
 
-
 @guidance
 def role_opener(lm, role_name, **kwargs):
     indent = getattr(lm, "indent_roles", True)
-    if not hasattr(lm, "get_role_start"):
-        raise Exception(
-            f"You need to use a chat model in order the use role blocks like `with {role_name}():`! Perhaps you meant to use the {type(lm).__name__}Chat class?"
-        )
+
 
     # Block start container (centers elements)
     if indent:
@@ -25,8 +21,17 @@ def role_opener(lm, role_name, **kwargs):
         lm += nodisp_start
     else:
         lm += span_start
-
-    lm += lm.get_role_start(role_name, **kwargs)
+    
+    # TODO [HN]: Temporary change while I instrument chat_template in transformers only.
+    # Eventually have all models use chat_template.
+    if hasattr(lm, "get_role_start"):
+        lm += lm.get_role_start(role_name, **kwargs)
+    elif hasattr(lm, "chat_template"):
+        lm += lm.chat_template.get_role_start(role_name)
+    else:
+        raise Exception(
+            f"You need to use a chat model in order the use role blocks like `with {role_name}():`! Perhaps you meant to use the {type(lm).__name__}Chat class?"
+        )
 
     # End of either debug or HTML no disp block
     if indent:
@@ -46,7 +51,12 @@ def role_closer(lm, role_name, **kwargs):
     else:
         lm += span_start
 
-    lm += lm.get_role_end(role_name)
+    # TODO [HN]: Temporary change while I instrument chat_template in transformers only.
+    # Eventually have all models use chat_template.
+    if hasattr(lm, "get_role_end"):
+        lm += lm.get_role_end(role_name)
+    elif hasattr(lm, "chat_template"):
+        lm += lm.chat_template.get_role_end(role_name)
 
     # End of either debug or HTML no disp block
     if indent:
@@ -60,7 +70,7 @@ def role_closer(lm, role_name, **kwargs):
 
     return lm
 
-
+# TODO HN: Add a docstring to better describe arbitrary role functions
 def role(role_name, text=None, **kwargs):
     if text is None:
         return block(

guidance/models/__init__.py

@@ -1,8 +1,8 @@
 from ._model import Model, Instruct, Chat
 
 # local models
-from .transformers._transformers import Transformers, TransformersChat
-from .llama_cpp import LlamaCpp, LlamaCppChat, MistralInstruct, MistralChat
+from .transformers._transformers import Transformers
+from .llama_cpp import LlamaCpp
 from ._mock import Mock, MockChat
 
 # grammarless models (we can't do constrained decoding for them)
@@ -15,15 +15,12 @@
 )
 from ._azure_openai import (
     AzureOpenAI,
-    AzureOpenAIChat,
-    AzureOpenAICompletion,
-    AzureOpenAIInstruct,
 )
 from ._azureai_studio import AzureAIStudioChat
-from ._openai import OpenAI, OpenAIChat, OpenAIInstruct, OpenAICompletion
+from ._openai import OpenAI
 from ._lite_llm import LiteLLM, LiteLLMChat, LiteLLMInstruct, LiteLLMCompletion
 from ._cohere import Cohere, CohereCompletion, CohereInstruct
-from ._anthropic import Anthropic, AnthropicChat
+from ._anthropic import Anthropic
 from ._googleai import GoogleAI, GoogleAIChat
 from ._togetherai import (
     TogetherAI,