[Build] Fix CI and Notebook builds (#940)

Get the CI and Notebook builds working again. The main fix involves the authentication to the AzureAI endpoints; we now have an Entra ID to use, rather than an `api_key`. The instructions are [on the Microsoft site](https://learn.microsoft.com/en-us/azure/developer/github/connect-from-azure).

Also update the chat templates again, for Phi3 and Mistral. The fix for Mistral is not ideal. It appears that it now 'supports' a `system` message by inserting it into the first `user` message. I don't think our approach can really support this.

Author: riedgar-ms

.github/workflows/ci_tests.yml

@@ -13,6 +13,10 @@ on:
     # Run at 1030 UTC every day
     - cron:  '30 10 * * *'
 
+permissions:
+  id-token: write
+  contents: read
+
 jobs:
   build:
 
@@ -24,15 +28,15 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
       - name: Install Rust
         shell: bash
         run: |
            curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain 1.75.0
            echo "$HOME/.cargo/bin" >> $GITHUB_PATH
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python-version }}
       - name: Show GPUs
         run: |
           nvidia-smi
@@ -57,15 +61,23 @@ jobs:
       - name: Check GPU available
         run: |
           python -c "import torch; assert torch.cuda.is_available()"
+      - name: 'Az CLI login'
+        uses: azure/login@v1
+        with:
+            client-id: ${{ secrets.AZURE_CLIENT_ID }}
+            tenant-id: ${{ secrets.AZURE_TENANT_ID }}
+            subscription-id: ${{ secrets.AZURE_SUBSCRIPTION_ID }}
+      - name: 'Run Azure CLI commands'
+        run: |
+            az account show
+            az group list
       - name: Test with pytest
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
           # Configure endpoints for Azure OpenAI
           AZUREAI_CHAT_ENDPOINT: ${{ secrets.AZUREAI_CHAT_ENDPOINT }}
-          AZUREAI_CHAT_KEY: ${{ secrets.AZUREAI_CHAT_KEY }}
           AZUREAI_CHAT_MODEL: ${{ secrets.AZUREAI_CHAT_MODEL }}
           AZUREAI_COMPLETION_ENDPOINT: ${{ secrets.AZUREAI_COMPLETION_ENDPOINT }}
-          AZUREAI_COMPLETION_KEY: ${{ secrets.AZUREAI_COMPLETION_KEY }}
           AZUREAI_COMPLETION_MODEL: ${{ secrets.AZUREAI_COMPLETION_MODEL }}
           # Configure endpoints for Azure AI Studio
           AZURE_AI_STUDIO_PHI3_ENDPOINT: ${{ vars.AZURE_AI_STUDIO_PHI3_ENDPOINT }}

.github/workflows/notebook_tests.yml

@@ -16,6 +16,10 @@ on:
     # Run at 0830 UTC every day
     - cron:  '30 08 * * *'
 
+permissions:
+  id-token: write
+  contents: read
+
 jobs:
   build:
 
@@ -60,15 +64,23 @@ jobs:
       - name: Check GPU available
         run: |
           python -c "import torch; assert torch.cuda.is_available()"
+      - name: 'Az CLI login'
+        uses: azure/login@v1
+        with:
+            client-id: ${{ secrets.AZURE_CLIENT_ID }}
+            tenant-id: ${{ secrets.AZURE_TENANT_ID }}
+            subscription-id: ${{ secrets.AZURE_SUBSCRIPTION_ID }}
+      - name: 'Run Azure CLI commands'
+        run: |
+            az account show
+            az group list
       - name: Test with pytest
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
           # Configure endpoints for Azure OpenAI
           AZUREAI_CHAT_ENDPOINT: ${{ secrets.AZUREAI_CHAT_ENDPOINT }}
-          AZUREAI_CHAT_KEY: ${{ secrets.AZUREAI_CHAT_KEY }}
           AZUREAI_CHAT_MODEL: ${{ secrets.AZUREAI_CHAT_MODEL }}
           AZUREAI_COMPLETION_ENDPOINT: ${{ secrets.AZUREAI_COMPLETION_ENDPOINT }}
-          AZUREAI_COMPLETION_KEY: ${{ secrets.AZUREAI_COMPLETION_KEY }}
           AZUREAI_COMPLETION_MODEL: ${{ secrets.AZUREAI_COMPLETION_MODEL }}
           # Configure endpoints for Azure AI Studio
           AZURE_AI_STUDIO_PHI3_ENDPOINT: ${{ vars.AZURE_AI_STUDIO_PHI3_ENDPOINT }}

guidance/chat.py

@@ -37,7 +37,8 @@ def __contains__(self, key: str):
 
 # 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() 
+CHAT_TEMPLATE_CACHE = ChatTemplateCache()
+
 
 class UnsupportedRoleException(Exception):
     def __init__(self, role_name, instance):
@@ -46,11 +47,12 @@ def __init__(self, role_name, 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. ")
+        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
@@ -60,23 +62,27 @@ def load_template_class(chat_template=None):
     """
     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.")
+            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.
+        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.""")
-    
+                        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
 
@@ -94,15 +100,18 @@ def _template_class_from_string(template_str):
 # --------------------------------------------------
 # 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
 
 
@@ -111,6 +120,8 @@ def get_role_end(self, role_name=None):
 # --------------------------------------------------
 # [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
@@ -124,7 +135,7 @@ def get_role_start(self, role_name):
             return " "
         else:
             raise UnsupportedRoleException(role_name, self)
-        
+
     def get_role_end(self, role_name=None):
         if role_name == "system":
             return "\n<</SYS>"
@@ -135,6 +146,7 @@ def get_role_end(self, role_name=None):
         else:
             raise UnsupportedRoleException(role_name, self)
 
+
 CHAT_TEMPLATE_CACHE[llama2_template] = Llama2ChatTemplate
 
 
@@ -143,6 +155,8 @@ def get_role_end(self, role_name=None):
 # --------------------------------------------------
 # [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
@@ -156,52 +170,89 @@ def get_role_start(self, role_name):
             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):
+phi3_mini_template = "{% for message in messages %}{% if message['role'] == 'system' %}{{'<|system|>\n' + message['content'] + '<|end|>\n'}}{% elif message['role'] == 'user' %}{{'<|user|>\n' + message['content'] + '<|end|>\n'}}{% elif message['role'] == 'assistant' %}{{'<|assistant|>\n' + message['content'] + '<|end|>\n'}}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|assistant|>\n' }}{% else %}{{ eos_token }}{% endif %}"
+
+
+class Phi3MiniChatTemplate(ChatTemplate):
     # available_roles = ["user", "assistant"]
-    template_str = phi3_template
+    template_str = phi3_mini_template
 
     def get_role_start(self, role_name):
         if role_name == "user":
             return "<|user|>"
         elif role_name == "assistant":
             return "<|assistant|>"
+        elif role_name == "system":
+            return "<|system|>"
         else:
             raise UnsupportedRoleException(role_name, self)
-        
+
     def get_role_end(self, role_name=None):
         return "<|end|>"
 
-CHAT_TEMPLATE_CACHE[phi3_template] = Phi3ChatTemplate
 
+CHAT_TEMPLATE_CACHE[phi3_mini_template] = Phi3MiniChatTemplate
+
+# https://huggingface.co/microsoft/Phi-3-small-8k-instruct/blob/main/tokenizer_config.json
+phi3_small_template = "{{ bos_token }}{% for message in messages %}{{'<|' + message['role'] + '|>' + '\n' + message['content'] + '<|end|>\n' }}{% endfor %}{% if add_generation_prompt %}{{ '<|assistant|>\n' }}{% else %}{{ eos_token }}{% endif %}"
+
+
+# https://huggingface.co/microsoft/Phi-3-medium-4k-instruct/blob/main/tokenizer_config.json#L119
+phi3_medium_template = "{% 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 %}"
+
+
+# Although the templates are different, the roles are the same between medium and small (for now)
+class Phi3SmallMediumChatTemplate(ChatTemplate):
+    # available_roles = ["user", "assistant"]
+    template_str = phi3_small_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_small_template] = Phi3SmallMediumChatTemplate
+CHAT_TEMPLATE_CACHE[phi3_medium_template] = Phi3SmallMediumChatTemplate
 
 # --------------------------------------------------
 # @@@@ 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 %}"
+mistral_7b_instruct_template = "{%- if messages[0]['role'] == 'system' %}\n    {%- set system_message = messages[0]['content'] %}\n    {%- set loop_messages = messages[1:] %}\n{%- else %}\n    {%- set loop_messages = messages %}\n{%- endif %}\n\n{{- bos_token }}\n{%- for message in loop_messages %}\n    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}\n        {{- raise_exception('After the optional system message, conversation roles must alternate user/assistant/user/assistant/...') }}\n    {%- endif %}\n    {%- if message['role'] == 'user' %}\n        {%- if loop.first and system_message is defined %}\n            {{- ' [INST] ' + system_message + '\\n\\n' + message['content'] + ' [/INST]' }}\n        {%- else %}\n            {{- ' [INST] ' + message['content'] + ' [/INST]' }}\n        {%- endif %}\n    {%- elif message['role'] == 'assistant' %}\n        {{- ' ' + message['content'] + eos_token}}\n    {%- else %}\n        {{- raise_exception('Only user and assistant roles are supported, with the exception of an initial optional system message!') }}\n    {%- endif %}\n{%- endfor %}\n"
+
+
 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] "
+            return " [INST] "
         elif role_name == "assistant":
             return " "
+        elif role_name == "system":
+            raise ValueError("Please include system instructions in the first user message")
         else:
             raise UnsupportedRoleException(role_name, self)
-        
+
     def get_role_end(self, role_name=None):
         if role_name == "user":
             return " [/INST]"
@@ -210,4 +261,5 @@ def get_role_end(self, role_name=None):
         else:
             raise UnsupportedRoleException(role_name, self)
 
+
 CHAT_TEMPLATE_CACHE[mistral_7b_instruct_template] = Mistral7BInstructChatTemplate

guidance/models/_azure_guidance.py

@@ -4,7 +4,7 @@
 import json
 import urllib.parse
 from ._model import Engine, Model, EngineCallResponse
-from ..chat import Phi3ChatTemplate
+from ..chat import Phi3MiniChatTemplate
 from ._byte_tokenizer import ByteTokenizer
 
 
@@ -30,7 +30,7 @@ def __init__(self, server_url, max_streaming_tokens=1000, chat_template=None):
 
         if chat_template is None:
             # TODO [PK]: obtain this from the server
-            chat_template=Phi3ChatTemplate
+            chat_template=Phi3MiniChatTemplate
 
         tokenizer = ByteTokenizer(chat_template)
 

notebooks/api_examples/models/AzureOpenAI.ipynb

@@ -36,7 +36,7 @@
     "import os\n",
     "\n",
     "# Uncomment if using DefaultAzureCredential below\n",
-    "# from azure.identity import DefaultAzureCredential, get_bearer_token_provider\n",
+    "from azure.identity import DefaultAzureCredential, get_bearer_token_provider\n",
     "\n",
     "# This is the name of the model deployed, such as 'gpt-4' or 'gpt-3.5-turbo\n",
     "model = os.getenv(\"AZUREAI_CHAT_MODEL\", \"Please set the model\")\n",
@@ -52,13 +52,13 @@
     "azure_api_version = os.getenv(\"AZUREAI_CHAT_API_VERSION\", \"Please set the API version\")\n",
     "\n",
     "# The environment variable should be set to the API key from the Azure AI playground:\n",
-    "api_key=os.getenv(\"AZUREAI_CHAT_KEY\", \"Please set API key\")\n",
+    "# api_key=os.getenv(\"AZUREAI_CHAT_KEY\", \"Please set API key\")\n",
     "\n",
     "# Alternatively, we can use Entra authentication\n",
-    "# token_provider = get_bearer_token_provider(\n",
-    "#     DefaultAzureCredential(),\n",
-    "#     \"https://cognitiveservices.azure.com/.default\"\n",
-    "#)"
+    "token_provider = get_bearer_token_provider(\n",
+    "     DefaultAzureCredential(),\n",
+    "     \"https://cognitiveservices.azure.com/.default\"\n",
+    ")"
    ]
   },
   {
@@ -84,9 +84,9 @@
     "    azure_deployment=azure_deployment,\n",
     "    version=azure_api_version,\n",
     "    # For authentication, use either\n",
-    "    api_key=api_key\n",
+    "    # api_key=api_key\n",
     "    # or\n",
-    "    # azure_ad_token_provider=token_provider\n",
+    "    azure_ad_token_provider=token_provider\n",
     ")"
    ]
   },