-
-
Notifications
You must be signed in to change notification settings - Fork 780
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add: mlflow for experiment tracking (#1059) [skip ci]
* Update requirements.txt adding mlflow * Update __init__.py Imports for mlflow * Update README.md * Create mlflow_.py (#1) * Update README.md * fix precommits * Update README.md Update mlflow_tracking_uri * Update trainer_builder.py update trainer building * chore: lint * make ternary a bit more readable --------- Co-authored-by: Wing Lian <wing.lian@gmail.com>
- Loading branch information
Showing
5 changed files
with
34 additions
and
2 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -22,6 +22,7 @@ hf_transfer | |
colorama | ||
numba | ||
numpy>=1.24.4 | ||
mlflow | ||
# qlora things | ||
bert-score==0.3.13 | ||
evaluate==0.4.0 | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,18 @@ | ||
"""Module for mlflow utilities""" | ||
|
||
import os | ||
|
||
from axolotl.utils.dict import DictDefault | ||
|
||
|
||
def setup_mlflow_env_vars(cfg: DictDefault): | ||
for key in cfg.keys(): | ||
if key.startswith("mlflow_"): | ||
value = cfg.get(key, "") | ||
|
||
if value and isinstance(value, str) and len(value) > 0: | ||
os.environ[key.upper()] = value | ||
|
||
# Enable mlflow if experiment name is present | ||
if cfg.mlflow_experiment_name and len(cfg.mlflow_experiment_name) > 0: | ||
cfg.use_mlflow = True |