Add rich formatting to tracebacks (#927)

* added rich traceback * sorted imports * added rich to eval
mosaicml · Feb 1, 2024 · b928de8 · b928de8
1 parent 169b653
commit b928de8
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/scripts/eval/eval.py b/scripts/eval/eval.py
@@ -18,9 +18,11 @@
 from composer.utils import dist, get_device, reproducibility
 from omegaconf import DictConfig, ListConfig
 from omegaconf import OmegaConf as om
+from rich.traceback import install
 from transformers import (AutoModelForCausalLM, PreTrainedTokenizerBase,
                           T5ForConditionalGeneration)
 
+install()
 from llmfoundry.models import MPTForCausalLM
 from llmfoundry.models.model_registry import COMPOSER_MODEL_REGISTRY
 from llmfoundry.utils.builders import (add_metrics_to_eval_loaders,

diff --git a/scripts/train/train.py b/scripts/train/train.py
@@ -20,8 +20,10 @@
 from composer.utils import dist, get_device, reproducibility
 from omegaconf import DictConfig, ListConfig
 from omegaconf import OmegaConf as om
+from rich.traceback import install
 from transformers import PreTrainedTokenizerBase
 
+install()
 from llmfoundry import (COMPOSER_MODEL_REGISTRY, ComposerHFCausalLM,
                         MPTForCausalLM)
 from llmfoundry.callbacks import AsyncEval