[Test] Basic checks on chat templates (#944)

Author: riedgar-ms

guidance/chat.py

@@ -190,16 +190,16 @@ class Phi3MiniChatTemplate(ChatTemplate):
 
     def get_role_start(self, role_name):
         if role_name == "user":
-            return "<|user|>"
+            return "<|user|>\n"
         elif role_name == "assistant":
-            return "<|assistant|>"
+            return "<|assistant|>\n"
         elif role_name == "system":
-            return "<|system|>"
+            return "<|system|>\n"
         else:
             raise UnsupportedRoleException(role_name, self)
 
     def get_role_end(self, role_name=None):
-        return "<|end|>"
+        return "<|end|>\n"
 
 
 CHAT_TEMPLATE_CACHE[phi3_mini_template] = Phi3MiniChatTemplate
@@ -219,14 +219,14 @@ class Phi3SmallMediumChatTemplate(ChatTemplate):
 
     def get_role_start(self, role_name):
         if role_name == "user":
-            return "<|user|>"
+            return "<|user|>\n"
         elif role_name == "assistant":
-            return "<|assistant|>"
+            return "<|assistant|>\n"
         else:
             raise UnsupportedRoleException(role_name, self)
 
     def get_role_end(self, role_name=None):
-        return "<|end|>"
+        return "<|end|>\n"
 
 
 CHAT_TEMPLATE_CACHE[phi3_small_template] = Phi3SmallMediumChatTemplate

tests/need_credentials/test_chat_templates.py

@@ -1,7 +1,9 @@
 import pytest
+import transformers
+
+import guidance
 
 from guidance.chat import CHAT_TEMPLATE_CACHE
-import transformers
 
 from ..utils import env_or_fail
 
@@ -23,8 +25,6 @@ def test_popular_models_in_cache(model_id: str, should_pass: bool):
     # If this fails, the models have had their templates updated, and we need to fix the cache manually.
     hf_token = env_or_fail("HF_TOKEN")
 
-    # model_id, should_pass = model_info
-
     tokenizer = transformers.AutoTokenizer.from_pretrained(
         model_id, token=hf_token, trust_remote_code=True
     )
@@ -38,3 +38,75 @@ def test_popular_models_in_cache(model_id: str, should_pass: bool):
 
 # TODO: Expand testing to verify that tokenizer.apply_chat_template() produces same results as our ChatTemplate subclasses
 # once I hook up the new ChatTemplate to guidance.models.Transformers and guidance.models.LlamaCPP, we can do this
+
+
+@pytest.mark.parametrize(
+    "model_id",
+    [
+        "microsoft/Phi-3-mini-4k-instruct",
+        "microsoft/Phi-3-small-8k-instruct",
+        "microsoft/Phi-3-medium-4k-instruct",
+        "meta-llama/Meta-Llama-3-8B-Instruct",
+        "meta-llama/Llama-2-7b-chat-hf",
+        "mistralai/Mistral-7B-Instruct-v0.2",
+    ],
+)
+def test_chat_format_smoke(model_id: str):
+    hf_token = env_or_fail("HF_TOKEN")
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        model_id, token=hf_token, trust_remote_code=True
+    )
+    model_chat_template = tokenizer.chat_template
+
+    lm = guidance.models.Mock("")
+    lm.chat_template = CHAT_TEMPLATE_CACHE[model_chat_template]()
+
+    messages = [
+        {"role": "user", "content": "Good day to you!"},
+        {"role": "assistant", "content": "Hello!"},
+    ]
+    tokeniser_render = tokenizer.apply_chat_template(messages, tokenize=False)
+
+    with guidance.user():
+        lm += "Good day to you!"
+    with guidance.assistant():
+        lm += "Hello!"
+    # Only check substring due to BOS/EOS tokens
+    assert str(lm) in tokeniser_render
+
+
+@pytest.mark.parametrize(
+    "model_id",
+    [
+        "microsoft/Phi-3-mini-4k-instruct",
+        "meta-llama/Meta-Llama-3-8B-Instruct",
+        "meta-llama/Llama-2-7b-chat-hf",
+    ],
+)
+def test_chat_format_smoke_with_system(model_id: str):
+    hf_token = env_or_fail("HF_TOKEN")
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        model_id, token=hf_token, trust_remote_code=True
+    )
+    model_chat_template = tokenizer.chat_template
+
+    lm = guidance.models.Mock("")
+    lm.chat_template = CHAT_TEMPLATE_CACHE[model_chat_template]()
+
+    messages = [
+        {"role": "system", "content": "You are an LLM"},
+        {"role": "user", "content": "Good day to you!"},
+        {"role": "assistant", "content": "Hello!"},
+    ]
+    tokeniser_render = tokenizer.apply_chat_template(messages, tokenize=False)
+
+    with guidance.system():
+        lm += "You are an LLM"
+    with guidance.user():
+        lm += "Good day to you!"
+    with guidance.assistant():
+        lm += "Hello!"
+    # Only check substring due to BOS/EOS tokens
+    assert str(lm) in tokeniser_render