With the rise of Generative AI, I’m curious if MLflow has tools specifically for tracking LLM inputs and outputs. I want to compare different prompts for the same model and see which ones yield better factual accuracy. Is there a specific flavor or UI component designed for this type of unstructured "run" data, or is it strictly for scalar metrics?
3 answers
Absolutely! MLflow has introduced "MLflow Tracing" and specialized LLM tracking features specifically for this purpose. You can use the mlflow.llm.log_predictions() function to store prompts, outputs, and even the intermediate chain steps. The UI has been updated with a dedicated "Evaluation" tab where you can view these side-by-side. It even supports "LLM-as-a-judge" metrics, where you can use a stronger model (like GPT-4) to automatically score the outputs of your smaller models based on relevance or safety. This makes MLflow a very strong contender for LLMOps, moving way beyond just tracking simple training loss curves
How do you handle the cost of using an 'LLM-as-a-judge' when you have thousands of test cases in your evaluation set?
The new 'Prompt Engineering' UI in MLflow is a lifesaver. You can iterate on prompts directly in the browser and see the results instantly.
I was just about to mention that! Being able to tweak a system message and see how it affects a whole batch of responses in one view is a game changer.
Cost management is tricky, George. In MLflow, the best strategy is to run your automated judges on a representative sample rather than the full set. You can also use smaller, open-source models like Llama-3 or Mistral as the evaluator to keep costs at zero. The tracking server doesn't care which model provides the score; it just records the 'accuracy' or 'faithfulness' value you give it, allowing you to optimize your budget.