7
7
import json
8
8
import requests
9
9
import threading
10
+ import logging
10
11
from urllib .parse import urlparse
11
12
12
13
@@ -110,6 +111,8 @@ class SafetensorRemote:
110
111
print(data)
111
112
"""
112
113
114
+ logger = logging .getLogger ("safetensor_remote" )
115
+
113
116
BASE_DOMAIN = "https://huggingface.co"
114
117
ALIGNMENT = 8 # bytes
115
118
@@ -243,7 +246,7 @@ def get_data_by_range(cls, url: str, start: int, size: int = -1) -> bytes:
243
246
244
247
# --- Multithreading Path ---
245
248
if size >= cls .MULTITHREAD_THREDSHOLD and cls .MULTITHREAD_COUNT > 1 :
246
- print (f"Using { cls .MULTITHREAD_COUNT } threads to download range of { size / (1024 * 1024 ):.2f} MB" )
249
+ cls . logger . info (f"Using { cls .MULTITHREAD_COUNT } threads to download range of { size / (1024 * 1024 ):.2f} MB" )
247
250
num_threads = cls .MULTITHREAD_COUNT
248
251
results : list [Any ] = [None ] * num_threads # Store results or exceptions
249
252
threads : list [threading .Thread ] = []
@@ -308,9 +311,9 @@ def download_chunk(chunk_url: str, chunk_start: int, chunk_size: int, index: int
308
311
# Check if it was supposed to download anything
309
312
expected_chunk_size = base_chunk_size + (1 if i < remainder else 0 )
310
313
if expected_chunk_size > 0 :
311
- raise RuntimeError (f"Thread { i } finished without providing data or exception for a non-zero chunk." )
314
+ raise RuntimeError (f"Thread { i } finished without providing data or exception for a non-zero chunk." )
312
315
else :
313
- final_data_parts .append (b"" ) # Append empty bytes for zero-size chunk
316
+ final_data_parts .append (b"" ) # Append empty bytes for zero-size chunk
314
317
else :
315
318
final_data_parts .append (result )
316
319
@@ -319,7 +322,7 @@ def download_chunk(chunk_url: str, chunk_start: int, chunk_size: int, index: int
319
322
320
323
# Final validation: Does the combined size match the requested size?
321
324
if len (final_data ) != size :
322
- raise IOError (f"Final assembled data size mismatch. Expected { size } , got { len (final_data )} . URL: { url } , Range: { start } -{ start + size - 1 } " )
325
+ raise IOError (f"Final assembled data size mismatch. Expected { size } , got { len (final_data )} . URL: { url } , Range: { start } -{ start + size - 1 } " )
323
326
324
327
return final_data
325
328
0 commit comments