From 0ea4ea5750fe4094cfb04d90321632b129a7a1b6 Mon Sep 17 00:00:00 2001
From: Benjamin Ye <benjamin.ye@georgian.io>
Date: Wed, 10 Apr 2024 04:15:10 +0000
Subject: [PATCH] suppress warnings

---
 llmtune/cli/toolkit.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/llmtune/cli/toolkit.py b/llmtune/cli/toolkit.py
index 749e557..ca06680 100644
--- 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()