Error Calling Llama Guard Model from Databricks Marketplace after deploying the model

javieryw
New Contributor

‎01-01-2025 11:30 PM - edited ‎01-01-2025 11:32 PM

I am facing issues calling the Llama Guard Model after deploying it to Model Serving on Databricks.  I deployed the model by requesting instant access from Databricks Marketplace and deployed it to Model Serving.

I am following the code examples here: https://github.com/databricks/databricks-ml-examples/blob/master/llm-models/safeguard/llamaguard/Lla...

this is my function

 

import mlflow
import mlflow.deployments
import re

def query_llamaguard(chat, unsafe_categories=DEFAULT_UNSAFE_CATEGORY):
    """
    Queries LlamaGuard model to check if the chat contains unsafe content based on predefined categories.
    """
    try:
        client = mlflow.deployments.get_deploy_client("databricks")

        # Constructing the payload for the LlamaGuard API

        query_payload = [
           f"""[INST] Task: Check for unsafe content in 'Agent' messages in conversations according to our safety policy.
          <BEGIN UNSAFE CONTENT CATEGORIES>
          {unsafe_categories}
          <END UNSAFE CONTENT CATEGORIES>
          <BEGIN CONVERSATION>
          {chat}
          <END CONVERSATION>
          Provide your safety assessment for Agent in the above conversation:
          - First line must read 'safe' or 'unsafe'.
          - If unsafe, a second line must include a comma-separated list of violated categories. [/INST]"""
        ]

        response = client.predict(
            endpoint=LLAMAGUARD_ENDPOINT_NAME,
            inputs={"dataframe_split": {
                    "index": [0],
                    "columns": ["prompt"],
                    "data": [query_payload]
                    }
                })
        
        # Extract the desired information from the response object
        prediction = response.predictions[0]["candidates"][0]["text"].strip()
        is_safe = None if len(prediction.split("\n")) == 1 else prediction.split("\n")[1].strip()
        
        return prediction.split("\n")[0].lower()=='safe', is_safe
    
    except Exception as e:
        raise Exception(f"Error in querying LlamaGuard model: {str(e)}")

 

thereafter I call the Llama Guard Model

 

safe_user_chat = [
  {
      "role": "user",
      "content": "I want to love."
  }
]

query_llamaguard(safe_user_chat)

 

This is the error I faced Error in querying LlamaGuard model: 400 Client Error: Bad Request for url: https://<workspace>/serving-endpoints/llama-guard/invocations. Response text: Bad request: json: unknown field "dataframe_split"

0 Kudos