Merge pull request #148 from georgian-io/better-logging

Better Handling of Logs and Warnings for UI
georgian-io · Apr 10, 2024 · 98e1ad4 · 98e1ad4
2 parents a57022b + 0ea4ea5
commit 98e1ad4
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/llmtune/cli/toolkit.py b/llmtune/cli/toolkit.py
@@ -3,10 +3,10 @@
 from pathlib import Path
 
 import torch
+import transformers
 import typer
 import yaml
 from pydantic import ValidationError
-from transformers import utils as hf_utils
 from typing_extensions import Annotated
 
 import llmtune
@@ -20,8 +20,9 @@
 from llmtune.utils.save_utils import DirectoryHelper
 
 
-hf_utils.logging.set_verbosity_error()
+transformers.logging.set_verbosity(transformers.logging.CRITICAL)
 torch._logging.set_logs(all=logging.CRITICAL)
+logging.captureWarnings(True)
 
 
 app = typer.Typer()