Skip to content

Commit 27d20ee

Browse files
romanlutzCopilot
romanlutz
and
Copilot
authored
FIX: Load HuggingFace models without blocking the event loop (microsoft#2211)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent 7c7a373 commit 27d20ee

6 files changed

Lines changed: 95 additions & 213 deletions

File tree

‎doc/code/targets/11_message_normalizer.ipynb‎

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -320,13 +320,6 @@
320320
"No HuggingFace token provided. Gated models may fail to load without authentication.\n"
321321
]
322322
},
323-
{
324-
"name": "stderr",
325-
"output_type": "stream",
326-
"text": [
327-
"Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n"
328-
]
329-
},
330323
{
331324
"name": "stdout",
332325
"output_type": "stream",
@@ -535,7 +528,7 @@
535528
"name": "python",
536529
"nbconvert_exporter": "python",
537530
"pygments_lexer": "ipython3",
538-
"version": "3.12.12"
531+
"version": "3.14.4"
539532
}
540533
},
541534
"nbformat": 4,

‎doc/code/targets/use_huggingface_chat_target.ipynb‎

Lines changed: 61 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,14 @@
4141
"id": "1",
4242
"metadata": {},
4343
"outputs": [
44+
{
45+
"name": "stderr",
46+
"output_type": "stream",
47+
"text": [
48+
"./git/copilot-worktrees/PyRIT/romanlutz-cautious-meme/.venv/Lib/site-packages/confusables/__init__.py:46: SyntaxWarning: \"\\*\" is an invalid escape sequence. Such sequences will not work in the future. Did you mean \"\\\\*\"? A raw string is also an option.\n",
49+
" space_regex = \"[\\*_~|`\\-\\.]*\" if include_character_padding else ''\n"
50+
]
51+
},
4452
{
4553
"name": "stdout",
4654
"output_type": "stream",
@@ -54,14 +62,48 @@
5462
"name": "stdout",
5563
"output_type": "stream",
5664
"text": [
57-
"No new upgrade operations detected.\n",
65+
"[pyrit:alembic] No new upgrade operations detected.\n"
66+
]
67+
},
68+
{
69+
"name": "stdout",
70+
"output_type": "stream",
71+
"text": [
5872
"Running model: Qwen/Qwen2-0.5B-Instruct\n"
5973
]
6074
},
6175
{
6276
"data": {
6377
"application/vnd.jupyter.widget-view+json": {
64-
"model_id": "f637f0a22e814d5390fcfab383ecff0b",
78+
"model_id": "2cb1ed2b5d6c4fa98456d1c217564010",
79+
"version_major": 2,
80+
"version_minor": 0
81+
},
82+
"text/plain": [
83+
"Downloading (incomplete total...): 0.00B [00:00, ?B/s]"
84+
]
85+
},
86+
"metadata": {},
87+
"output_type": "display_data"
88+
},
89+
{
90+
"data": {
91+
"application/vnd.jupyter.widget-view+json": {
92+
"model_id": "76523dac8fe04b87921ad444853385c3",
93+
"version_major": 2,
94+
"version_minor": 0
95+
},
96+
"text/plain": [
97+
"Fetching 10 files: 0%| | 0/10 [00:00<?, ?it/s]"
98+
]
99+
},
100+
"metadata": {},
101+
"output_type": "display_data"
102+
},
103+
{
104+
"data": {
105+
"application/vnd.jupyter.widget-view+json": {
106+
"model_id": "650a316e795b4e84abdc46b06bc79cc1",
65107
"version_major": 2,
66108
"version_minor": 0
67109
},
@@ -76,7 +118,7 @@
76118
"name": "stdout",
77119
"output_type": "stream",
78120
"text": [
79-
"Average response time for Qwen/Qwen2-0.5B-Instruct: 32.08 seconds\n",
121+
"Average response time for Qwen/Qwen2-0.5B-Instruct: 7.43 seconds\n",
80122
"\n",
81123
"\n",
82124
"\u001b[33m════════════════════════════════════════════════════════════════════════════════════════════════════\u001b[0m\n",
@@ -88,11 +130,11 @@
88130
"\u001b[1m 📋 Basic Information\u001b[0m\n",
89131
"\u001b[36m • Objective: What is 3*3? Give me the solution.\u001b[0m\n",
90132
"\u001b[36m • Attack Type: PromptSendingAttack\u001b[0m\n",
91-
"\u001b[36m • Conversation ID: 434c7b92-3756-4f64-8c9d-f61a181e1ed8\u001b[0m\n",
133+
"\u001b[36m • Conversation ID: 3e01b421-e41b-4c47-a51b-ea192f7cdb24\u001b[0m\n",
92134
"\n",
93135
"\u001b[1m ⚡ Execution Metrics\u001b[0m\n",
94136
"\u001b[32m • Turns Executed: 1\u001b[0m\n",
95-
"\u001b[32m • Execution Time: 1m 2s\u001b[0m\n",
137+
"\u001b[32m • Execution Time: 12.20s\u001b[0m\n",
96138
"\n",
97139
"\u001b[1m 🎯 Outcome\u001b[0m\n",
98140
"\u001b[33m • Status: ❓ UNDETERMINED\u001b[0m\n",
@@ -109,12 +151,14 @@
109151
"\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
110152
"\u001b[1m\u001b[33m🔸 ASSISTANT\u001b[0m\n",
111153
"\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
112-
"\u001b[33m The solution, after expanding \\(3 \\cdot 3\\), is that it equals 9.\u001b[0m\n",
154+
"\u001b[33m The answer to \"What is 3 * 3?\" is 9.\u001b[0m\n",
155+
"\u001b[33m \u001b[0m\n",
156+
"\u001b[33m This is because the product of two numbers, where the first number (\u001b[0m\n",
113157
"\n",
114158
"\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
115159
"\n",
116160
"\u001b[2m\u001b[37m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
117-
"\u001b[2m\u001b[37m Report generated at: 2026-05-21 23:13:40 UTC \u001b[0m\n",
161+
"\u001b[2m\u001b[37m Report generated at: 2026-07-16 05:18:39 UTC \u001b[0m\n",
118162
"\n",
119163
"\u001b[33m════════════════════════════════════════════════════════════════════════════════════════════════════\u001b[0m\n",
120164
"\u001b[1m\u001b[33m ❓ ATTACK RESULT: UNDETERMINED ❓ \u001b[0m\n",
@@ -125,11 +169,11 @@
125169
"\u001b[1m 📋 Basic Information\u001b[0m\n",
126170
"\u001b[36m • Objective: What is 4*4? Give me the solution.\u001b[0m\n",
127171
"\u001b[36m • Attack Type: PromptSendingAttack\u001b[0m\n",
128-
"\u001b[36m • Conversation ID: 3b782d88-adb6-47b1-bb3f-699ec338d66b\u001b[0m\n",
172+
"\u001b[36m • Conversation ID: 76b845b6-0e72-4a80-a793-09652c4d7405\u001b[0m\n",
129173
"\n",
130174
"\u001b[1m ⚡ Execution Metrics\u001b[0m\n",
131175
"\u001b[32m • Turns Executed: 1\u001b[0m\n",
132-
"\u001b[32m • Execution Time: 1.40s\u001b[0m\n",
176+
"\u001b[32m • Execution Time: 2.63s\u001b[0m\n",
133177
"\n",
134178
"\u001b[1m 🎯 Outcome\u001b[0m\n",
135179
"\u001b[33m • Status: ❓ UNDETERMINED\u001b[0m\n",
@@ -146,13 +190,17 @@
146190
"\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
147191
"\u001b[1m\u001b[33m🔸 ASSISTANT\u001b[0m\n",
148192
"\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
149-
"\u001b[33m The answer to \"4 * 4\" is 16.\u001b[0m\n",
193+
"\u001b[33m The result of multiplying 4 by itself four times is:\u001b[0m\n",
194+
"\u001b[33m 256.\u001b[0m\n",
195+
"\u001b[33m \u001b[0m\n",
196+
"\u001b[33m Here's why:\u001b[0m\n",
197+
"\u001b[33m First, we multiply 4 and\u001b[0m\n",
150198
"\n",
151199
"\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
152200
"\n",
153201
"\u001b[2m\u001b[37m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
154-
"\u001b[2m\u001b[37m Report generated at: 2026-05-21 23:13:40 UTC \u001b[0m\n",
155-
"Qwen/Qwen2-0.5B-Instruct: 32.08 seconds\n"
202+
"\u001b[2m\u001b[37m Report generated at: 2026-07-16 05:18:39 UTC \u001b[0m\n",
203+
"Qwen/Qwen2-0.5B-Instruct: 7.43 seconds\n"
156204
]
157205
}
158206
],
@@ -236,7 +284,7 @@
236284
"name": "python",
237285
"nbconvert_exporter": "python",
238286
"pygments_lexer": "ipython3",
239-
"version": "3.12.12"
287+
"version": "3.14.4"
240288
}
241289
},
242290
"nbformat": 4,

‎pyrit/common/download_hf_model.py‎

Lines changed: 10 additions & 111 deletions
Original file line numberDiff line numberDiff line change
@@ -2,125 +2,24 @@
22
# Licensed under the MIT license.
33

44
import asyncio
5-
import logging
65
from pathlib import Path
76

8-
import aiofiles
9-
import httpx
10-
from huggingface_hub import HfApi
11-
12-
logger = logging.getLogger(__name__)
13-
14-
15-
def get_available_files(model_id: str, token: str) -> list[str]:
16-
"""
17-
Fetch available files for a model from the Hugging Face repository.
18-
19-
Returns:
20-
List of available file names.
21-
22-
Raises:
23-
ValueError: If no files are found for the model.
24-
"""
25-
api = HfApi()
26-
try:
27-
model_info = api.model_info(model_id, token=token)
28-
available_files = [file.rfilename for file in (model_info.siblings or [])]
29-
30-
# Perform simple validation: raise a ValueError if no files are available
31-
if not len(available_files):
32-
raise ValueError(f"No available files found for the model: {model_id}")
33-
34-
return available_files
35-
except Exception as e:
36-
logger.info(f"Error fetching model files for {model_id}: {e}")
37-
return []
7+
from huggingface_hub import snapshot_download
388

399

4010
async def download_specific_files_async(
4111
model_id: str, file_patterns: list[str] | None, token: str, cache_dir: Path
4212
) -> None:
4313
"""
44-
Download specific files from a Hugging Face model repository.
14+
Download a Hugging Face model snapshot without blocking the event loop.
15+
4516
If file_patterns is None, downloads all files.
4617
"""
4718
cache_dir.mkdir(parents=True, exist_ok=True)
48-
49-
available_files = get_available_files(model_id, token)
50-
# If no file patterns are provided, download all available files
51-
if file_patterns is None:
52-
files_to_download = available_files
53-
logger.info(f"Downloading all files for model {model_id}.")
54-
else:
55-
# Filter files based on the patterns provided
56-
files_to_download = [file for file in available_files if any(pattern in file for pattern in file_patterns)]
57-
if not files_to_download:
58-
logger.info(f"No files matched the patterns provided for model {model_id}.")
59-
return
60-
61-
# Generate download URLs directly
62-
base_url = f"https://huggingface.co/{model_id}/resolve/main/"
63-
urls = [base_url + file for file in files_to_download]
64-
65-
# Download the files
66-
await download_files_async(urls, token, cache_dir)
67-
68-
69-
async def download_chunk_async(
70-
url: str, headers: dict[str, str], start: int, end: int, client: httpx.AsyncClient
71-
) -> bytes:
72-
"""
73-
Download a chunk of the file with a specified byte range.
74-
75-
Returns:
76-
The content of the downloaded chunk.
77-
"""
78-
range_header = {"Range": f"bytes={start}-{end}", **headers}
79-
response = await client.get(url, headers=range_header)
80-
response.raise_for_status()
81-
return response.content
82-
83-
84-
async def download_file_async(url: str, token: str, download_dir: Path, num_splits: int) -> None:
85-
"""Download a file in multiple segments (splits) using byte-range requests."""
86-
headers = {"Authorization": f"Bearer {token}"}
87-
async with httpx.AsyncClient(follow_redirects=True) as client:
88-
# Get the file size to determine chunk size
89-
response = await client.head(url, headers=headers)
90-
response.raise_for_status()
91-
file_size = int(response.headers["Content-Length"])
92-
chunk_size = file_size // num_splits
93-
94-
# Prepare tasks for each chunk
95-
tasks = []
96-
file_name = url.split("/")[-1]
97-
file_path = Path(download_dir, file_name)
98-
99-
for i in range(num_splits):
100-
start = i * chunk_size
101-
end = start + chunk_size - 1 if i < num_splits - 1 else file_size - 1
102-
tasks.append(download_chunk_async(url, headers, start, end, client))
103-
104-
# Download all chunks concurrently
105-
chunks = await asyncio.gather(*tasks)
106-
107-
# Write chunks to the file in order
108-
async with aiofiles.open(file_path, "wb") as f:
109-
for chunk in chunks:
110-
await f.write(chunk)
111-
logger.info(f"Downloaded {file_name} to {file_path}")
112-
113-
114-
async def download_files_async(
115-
urls: list[str], token: str, download_dir: Path, num_splits: int = 3, parallel_downloads: int = 4
116-
) -> None:
117-
"""Download multiple files with parallel downloads and segmented downloading."""
118-
# Limit the number of parallel downloads
119-
semaphore = asyncio.Semaphore(parallel_downloads)
120-
121-
async def download_with_limit_async(url: str) -> None:
122-
async with semaphore:
123-
await download_file_async(url, token, download_dir, num_splits)
124-
125-
# Run downloads concurrently, but limit to parallel_downloads at a time
126-
await asyncio.gather(*(download_with_limit_async(url) for url in urls))
19+
await asyncio.to_thread(
20+
snapshot_download,
21+
repo_id=model_id,
22+
allow_patterns=file_patterns,
23+
token=token,
24+
local_dir=cache_dir,
25+
)

‎pyrit/prompt_target/hugging_face/hugging_face_chat_target.py‎

Lines changed: 3 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -295,20 +295,12 @@ async def load_model_and_tokenizer_async(self) -> None:
295295
Path(cache_dir),
296296
)
297297

298-
# Load the tokenizer and model from the specified directory
298+
# Load the tokenizer and model from the downloaded local snapshot.
299299
logger.info(f"Loading model {self.model_id} from cache path: {cache_dir}...")
300-
self.tokenizer = AutoTokenizer.from_pretrained(
301-
self.model_id or "", cache_dir=cache_dir, trust_remote_code=self.trust_remote_code
302-
)
303-
self.model = AutoModelForCausalLM.from_pretrained(
304-
self.model_id or "",
305-
cache_dir=cache_dir,
306-
trust_remote_code=self.trust_remote_code,
307-
**optional_model_kwargs,
308-
)
300+
self._load_from_path(str(cache_dir), **optional_model_kwargs)
309301

310302
# Move the model to the correct device
311-
self.model = cast("Any", self.model).to(self.device)
303+
self.model = self.model.to(self.device)
312304

313305
# Debug prints to check types
314306
logger.info(f"Model loaded: {type(self.model)}")

0 commit comments

Comments
 (0)