textattack.llms package

Large Language Models

TextAttack can generate responses to prompts using LLMs, which take in a list of strings and outputs a list of responses.

We’ve provided an implementation around two common LLM patterns:

  1. HuggingFaceLLMWrapper for LLMs in HuggingFace

  2. ChatGptWrapper for OpenAI’s ChatGPT model

class textattack.llms.huggingface_llm_wrapper.HuggingFaceLLMWrapper(model, tokenizer)[source]

Bases: ModelWrapper

A wrapper around HuggingFace for LLMs.

Parameters:
  • model – A HuggingFace pretrained LLM

  • tokenizer – A HuggingFace pretrained tokenizer

class textattack.llms.chat_gpt_wrapper.ChatGptWrapper(model_name='gpt-3.5-turbo', key_environment_variable='OPENAI_API_KEY')[source]

Bases: ModelWrapper

A wrapper around OpenAI’s ChatGPT model. Note that you must provide your own API key to use this wrapper.

Parameters:
  • model_name (str) – The name of the GPT model to use. See the OpenAI documentation for a list of latest model names

  • key_environment_variable (str, ‘optional`, defaults to OPENAI_API_KEY) – The environment variable that the API key is set to