Callback Handler for MLflow #4150
Conversation
Implementation of callback handler for one of the most popular mlops platform - mlflow. Inspired by wandb_callback handler. This handler logs the metrics and artifacts to the mlflow server.
Update code to report metrics in batch Support to include 'step' in the metric
- Added support to report chain assets using mlflow.langchain.log_model. Using log_artifact for other types. - Fixed 'input' key error when prompt templates are used. - Improved the reporting logic by using mlflow active run context
- Added support to report chain assets using mlflow.langchain.log_model. Using log_artifact for other types. - Fixed 'input' key error when prompt templates are used. - Improved the reporting logic by using mlflow active run context Add example notebook and fix
|
@mahmedk if you don't mind, I didn't have write permissions on your fork so I added the sample notebook while preserving your commits so you get the credit in the commit history :) EDIT: and I've added you as a collaborator on my fork in case you want to update it! @hwchase17 this includes the sample notebook you requested and I've tested it works with the local MLflow tracking store |
|
@agola11 thanks for the review and apologies for not following the contributions instructions earlier - I've run the make lint + format now so things should be happy (and removed my change from Chain base). After this PR, I am willing to add a few more tests in a follow up - are you willing to add mlflow + azureml-mlflow packages as test dependencies or should we explore some other options? |
There was a problem hiding this comment.
LGTM with a nit comment: https://github.com/hwchase17/langchain/pull/4150/files#r1190588611. Thanks @akshaya-a @hwchase17 !
There was a problem hiding this comment.
LGTM with a nit comment: https://github.com/hwchase17/langchain/pull/4150/files#r1190588611. Thanks @akshaya-a @hwchase17 !
Co-authored-by: Corey Zumar <39497902+dbczumar@users.noreply.github.com>
Rebased Mahmedk's PR with the callback refactor and added the example requested by hwchase plus a couple minor fixes --------- Co-authored-by: Ahmed K <77802633+mahmedk@users.noreply.github.com> Co-authored-by: Ahmed K <mda3k27@gmail.com> Co-authored-by: Davis Chase <130488702+dev2049@users.noreply.github.com> Co-authored-by: Corey Zumar <39497902+dbczumar@users.noreply.github.com> Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>
Rebased Mahmedk's PR with the callback refactor and added the example requested by hwchase plus a couple minor fixes --------- Co-authored-by: Ahmed K <77802633+mahmedk@users.noreply.github.com> Co-authored-by: Ahmed K <mda3k27@gmail.com> Co-authored-by: Davis Chase <130488702+dev2049@users.noreply.github.com> Co-authored-by: Corey Zumar <39497902+dbczumar@users.noreply.github.com> Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>
@akshaya-a thanks much ! Sorry, I could not attend to this work for past couple of weeks due to internal commitments. Glad to see the PR getting merged and mlflow support in langchain. Thanks Again :-) |
Rebased Mahmedk's PR with the callback refactor and added the example requested by hwchase plus a couple minor fixes