LLM-as-a-judge with MLflow 3.0

MLflow 3.0 introduces expanded support for LLM-as-a-judge workflows, allowing developers to build both out-of-the-box judges for toxicity and safety and domain-specific custom judges, according to a technical walkthrough of the release. The post details three judge types and shows a code example using the mlflow.genai.judges make_judge function with a response_quality judge that scores answers on a 0-to-1 rubric. The author notes the Databricks Generative AI Certification Exam covers only MLflow 2.0, so the examples use version 3.0 for its added LLMOps flexibility.

I recently finished taking my Databricks Generative AI Certification Exam. https://medium.com/@parksjr/databricks-generative-ai-certification-studying-ee30943d201e In short, it was filled with great content and I learned a good amount. It ranged from reminders of the Gen AI basics, to what tools to use in Databricks for AI, to how a robust AI system should be created. In the meantime, one of the most valuable takeaways for me was the section on LLMOps. While many people today can make a notebook that uses an LLM or vibe code their way to it, I was looking to learn how to take my AI projects to the next level. The questions I wanted answered were: While the Databricks certification only showed how to use MLflow 2.0 to answer these questions, MLflow did go through a recent major version update to MLflow 3.0 https://mlflow.org/releases/3 . This newer version allows more flexibility in functions to use with AI related projects, models, evaluations, along with functions specific to handling LLMOps. So while the certification provided examples for how to use MLflow version 2.0, for my project and the examples I will be providing here, I am using version 3.0. As you can imagine, I used the newest version since it is best practice and because of the increase in ease of use for LLMOps. For this series, I will look specifically at how MLflow 3.0 can be used to create an LLM-as-a-judge. This will cover the range of need for judges from out-of-the-box judges that assess things like toxicity, safety, etc. to creating domain specific custom judges. As we all know, LLMs aren’t perfect. I recently had one misspell “commented” as “commmented”. While a human may make that typo when typing, I cannot think of a human that would actually think an English word has three “m”s together. The two “m”s alone get the job done. As a result, it’s important we evaluate our LLM outputs; they should not be taken at face value. It is easy to validate an output if you are only evaluating a few at a time, but what if you have thousands of outputs coming in each day? It is not possible for a human to go through and check each and everyone of of these. An LLM-as-a-judge is a where an LLM is used to evaluate the output of another LLM. These evaluation can look at if the output is safe, if it is grammatically correct, or even check domain specific materials, such as if it included accurate domain information. As mentioned earlier, MLflow 3.0 has created more flexibility for integrating it with AI, including creating different types of judges. Below, we will highlight three different judges and what they should be used for. Here is a brief explanation of what each label means: Please note, there are best practices around LLM-as-a-judge that aren’t not covered in this post. Flexibility : High Amount of Time : High Requirements: Sample labeled dataset, prompt, judge name Returns: User provided scoring metric When to Use: When extensive prompt, few-shot examples, many rules, etc. are needed to produce accurate results Notes : To improve performance of this judge, it is helpful to include examples of inputs and outputs that would receive certain scores and rationales; this is shown in the example below. Please see the hyperlink for additional parameters. Here is the definition of main ones used: python import mlflowfrom mlflow.genai.judges import make judgefrom typing import Literal Create a judge that evaluates response quality using template variablesquality judge = make judge name="response quality", description="Evaluate how well the response answers the question", instructions= """ Instructions Evaluate if the response in {{ outputs }} correctly answers the question in {{ inputs }}. The response should be accurate, complete, and professional. Rubric 0: Does not answer question at all 0.5: Partially answers question 1: Completely answers question Examples Input: What is the capitol of France? Output: Paris Score: 1 Rationale: The input only required one word answer and is correct Input: How many schools are in the Big Ten and which ones are in Michigan? Output: There are 18 schools in the Big Ten. Score: 0.5 Rationale: The response correctly answered the number of schools in the Big Ten but did not answer the other question asking how many are in Michigan. """ , model="openai:/gpt-4", feedback value type=int Flexibility : Medium Amount of Time : Medium Returns : Pass/Fail When to Use: When a prompt with a few examples or rules are needed; stick a few paragraphs max python import mlflowfrom mlflow.genai.scorers import Guidelinesenglish = Guidelines name="english", guidelines= "The response must be in English" Available out-of-the-box judges listed in hyperlinked title above. Flexibility : Least Time : Low Requirements : None When to Use: When standard judge topics cover data evaluation needs python from mlflow.genai.scorers import Safety, Correctnessresults = mlflow.genai.evaluate data=sample data, scorers= Safety , Correctness , , Lastly, in order to actually implement these judges, you will need to have a table with inputs and outputs as their own columns. MLFlow will then iterate through each row and evaluate the data using the judge s . In order to set up your evaluation data correctly, it must follow the convention below: create evaluationd data from dummy data frame df with columns data, responsesample data = pd.DataFrame { "inputs": {"data":x} for x in df.data, "outputs": {"response":y} for y in df.response } generate results for all judgesresults = mlflow.genai.evaluate data=sample data, scorers= Safety , Correctness , english, quality judge , LLMs-as-a-judge are not only great tools to evaluate qualitative responses, but can also be reviewing large amounts of data for specific errors that would require many manual hours to carefully review and try to catch. With the various ways to set up these judges with MLFlow, you can speed up this part of your pipeline. By having pre-defined structures, you can quickly fill in the judge with your data, all while having the flexibility to evaluate the entire response down to aspects as small as punctuation and specific word use. This is part of my Databricks Generative AI Certification Exam Series. Here are some other articles from it: Databricks Generative AI Certification: Study https://medium.com/@parksjr/databricks-generative-ai-certification-studying-ee30943d201e Databricks Generative AI Certification: Exam https://medium.com/@parksjr/databricks-generative-ai-certification-exam-b9fbcce80c36 LLM-as-a-judge with MLflow 3.0 https://pub.towardsai.net/llm-as-a-judge-with-mlflow-3-0-9bd80c826735 was originally published in Towards AI https://pub.towardsai.net on Medium, where people are continuing the conversation by highlighting and responding to this story.