Skip to content

Commit c3a0a92

Browse files
UtkarshTheDevUtkarshTheDev
committed
Created Config and Used New Approach for Model Download Bars Display
1 parent ec8da0b commit c3a0a92

File tree

5 files changed

+195
-100
lines changed

5 files changed

+195
-100
lines changed

‎CHANGELOG.md

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,20 @@
22

33
All notable changes to LocalLab will be documented in this file.
44

5+
## [0.6.4] - 2025-05-16
6+
7+
### Improved
8+
9+
- Completely redesigned model downloading experience with proper native Hugging Face progress bars
10+
- Created new early configuration system to set up logging before any Hugging Face libraries are imported
11+
- Implemented StdoutRedirector to ensure proper display of progress bars during model downloads
12+
- Temporarily disabled all logging handlers during model downloads to prevent interference
13+
- Added clear visual separation between LocalLab logs and Hugging Face progress bars
14+
- Set environment variables to optimize Hugging Face download experience
15+
- Configured transformers library to use native progress bars for model downloads
16+
- Added informative messages before and after model downloads for better user experience
17+
- Ensured consistent progress bar display across different model types and sizes
18+
519
## [0.6.3] - 2025-05-16
620

721
### Improved
@@ -12,9 +26,6 @@ All notable changes to LocalLab will be documented in this file.
1226
- Improved configuration of Hugging Face Hub progress bars for better visual experience
1327
- Completely bypassed custom logging for Hugging Face download-related logs
1428
- Configured transformers library to use native progress bars for model downloads
15-
- Disabled logger propagation for HuggingFace-related modules during downloads
16-
- Added proper spacing before and after progress bars for better readability
17-
- Enhanced progress bar detection to catch all download-related progress indicators
1829

1930
## [0.6.2] - 2025-05-04
2031

‎locallab/__init__.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,21 @@
22
LocalLab - A lightweight AI inference server for running LLMs locally
33
"""
44

5-
__version__ = "0.6.3" # Updated to improve model downloading experience and fix CLI settings
5+
# Import early configuration first to set up logging and environment variables
6+
# This ensures Hugging Face's progress bars are displayed correctly
7+
from .utils.early_config import enable_hf_progress_bars
8+
9+
__version__ = "0.6.4" # Updated to improve model downloading experience and fix CLI settings
610

711
# Only import what's necessary initially, lazy-load the rest
812
from .logger import get_logger
913

1014
# Explicitly expose start_server for direct import
1115
from .server import start_server, cli
1216

17+
# Enable Hugging Face progress bars with native display
18+
enable_hf_progress_bars()
19+
1320
# Other imports will be lazy-loaded when needed
1421
# from .config import MODEL_REGISTRY, DEFAULT_MODEL
1522
# from .model_manager import ModelManager

‎locallab/model_manager.py

Lines changed: 72 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
1+
# Import early configuration module first to set up logging and environment variables
2+
# This ensures Hugging Face's progress bars are displayed correctly
3+
from .utils.early_config import enable_hf_progress_bars, StdoutRedirector
4+
15
from .config import HF_TOKEN_ENV, get_env_var, set_env_var
26
import os
37
import logging
48
import torch
5-
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
69
from typing import Optional, Generator, Dict, Any, List, Union, Callable, AsyncGenerator
710
from fastapi import HTTPException
811
import time
@@ -13,7 +16,6 @@
1316
)
1417
from .logger.logger import logger, log_model_loaded, log_model_unloaded
1518
from .utils import check_resource_availability, get_device, format_model_size
16-
from .utils.progress import configure_hf_hub_progress
1719
import gc
1820
from colorama import Fore, Style
1921
import asyncio
@@ -22,20 +24,12 @@
2224
import tempfile
2325
import json
2426

25-
# Configure HuggingFace Hub progress bars to use native display
27+
# Enable Hugging Face progress bars with native display
2628
# This ensures we see the visually appealing progress bars from HuggingFace
27-
configure_hf_hub_progress()
28-
29-
# Also configure transformers to use HuggingFace Hub's progress bars
30-
try:
31-
import transformers
32-
transformers.utils.logging.enable_progress_bar()
33-
# Set transformers logging to only show warnings and errors
34-
transformers.logging.set_verbosity_warning()
35-
except ImportError:
36-
logger.debug("Could not configure transformers progress bars")
37-
except Exception as e:
38-
logger.debug(f"Error configuring transformers progress bars: {str(e)}")
29+
enable_hf_progress_bars()
30+
31+
# Import transformers after configuring logging to ensure proper display
32+
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
3933

4034
QUANTIZATION_SETTINGS = {
4135
"fp16": {
@@ -280,52 +274,41 @@ async def _load_model_with_optimizations(self, model_id: str):
280274
# Add an empty line to separate from HuggingFace progress bars
281275
print("")
282276

283-
# Set a flag to indicate we're downloading a model
284-
# This will help our logger know to let HuggingFace's progress bars through
285-
try:
286-
# Access the module's global variable
287-
import locallab.utils.progress
288-
locallab.utils.progress.is_downloading = True
289-
290-
# Ensure HuggingFace Hub's progress bars are enabled
291-
from huggingface_hub.utils import logging as hf_logging
292-
hf_logging.enable_progress_bars()
293-
294-
# Configure transformers to use progress bars
295-
import transformers
296-
transformers.utils.logging.enable_progress_bar()
297-
298-
# Also ensure tqdm is properly configured for nice display
299-
import tqdm
300-
tqdm.tqdm.monitor_interval = 0 # Disable monitor thread which can cause issues
277+
# Add an empty line before progress bars start
278+
print(f"\n{Fore.CYAN}Starting model download - native progress bars will appear below{Style.RESET_ALL}\n")
279+
280+
# Enable Hugging Face progress bars again to ensure they're properly configured
281+
enable_hf_progress_bars()
282+
283+
# Use a context manager to ensure proper display of Hugging Face progress bars
284+
with StdoutRedirector(disable_logging=True):
285+
# Load tokenizer first
286+
logger.info(f"Loading tokenizer for {model_id}...")
287+
self.tokenizer = AutoTokenizer.from_pretrained(
288+
model_id,
289+
token=hf_token if hf_token else None
290+
)
291+
logger.info(f"Tokenizer loaded successfully")
301292

302-
# Temporarily disable our custom logger for HuggingFace logs
303-
import logging
304-
for logger_name in ['tqdm', 'huggingface_hub', 'transformers', 'filelock']:
305-
logging.getLogger(logger_name).handlers = [] # Remove any handlers
306-
logging.getLogger(logger_name).propagate = False # Don't propagate to parent loggers
307-
except:
308-
# Fallback if import fails
309-
pass
293+
# Load model with optimizations
294+
logger.info(f"Loading model weights for {model_id}...")
310295

311-
# Add an empty line before progress bars start
312-
print("\n")
296+
# This is the critical part where we want to see nice progress bars
297+
# We'll temporarily disable our logger's handlers to prevent interference
298+
root_logger = logging.getLogger()
299+
original_handlers = root_logger.handlers.copy()
300+
root_logger.handlers = []
313301

314-
# Load tokenizer first
315-
logger.info(f"Loading tokenizer for {model_id}...")
316-
self.tokenizer = AutoTokenizer.from_pretrained(
317-
model_id,
318-
token=hf_token if hf_token else None
319-
)
320-
logger.info(f"Tokenizer loaded successfully")
321-
322-
# Load model with optimizations
323-
logger.info(f"Loading model weights for {model_id}...")
324-
self.model = AutoModelForCausalLM.from_pretrained(
325-
model_id,
326-
token=hf_token if hf_token else None,
327-
**quant_config
328-
)
302+
try:
303+
# Load the model with Hugging Face's native progress bars
304+
self.model = AutoModelForCausalLM.from_pretrained(
305+
model_id,
306+
token=hf_token if hf_token else None,
307+
**quant_config
308+
)
309+
finally:
310+
# Restore our logger's handlers
311+
root_logger.handlers = original_handlers
329312
# Reset the downloading flag
330313
try:
331314
# Access the module's global variable
@@ -1106,47 +1089,41 @@ async def load_custom_model(self, model_name: str, fallback_model: Optional[str]
11061089
# Add an empty line to separate from HuggingFace progress bars
11071090
print("")
11081091

1109-
# Set a flag to indicate we're downloading a model
1110-
try:
1111-
# Access the module's global variable
1112-
import locallab.utils.progress
1113-
locallab.utils.progress.is_downloading = True
1114-
1115-
# Ensure HuggingFace Hub's progress bars are enabled
1116-
from huggingface_hub.utils import logging as hf_logging
1117-
hf_logging.enable_progress_bars()
1092+
# Add an empty line before progress bars start
1093+
print(f"\n{Fore.CYAN}Starting custom model download - native progress bars will appear below{Style.RESET_ALL}\n")
11181094

1119-
# Configure transformers to use progress bars
1120-
import transformers
1121-
transformers.utils.logging.enable_progress_bar()
1095+
# Enable Hugging Face progress bars again to ensure they're properly configured
1096+
enable_hf_progress_bars()
11221097

1123-
# Also ensure tqdm is properly configured for nice display
1124-
import tqdm
1125-
tqdm.tqdm.monitor_interval = 0 # Disable monitor thread which can cause issues
1098+
# Use a context manager to ensure proper display of Hugging Face progress bars
1099+
with StdoutRedirector(disable_logging=True):
1100+
# Load tokenizer first
1101+
logger.info(f"Loading tokenizer for custom model {model_name}...")
11261102

1127-
# Temporarily disable our custom logger for HuggingFace logs
1128-
import logging
1129-
for logger_name in ['tqdm', 'huggingface_hub', 'transformers', 'filelock']:
1130-
logging.getLogger(logger_name).handlers = [] # Remove any handlers
1131-
logging.getLogger(logger_name).propagate = False # Don't propagate to parent loggers
1132-
except:
1133-
# Fallback if import fails
1134-
pass
1103+
# This is the critical part where we want to see nice progress bars
1104+
# We'll temporarily disable our logger's handlers to prevent interference
1105+
root_logger = logging.getLogger()
1106+
original_handlers = root_logger.handlers.copy()
1107+
root_logger.handlers = []
11351108

1136-
# Add an empty line before progress bars start
1137-
print("\n")
1138-
1139-
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
1140-
logger.info(f"Tokenizer loaded successfully")
1141-
1142-
# Load model with optimizations
1143-
logger.info(f"Loading model weights for custom model {model_name}...")
1144-
self.model = AutoModelForCausalLM.from_pretrained(
1145-
model_name,
1146-
torch_dtype=torch.float16,
1147-
device_map="auto",
1148-
quantization_config=quant_config
1149-
)
1109+
try:
1110+
# Load tokenizer with Hugging Face's native progress bars
1111+
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
1112+
logger.info(f"Tokenizer loaded successfully")
1113+
1114+
# Load model with optimizations
1115+
logger.info(f"Loading model weights for custom model {model_name}...")
1116+
1117+
# Load the model with Hugging Face's native progress bars
1118+
self.model = AutoModelForCausalLM.from_pretrained(
1119+
model_name,
1120+
torch_dtype=torch.float16,
1121+
device_map="auto",
1122+
quantization_config=quant_config
1123+
)
1124+
finally:
1125+
# Restore our logger's handlers
1126+
root_logger.handlers = original_handlers
11501127
# Reset the downloading flag
11511128
try:
11521129
# Access the module's global variable

‎locallab/utils/early_config.py

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
"""
2+
Early configuration module for LocalLab.
3+
This module is imported before any other modules to configure logging and environment variables.
4+
"""
5+
6+
import os
7+
import sys
8+
import logging
9+
import warnings
10+
11+
# Configure environment variables for Hugging Face
12+
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1" # Enable HF Transfer for better downloads
13+
os.environ["TOKENIZERS_PARALLELISM"] = "true" # Enable parallelism for tokenizers
14+
os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1" # Disable advisory warnings
15+
os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1" # Disable telemetry
16+
17+
# Configure tqdm to use the best available display method
18+
os.environ["TQDM_DISABLE"] = "0" # Ensure tqdm is not disabled
19+
os.environ["TQDM_MININTERVAL"] = "0.1" # Update progress bars more frequently
20+
21+
# Configure Hugging Face logging before importing any HF libraries
22+
def configure_hf_logging():
23+
"""
24+
Configure Hugging Face logging before any HF libraries are imported.
25+
This ensures that HF's progress bars are displayed correctly.
26+
"""
27+
# Disable all warnings
28+
warnings.filterwarnings("ignore")
29+
30+
# Configure logging for Hugging Face libraries
31+
for logger_name in ["transformers", "huggingface_hub", "accelerate", "tqdm", "filelock"]:
32+
hf_logger = logging.getLogger(logger_name)
33+
hf_logger.setLevel(logging.WARNING) # Only show warnings and errors
34+
hf_logger.propagate = False # Don't propagate to parent loggers
35+
36+
# Remove any existing handlers
37+
for handler in hf_logger.handlers[:]:
38+
hf_logger.removeHandler(handler)
39+
40+
# Add a null handler to prevent warnings about no handlers
41+
hf_logger.addHandler(logging.NullHandler())
42+
43+
# Run configuration immediately on import
44+
configure_hf_logging()
45+
46+
# Function to temporarily redirect stdout/stderr during model downloads
47+
class StdoutRedirector:
48+
"""
49+
Context manager to temporarily redirect stdout/stderr during model downloads.
50+
This ensures that tqdm progress bars are displayed correctly.
51+
"""
52+
def __init__(self, disable_logging=True):
53+
self.disable_logging = disable_logging
54+
self.original_stdout = sys.stdout
55+
self.original_stderr = sys.stderr
56+
self.original_log_levels = {}
57+
58+
def __enter__(self):
59+
# Store original log levels
60+
if self.disable_logging:
61+
for logger_name in ["transformers", "huggingface_hub", "accelerate", "tqdm", "filelock"]:
62+
logger = logging.getLogger(logger_name)
63+
self.original_log_levels[logger_name] = logger.level
64+
logger.setLevel(logging.WARNING)
65+
66+
return self
67+
68+
def __exit__(self, exc_type, exc_val, exc_tb):
69+
# Restore original log levels
70+
if self.disable_logging:
71+
for logger_name, level in self.original_log_levels.items():
72+
logging.getLogger(logger_name).setLevel(level)
73+
74+
# Function to enable Hugging Face progress bars
75+
def enable_hf_progress_bars():
76+
"""
77+
Enable Hugging Face progress bars.
78+
Call this function before downloading models.
79+
"""
80+
# Configure tqdm
81+
try:
82+
import tqdm
83+
tqdm.tqdm.monitor_interval = 0 # Disable monitor thread
84+
except ImportError:
85+
pass
86+
87+
# Configure huggingface_hub
88+
try:
89+
import huggingface_hub
90+
huggingface_hub.enable_progress_bars()
91+
except ImportError:
92+
pass
93+
94+
# Configure transformers
95+
try:
96+
import transformers
97+
transformers.utils.logging.enable_progress_bar()
98+
transformers.logging.set_verbosity_warning()
99+
except ImportError:
100+
pass

‎setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@
4747

4848
setup(
4949
name="locallab",
50-
version="0.6.3",
50+
version="0.6.4",
5151
packages=find_packages(include=["locallab", "locallab.*"]),
5252
install_requires=install_requires,
5353
extras_require={

0 commit comments

Comments
 (0)