diff --git a/doc/code/targets/11_message_normalizer.ipynb b/doc/code/targets/11_message_normalizer.ipynb index 7f896e960e..a355fd3569 100644 --- a/doc/code/targets/11_message_normalizer.ipynb +++ b/doc/code/targets/11_message_normalizer.ipynb @@ -320,13 +320,6 @@ "No HuggingFace token provided. Gated models may fail to load without authentication.\n" ] }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n" - ] - }, { "name": "stdout", "output_type": "stream", @@ -535,7 +528,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.12" + "version": "3.14.4" } }, "nbformat": 4, diff --git a/doc/code/targets/use_huggingface_chat_target.ipynb b/doc/code/targets/use_huggingface_chat_target.ipynb index 08e329dc3b..02d34c4365 100644 --- a/doc/code/targets/use_huggingface_chat_target.ipynb +++ b/doc/code/targets/use_huggingface_chat_target.ipynb @@ -41,6 +41,14 @@ "id": "1", "metadata": {}, "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "./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", + " space_regex = \"[\\*_~|`\\-\\.]*\" if include_character_padding else ''\n" + ] + }, { "name": "stdout", "output_type": "stream", @@ -54,14 +62,48 @@ "name": "stdout", "output_type": "stream", "text": [ - "No new upgrade operations detected.\n", + "[pyrit:alembic] No new upgrade operations detected.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ "Running model: Qwen/Qwen2-0.5B-Instruct\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "f637f0a22e814d5390fcfab383ecff0b", + "model_id": "2cb1ed2b5d6c4fa98456d1c217564010", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Downloading (incomplete total...): 0.00B [00:00, ?B/s]" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "76523dac8fe04b87921ad444853385c3", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Fetching 10 files: 0%| | 0/10 [00:00 list[str]: - """ - Fetch available files for a model from the Hugging Face repository. - - Returns: - List of available file names. - - Raises: - ValueError: If no files are found for the model. - """ - api = HfApi() - try: - model_info = api.model_info(model_id, token=token) - available_files = [file.rfilename for file in (model_info.siblings or [])] - - # Perform simple validation: raise a ValueError if no files are available - if not len(available_files): - raise ValueError(f"No available files found for the model: {model_id}") - - return available_files - except Exception as e: - logger.info(f"Error fetching model files for {model_id}: {e}") - return [] +from huggingface_hub import snapshot_download async def download_specific_files_async( model_id: str, file_patterns: list[str] | None, token: str, cache_dir: Path ) -> None: """ - Download specific files from a Hugging Face model repository. + Download a Hugging Face model snapshot without blocking the event loop. + If file_patterns is None, downloads all files. """ cache_dir.mkdir(parents=True, exist_ok=True) - - available_files = get_available_files(model_id, token) - # If no file patterns are provided, download all available files - if file_patterns is None: - files_to_download = available_files - logger.info(f"Downloading all files for model {model_id}.") - else: - # Filter files based on the patterns provided - files_to_download = [file for file in available_files if any(pattern in file for pattern in file_patterns)] - if not files_to_download: - logger.info(f"No files matched the patterns provided for model {model_id}.") - return - - # Generate download URLs directly - base_url = f"https://huggingface.co/{model_id}/resolve/main/" - urls = [base_url + file for file in files_to_download] - - # Download the files - await download_files_async(urls, token, cache_dir) - - -async def download_chunk_async( - url: str, headers: dict[str, str], start: int, end: int, client: httpx.AsyncClient -) -> bytes: - """ - Download a chunk of the file with a specified byte range. - - Returns: - The content of the downloaded chunk. - """ - range_header = {"Range": f"bytes={start}-{end}", **headers} - response = await client.get(url, headers=range_header) - response.raise_for_status() - return response.content - - -async def download_file_async(url: str, token: str, download_dir: Path, num_splits: int) -> None: - """Download a file in multiple segments (splits) using byte-range requests.""" - headers = {"Authorization": f"Bearer {token}"} - async with httpx.AsyncClient(follow_redirects=True) as client: - # Get the file size to determine chunk size - response = await client.head(url, headers=headers) - response.raise_for_status() - file_size = int(response.headers["Content-Length"]) - chunk_size = file_size // num_splits - - # Prepare tasks for each chunk - tasks = [] - file_name = url.split("/")[-1] - file_path = Path(download_dir, file_name) - - for i in range(num_splits): - start = i * chunk_size - end = start + chunk_size - 1 if i < num_splits - 1 else file_size - 1 - tasks.append(download_chunk_async(url, headers, start, end, client)) - - # Download all chunks concurrently - chunks = await asyncio.gather(*tasks) - - # Write chunks to the file in order - async with aiofiles.open(file_path, "wb") as f: - for chunk in chunks: - await f.write(chunk) - logger.info(f"Downloaded {file_name} to {file_path}") - - -async def download_files_async( - urls: list[str], token: str, download_dir: Path, num_splits: int = 3, parallel_downloads: int = 4 -) -> None: - """Download multiple files with parallel downloads and segmented downloading.""" - # Limit the number of parallel downloads - semaphore = asyncio.Semaphore(parallel_downloads) - - async def download_with_limit_async(url: str) -> None: - async with semaphore: - await download_file_async(url, token, download_dir, num_splits) - - # Run downloads concurrently, but limit to parallel_downloads at a time - await asyncio.gather(*(download_with_limit_async(url) for url in urls)) + await asyncio.to_thread( + snapshot_download, + repo_id=model_id, + allow_patterns=file_patterns, + token=token, + local_dir=cache_dir, + ) diff --git a/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py b/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py index d333e234d8..ad98a94969 100644 --- a/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py +++ b/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py @@ -295,20 +295,12 @@ async def load_model_and_tokenizer_async(self) -> None: Path(cache_dir), ) - # Load the tokenizer and model from the specified directory + # Load the tokenizer and model from the downloaded local snapshot. logger.info(f"Loading model {self.model_id} from cache path: {cache_dir}...") - self.tokenizer = AutoTokenizer.from_pretrained( - self.model_id or "", cache_dir=cache_dir, trust_remote_code=self.trust_remote_code - ) - self.model = AutoModelForCausalLM.from_pretrained( - self.model_id or "", - cache_dir=cache_dir, - trust_remote_code=self.trust_remote_code, - **optional_model_kwargs, - ) + self._load_from_path(str(cache_dir), **optional_model_kwargs) # Move the model to the correct device - self.model = cast("Any", self.model).to(self.device) + self.model = self.model.to(self.device) # Debug prints to check types logger.info(f"Model loaded: {type(self.model)}") diff --git a/tests/unit/common/test_hf_model_downloads.py b/tests/unit/common/test_hf_model_downloads.py index ce5d733d07..54abd49ad8 100644 --- a/tests/unit/common/test_hf_model_downloads.py +++ b/tests/unit/common/test_hf_model_downloads.py @@ -3,17 +3,11 @@ import os from pathlib import Path -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import patch import pytest -# Import functions to test from local application files -from pyrit.common.download_hf_model import ( - download_chunk_async, - download_file_async, - download_files_async, - download_specific_files_async, -) +from pyrit.common.download_hf_model import download_specific_files_async # Define constants for testing MODEL_ID = "microsoft/Phi-3-mini-4k-instruct" @@ -31,7 +25,6 @@ @pytest.fixture(scope="module") def setup_environment(): """Fixture to set up the environment for Hugging Face downloads.""" - # Check for Hugging Face token with patch.dict(os.environ, {"HUGGINGFACE_TOKEN": "mocked_token"}): token = os.getenv("HUGGINGFACE_TOKEN") yield token @@ -40,65 +33,17 @@ def setup_environment(): async def test_download_specific_files_async(setup_environment): """Test downloading specific files""" token = setup_environment # Get the token from the fixture - - with patch("os.makedirs"), patch("pyrit.common.download_hf_model.download_files_async"): - await download_specific_files_async(MODEL_ID, FILE_PATTERNS, token, Path("")) - - -async def test_download_files_async_dispatches_one_call_per_url(): - """Exercise the nested download_with_limit_async helper plus asyncio.gather.""" - seen_urls: list[str] = [] - - async def fake_file_async(url, token, download_dir, num_splits): - seen_urls.append(url) - - with patch("pyrit.common.download_hf_model.download_file_async", new=fake_file_async): - urls = ["https://example/a", "https://example/b", "https://example/c"] - await download_files_async(urls, "token", Path("/tmp"), num_splits=2, parallel_downloads=2) - - assert sorted(seen_urls) == sorted(urls) - - -async def test_download_file_async_schedules_one_chunk_per_split(tmp_path): - """Exercise the chunk-task assembly inside download_file_async.""" - num_splits = 3 - file_size = 30 - chunk_bytes = b"abcdefghij" - - head_response = MagicMock() - head_response.raise_for_status = MagicMock() - head_response.headers = {"Content-Length": str(file_size)} - - client_instance = MagicMock() - client_instance.head = AsyncMock(return_value=head_response) - - class _ClientCM: - async def __aenter__(self): - return client_instance - - async def __aexit__(self, exc_type, exc, tb): - return False + cache_dir = Path("model-cache") with ( - patch("pyrit.common.download_hf_model.httpx.AsyncClient", return_value=_ClientCM()), - patch("pyrit.common.download_hf_model.download_chunk_async", new_callable=AsyncMock) as mock_chunk_async, + patch("pathlib.Path.mkdir"), + patch("pyrit.common.download_hf_model.snapshot_download") as snapshot_download_mock, ): - mock_chunk_async.return_value = chunk_bytes - await download_file_async("https://example/myfile.bin", "token", tmp_path, num_splits) - - assert mock_chunk_async.await_count == num_splits - assert (tmp_path / "myfile.bin").read_bytes() == chunk_bytes * num_splits - - -async def test_download_chunk_async_returns_response_content(): - """Sanity-check the real download_chunk_async implementation.""" - response = MagicMock() - response.raise_for_status = MagicMock() - response.content = b"chunk-payload" - client = MagicMock() - client.get = AsyncMock(return_value=response) - - result = await download_chunk_async("https://example/file", {"Authorization": "Bearer t"}, 0, 9, client) - - assert result == b"chunk-payload" - client.get.assert_awaited_once() + await download_specific_files_async(MODEL_ID, FILE_PATTERNS, token, cache_dir) + + snapshot_download_mock.assert_called_once_with( + repo_id=MODEL_ID, + allow_patterns=FILE_PATTERNS, + token=token, + local_dir=cache_dir, + ) diff --git a/tests/unit/prompt_target/target/test_huggingface_chat_target.py b/tests/unit/prompt_target/target/test_huggingface_chat_target.py index 3d302c91f3..12177d7c44 100644 --- a/tests/unit/prompt_target/target/test_huggingface_chat_target.py +++ b/tests/unit/prompt_target/target/test_huggingface_chat_target.py @@ -3,6 +3,8 @@ import json from asyncio import Task +from collections.abc import Coroutine +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -92,9 +94,12 @@ async def _await(): @pytest.fixture(autouse=True) def mock_create_task(): + def _close_coroutine(coroutine: Coroutine[Any, Any, None]) -> AwaitableTask: + coroutine.close() + return AwaitableTask(spec=Task) + with patch("asyncio.create_task") as mock_task: - # Return an AwaitableTask that can be awaited - mock_task.return_value = AwaitableTask(spec=Task) + mock_task.side_effect = _close_coroutine yield mock_task