pyrit.prompt_target.HuggingFaceChatTarget

pyrit.prompt_target.HuggingFaceChatTarget#

class HuggingFaceChatTarget(*, model_id: str | None = None, model_path: str | None = None, hf_access_token: str | None = None, use_cuda: bool = False, tensor_format: str = 'pt', necessary_files: list = None, max_new_tokens: int = 20, temperature: float = 1.0, top_p: float = 1.0, skip_special_tokens: bool = True, trust_remote_code: bool = False, device_map: str | None = None, torch_dtype: torch.dtype | None = None, attn_implementation: str | None = None)[source]#

Bases: PromptChatTarget

The HuggingFaceChatTarget interacts with HuggingFace models, specifically for conducting red teaming activities. Inherits from PromptTarget to comply with the current design standards.

__init__(*, model_id: str | None = None, model_path: str | None = None, hf_access_token: str | None = None, use_cuda: bool = False, tensor_format: str = 'pt', necessary_files: list = None, max_new_tokens: int = 20, temperature: float = 1.0, top_p: float = 1.0, skip_special_tokens: bool = True, trust_remote_code: bool = False, device_map: str | None = None, torch_dtype: torch.dtype | None = None, attn_implementation: str | None = None) → None[source]#

Methods

__init__(*[, model_id, model_path, ...])

disable_cache()

Disables the class-level cache and clears the cache.

dispose_db_engine()

Dispose DuckDB database engine to release database connections and resources.

enable_cache()

Enables the class-level cache.

get_identifier()

is_model_id_valid()

Check if the HuggingFace model ID is valid.

load_model_and_tokenizer()

Loads the model and tokenizer, downloading if necessary.

send_chat_prompt_async(*, prompt, ...[, ...])

Sends a text prompt to the target without having to build the prompt request.

send_prompt_async(*, prompt_request)

Sends a normalized prompt asynchronously to the HuggingFace model.

set_system_prompt(*, system_prompt, ...[, ...])

Sets the system prompt for the prompt target.

Attributes

HUGGINGFACE_TOKEN_ENVIRONMENT_VARIABLE

supported_converters

HUGGINGFACE_TOKEN_ENVIRONMENT_VARIABLE = 'HUGGINGFACE_TOKEN'#
classmethod disable_cache()[source]#

Disables the class-level cache and clears the cache.

classmethod enable_cache()[source]#

Enables the class-level cache.

is_model_id_valid() → bool[source]#

Check if the HuggingFace model ID is valid. :return: True if valid, False otherwise.

async load_model_and_tokenizer()[source]#

Loads the model and tokenizer, downloading if necessary.

Downloads the model to the HF_MODELS_DIR folder if it does not exist, then loads it from there.

Raises:

Exception – If the model loading fails.

async send_prompt_async(*, prompt_request: PromptRequestResponse) → PromptRequestResponse[source]#

Sends a normalized prompt asynchronously to the HuggingFace model.