Add files via upload

easyjailbreak/models/__init__.py

@@ -0,0 +1,7 @@
+from .model_base import ModelBase, WhiteBoxModelBase, BlackBoxModelBase
+from .huggingface_model import HuggingfaceModel, from_pretrained
+from .openai_model import OpenaiModel
+from .anthropic_model import AnthropicModel
+from .wenxinyiyan_model import WenxinyiyanModel
+
+__all__ = ['ModelBase', 'WhiteBoxModelBase', 'BlackBoxModelBase', 'HuggingfaceModel', 'from_pretrained', 'OpenaiModel', 'WenxinyiyanModel', 'AnthropicModel']

easyjailbreak/models/anthropic_model.py

@@ -0,0 +1,57 @@
+import logging
+import warnings
+import anthropic
+from .model_base import BlackBoxModelBase
+from fastchat.conversation import get_conv_template
+from httpx import URL
+
+
+class AnthropicModel(BlackBoxModelBase):
+    def __init__(self, model_name: str, api_keys: str, base_url: str, generation_config=None):
+        """
+        Initializes the OpenAI model with necessary parameters.
+        :param str model_name: The name of the model to use.
+        :param str api_keys: API keys for accessing the OpenAI service.
+        :param str template_name: The name of the conversation template, defaults to 'chatgpt'.
+        :param dict generation_config: Configuration settings for generation, defaults to an empty dictionary.
+        :param str|URL base_url: The base URL for the OpenAI API, defaults to None.
+        """
+        self.client = anthropic.Anthropic(api_key=api_keys, base_url=base_url)
+        self.model_name = model_name
+        self.conversation = get_conv_template('claude-3-5-sonnet-20240620')
+        self.generation_config = generation_config if generation_config is not None else {}
+
+        self.base_url = base_url
+
+    def set_system_message(self, system_message: str):
+        """
+        Sets a system message for the conversation.
+        :param str system_message: The system message to set.
+        """
+        self.conversation.system_message = system_message
+
+    def generate(self, messages, clear_old_history=True, **kwargs):
+        """
+        Generates a response based on messages that include conversation history.
+        :param list[str]|str messages: A list of messages or a single message string.
+                                       User and assistant messages should alternate.
+        :param bool clear_old_history: If True, clears the old conversation history before adding new messages.
+        :return str: The response generated by the OpenAI model based on the conversation history.
+        """
+        if clear_old_history:
+            self.conversation.messages = []
+        if isinstance(messages, str):
+            messages = [messages]
+        for index, message in enumerate(messages):
+            self.conversation.append_message(self.conversation.roles[index % 2], message)
+        messages = self.conversation.to_openai_api_messages()
+        system_message = messages[0]['content']
+        messages = messages[1:]
+        response = self.client.messages.create(
+            model=self.model_name,
+            system=system_message,
+            messages=messages,
+            **kwargs,
+            **self.generation_config
+        )
+        return response.content[0].text

easyjailbreak/models/gemini_model.py

@@ -0,0 +1,53 @@
+import os
+import logging
+import warnings
+import google.generativeai as genai
+from .model_base import BlackBoxModelBase
+from fastchat.conversation import get_conv_template
+from httpx import URL
+
+
+class GenaiModel(BlackBoxModelBase):
+    def __init__(self, model_name: str, api_keys: str, generation_config=None):
+        """
+        Initializes the OpenAI model with necessary parameters.
+        :param str model_name: The name of the model to use.
+        :param str api_keys: API keys for accessing the OpenAI service.
+        :param str template_name: The name of the conversation template, defaults to 'chatgpt'.
+        :param dict generation_config: Configuration settings for generation, defaults to an empty dictionary.
+        :param str|URL base_url: The base URL for the OpenAI API, defaults to None.
+        """
+        genai.configure(api_key=api_keys)
+        self.model_name = model_name
+        self.model = genai.GenerativeModel(self.model_name)
+
+        self.conversation = get_conv_template('gemini')
+        self.generation_config = generation_config if generation_config is not None else {}
+
+    def set_system_message(self, system_message: str):
+        """
+        Sets a system message for the conversation.
+        :param str system_message: The system message to set.
+        """
+        self.conversation.system_message = system_message
+
+    def generate(self, messages, clear_old_history=True, **kwargs):
+        """
+        Generates a response based on messages that include conversation history.
+        :param list[str]|str messages: A list of messages or a single message string.
+                                       User and assistant messages should alternate.
+        :param bool clear_old_history: If True, clears the old conversation history before adding new messages.
+        :return str: The response generated by the OpenAI model based on the conversation history.
+        """
+        if clear_old_history:
+            self.conversation.messages = []
+        if isinstance(messages, str):
+            messages = [messages]
+        for index, message in enumerate(messages):
+            self.conversation.append_message(self.conversation.roles[index % 2], message)
+        messages = self.conversation.to_gemini_api_messages()
+        print(messages)
+        chat = self.model.start_chat()
+        response = chat.send_message("hello")
+        print(response)
+        return response.text

easyjailbreak/models/huggingface_model.py

@@ -0,0 +1,295 @@
+"""
+This file contains a wrapper for Huggingface models, implementing various methods used in downstream tasks.
+It includes the HuggingfaceModel class that extends the functionality of the WhiteBoxModelBase class.
+"""
+
+import sys
+from .model_base import WhiteBoxModelBase
+import warnings
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import functools
+import torch
+from fastchat.conversation import get_conv_template
+from typing import Optional, Dict, List, Any
+import logging
+
+class HuggingfaceModel(WhiteBoxModelBase):
+    """
+    HuggingfaceModel is a wrapper for Huggingface's transformers models.
+    It extends the WhiteBoxModelBase class and provides additional functionality specifically
+    for handling conversation generation tasks with various models.
+    This class supports custom conversation templates and formatting,
+    and offers configurable options for generation.
+    """
+
+    def __init__(
+        self,
+        model: Any,
+        tokenizer: Any,
+        model_name: str,
+        generation_config: Optional[Dict[str, Any]] = None
+    ):
+        """
+        Initializes the HuggingfaceModel with a specified model, tokenizer, and generation configuration.
+
+        :param Any model: A huggingface model.
+        :param Any tokenizer: A huggingface tokenizer.
+        :param str model_name: The name of the model being used. Refer to
+            `FastChat conversation.py <https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py>`_
+            for possible options and templates.
+        :param Optional[Dict[str, Any]] generation_config: A dictionary containing configuration settings for text generation.
+            If None, a default configuration is used.
+        """
+        super().__init__(model, tokenizer)
+        self.model_name = model_name
+        try:
+            self.conversation = get_conv_template(model_name)
+        except KeyError:
+            logging.error(f'Invalid model_name: {model_name}. Refer to '
+                          'https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py '
+                          'for possible options and templates.')
+            raise  # Continue raising the KeyError
+
+        if model_name == 'llama-2':
+            self.conversation.sep2 = self.conversation.sep2.strip()
+
+        if model_name == 'zero_shot':
+            self.conversation.roles = tuple(['### ' + r for r in self.conversation.template.roles])
+            self.conversation.sep = '\n'
+
+        self.format_str = self.create_format_str()
+
+        if generation_config is None:
+            generation_config = {}
+        self.generation_config = generation_config
+
+    def set_system_message(self, system_message: str):
+        r"""
+        Sets a system message to be used in the conversation.
+
+        :param str system_message: The system message to be set for the conversation.
+        """
+        # TODO check llama2 add system prompt
+        self.conversation.system_message = system_message
+
+    def create_format_str(self):
+        self.conversation.messages = []
+        self.conversation.append_message(self.conversation.roles[0], "{prompt}")
+        self.conversation.append_message(self.conversation.roles[1], "{response}")
+        format_str = self.conversation.get_prompt()
+        self.conversation.messages = []  # clear history
+        return format_str
+
+    def create_conversation_prompt(self, messages, clear_old_history=True):
+        r"""
+        Constructs a conversation prompt that includes the conversation history.
+
+        :param list[str] messages: A list of messages that form the conversation history.
+                                   Messages from the user and the assistant should alternate.
+        :param bool clear_old_history: If True, clears the previous conversation history before adding new messages.
+        :return: A string representing the conversation prompt including the history.
+        """
+        if clear_old_history:
+            self.conversation.messages = []
+        if isinstance(messages, str):
+            messages = [messages]
+        for index, message in enumerate(messages):
+            self.conversation.append_message(self.conversation.roles[index % 2], message)
+        self.conversation.append_message(self.conversation.roles[-1], None)
+        return self.conversation.get_prompt()
+
+    def clear_conversation(self):
+        r"""
+        Clears the current conversation history.
+        """
+        self.conversation.messages = []
+
+    def generate(self, messages, input_field_name='input_ids', clear_old_history=True, **kwargs):
+        r"""
+        Generates a response for the given messages within a single conversation.
+
+        :param list[str]|str messages: The text input by the user. Can be a list of messages or a single message.
+        :param str input_field_name: The parameter name for the input message in the model's generation function.
+        :param bool clear_old_history: If True, clears the conversation history before generating a response.
+        :param dict kwargs: Optional parameters for the model's generation function, such as 'temperature' and 'top_p'.
+        :return: A string representing the pure response from the model, containing only the text of the response.
+        """
+        if isinstance(messages, str):
+            messages = [messages]
+        prompt = self.create_conversation_prompt(messages, clear_old_history=clear_old_history)
+
+        input_ids = self.tokenizer(prompt,
+                                   return_tensors='pt',
+                                   add_special_tokens=False).input_ids.to(self.model.device.index)
+        input_length = len(input_ids[0])
+
+        kwargs.update({input_field_name: input_ids})
+        output_ids = self.model.generate(**kwargs, **self.generation_config)
+        output = self.tokenizer.decode(output_ids[0][input_length:], skip_special_tokens=True)
+
+        return output
+
+    def batch_generate(self, conversations, **kwargs)-> List[str]:
+        r"""
+        Generates responses for a batch of conversations.
+
+        :param list[list[str]]|list[str] conversations: A list of conversations. Each conversation can be a list of messages
+                                                         or a single message string. If a single string is provided, a warning
+                                                         is issued, and the string is treated as a single-message conversation.
+        :param dict kwargs: Optional parameters for the model's generation function.
+        :return: A list of responses, each corresponding to one of the input conversations.
+        """
+        prompt_list = []
+        for conversation in conversations:
+            if isinstance(conversation, str):
+                warnings.warn('If you want the model to generate batches based on several conversations, '
+                              'please construct a list[str] for each conversation, or they will be divided into individual sentences. '
+                              'Switch input type of batch_generate() to list[list[str]] to avoid this warning.')
+                conversation = [conversation]
+            prompt_list.append(self.create_conversation_prompt(conversation))
+        input_ids = self.tokenizer(prompt_list,
+                                   return_tensors='pt',
+                                   padding=True,
+                                   add_special_tokens=False)
+        input_ids = {k: v.to(self.model.device.index) for k, v in input_ids.items()}
+        kwargs.update(**input_ids)
+        output_ids = self.model.generate(**kwargs, **self.generation_config)
+        if not self.model.config.is_encoder_decoder:
+            output_ids = output_ids[:, input_ids["input_ids"].shape[1]:]
+        output_list = self.batch_decode(output_ids)
+
+        return output_list
+
+    def __call__(self, *args, **kwargs):
+        r"""
+        Allows the HuggingfaceModel instance to be called like a function, which internally calls the model's
+        __call__ method.
+
+        :return: The output from the model's __call__ method.
+        """
+        return self.model(*args, **kwargs)
+
+    def tokenize(self, *args, **kwargs):
+        r"""
+        Tokenizes the input using the model's tokenizer.
+
+        :return: The tokenized output.
+        """
+        return self.tokenizer.tokenize(*args, **kwargs)
+
+    def batch_encode(self, *args, **kwargs):
+        return self.tokenizer(*args, **kwargs)
+
+    def batch_decode(self, *args, **kwargs)-> List[str]:
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def format(self, **kwargs):
+        return self.format_str.format(**kwargs)
+
+    def format_instance(self, query, jailbreak_prompt, response):
+        prompt = jailbreak_prompt.replace('{query}', query) # 之所以不用str.format方法，是因为jailbreak_prompt可能是任意字符串，其中可能包含{...}，而str.format不支持忽略未提供的key
+        return self.format(prompt=prompt, response=response)
+
+    @property
+    def device(self):
+        return self.model.device
+
+    @property
+    def dtype(self):
+        return self.model.dtype
+
+    @property
+    def bos_token_id(self):
+        return self.tokenizer.bos_token_id
+
+    @property
+    def eos_token_id(self):
+        return self.tokenizer.eos_token_id
+
+    @property
+    def pad_token_id(self):
+        return self.tokenizer.pad_token_id
+
+    @property
+    #@functools.cache  # 记忆结果，避免重复计算。不影响子类的重载。
+    def embed_layer(self) -> Optional[torch.nn.Embedding]:
+        """
+        Retrieve the embedding layer object of the model.
+
+        This method provides two commonly used approaches to search for the embedding layer in Hugging Face models.
+        If these methods are not effective, users should consider manually overriding this method.
+
+        Returns:
+            torch.nn.Embedding or None: The embedding layer object if found, otherwise None.
+        """
+        for module in self.model.modules():
+            if isinstance(module, torch.nn.Embedding):
+                return module
+
+        for module in self.model.modules():
+            if all(hasattr(module, attr) for attr in ['bos_token_id', 'eos_token_id', 'encode', 'decode', 'tokenize']):
+                return module
+        return None
+
+    @property
+    #@functools.cache
+    def vocab_size(self) -> int:
+        """
+        Get the vocabulary size.
+
+        This method provides two commonly used approaches for obtaining the vocabulary size in Hugging Face models.
+        If these methods are not effective, users should consider manually overriding this method.
+
+        Returns:
+            int: The size of the vocabulary.
+        """
+        if hasattr(self, 'config') and hasattr(self.config, 'vocab_size'):
+            return self.config.vocab_size
+
+        return self.embed_layer.weight.size(0)
+
+
+def from_pretrained(model_name_or_path: str, model_name: str, tokenizer_name_or_path: Optional[str] = None,
+                    dtype: Optional[torch.dtype] = None, **generation_config: Dict[str, Any]) -> HuggingfaceModel:
+    """
+    Imports a Hugging Face model and tokenizer with a single function call.
+
+    :param str model_name_or_path: The identifier or path for the pre-trained model.
+    :param str model_name: The name of the model, used for generating conversation template.
+    :param Optional[str] tokenizer_name_or_path: The identifier or path for the pre-trained tokenizer.
+        Defaults to `model_name_or_path` if not specified separately.
+    :param Optional[torch.dtype] dtype: The data type to which the model should be cast.
+        Defaults to None.
+    :param generation_config: Additional configuration options for model generation.
+    :type generation_config: dict
+
+    :return HuggingfaceModel: An instance of the HuggingfaceModel class containing the imported model and tokenizer.
+
+    .. note::
+        The model is loaded for evaluation by default. If `dtype` is specified, the model is cast to the specified data type.
+        The `tokenizer.padding_side` is set to 'right' if not already specified.
+        If the tokenizer has no specified pad token, it is set to the EOS token, and the model's config is updated accordingly.
+
+    **Example**
+
+    .. code-block:: python
+
+        model = from_pretrained('bert-base-uncased', 'bert-model', dtype=torch.float32, max_length=512)
+    """
+    if dtype is None:
+        dtype = 'auto'
+    model = AutoModelForCausalLM.from_pretrained(model_name_or_path, device_map='auto', trust_remote_code=True, low_cpu_mem_usage=True, torch_dtype=dtype).eval()
+    if tokenizer_name_or_path is None:
+        tokenizer_name_or_path = model_name_or_path
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, trust_remote_code=True)
+
+    if tokenizer.padding_side is None:
+        tokenizer.padding_side = 'right'
+
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+        model.config.pad_token_id = tokenizer.pad_token_id
+        model.generation_config.pad_token_id = tokenizer.pad_token_id
+
+    return HuggingfaceModel(model, tokenizer, model_name=model_name, generation_config=generation_config)

easyjailbreak/models/model_base.py

@@ -0,0 +1,117 @@
+"""
+Defines base classes for black-box and white-box model interface standards.
+The purpose is to unify models from various sources.
+There should be no specific algorithm implementations here.
+"""
+class ModelBase:
+    """
+    Defines a common model interface.
+    This base class is intended to provide a standardized interface for different types of models.
+    """
+    def generate(self, *args, **kwargs) -> str:
+        """
+        Generates output based on input arguments. This method must be implemented by subclasses.
+        :return str: The generated output.
+        """
+        raise NotImplementedError
+
+class WhiteBoxModelBase(ModelBase):
+    """
+    Defines the interface that white-box models should possess. Any user-defined white-box model should inherit from this class.
+    These models could be Hugging Face models or custom models written in PyTorch/TensorFlow, etc.
+    To maintain consistency with black-box models, this class integrates a tokenizer.
+    """
+    def __init__(self, model, tokenizer):
+        """
+        Initializes the white-box model with a model and a tokenizer.
+        :param model: The underlying model for generation.
+        :param tokenizer: The tokenizer used for processing input and output.
+        """
+        super().__init__()
+        self.model = model
+        self.tokenizer = tokenizer
+
+    def instance2str(self, instance, *args, **kwargs):
+        """
+        Converts an instance to a string. This method must be implemented by subclasses.
+        :param instance: The instance to be converted.
+        :return: A string representation of the instance.
+        """
+        raise NotImplementedError
+
+    @property
+    def device(self):
+        """
+        Returns the device on which the model is running.
+        :return: The device used by the model.
+        """
+        raise NotImplementedError
+
+    @property
+    def embed_layer(self):
+        """
+        Provides access to the embedding layer of the model.
+        :return: The embedding layer of the model.
+        """
+        raise NotImplementedError
+
+    @property
+    def vocab_size(self):
+        """
+        Returns the vocabulary size of the model.
+        :return: The size of the model's vocabulary.
+        """
+        raise NotImplementedError
+
+    @property
+    def bos_token_id(self):
+        """
+        Returns the Beginning-Of-Sequence token ID.
+        :return: The BOS token ID.
+        """
+        raise NotImplementedError
+
+    @property
+    def eos_token_id(self):
+        """
+        Returns the End-Of-Sequence token ID.
+        :return: The EOS token ID.
+        """
+        raise NotImplementedError
+
+    @property
+    def pad_token_id(self):
+        """
+        Returns the padding token ID.
+        :return: The padding token ID.
+        """
+        raise NotImplementedError
+
+    def __call__(self, *args, **kwargs):
+        """
+        Used to get logits, loss, and perform backpropagation, etc. This method must be implemented by subclasses.
+        """
+        raise NotImplementedError
+
+    def batch_encode(self, *args, **kwargs):
+        """
+        Encodes a batch of inputs. This method must be implemented by subclasses.
+        """
+        raise NotImplementedError
+
+    def batch_decode(self, *args, **kwargs):
+        """
+        Decodes a batch of outputs. This method must be implemented by subclasses.
+        """
+        raise NotImplementedError
+
+class BlackBoxModelBase(ModelBase):
+    """
+    Defines the interface that black-box models should possess. Any user-defined black-box model should inherit from this class.
+    These models could be like OpenAI's API or based on HTTP request services from third parties or self-built APIs.
+    """
+    def batch_generate(self, *args, **kwargs):
+        """
+        Uses asynchronous requests or multithreading to efficiently obtain batch responses. This method must be implemented by subclasses.
+        """
+        raise NotImplementedError

easyjailbreak/models/openai_model.py

@@ -0,0 +1,69 @@
+import logging
+import warnings
+from .model_base import BlackBoxModelBase
+from openai import OpenAI
+from fastchat.conversation import get_conv_template
+from httpx import URL
+
+
+class OpenaiModel(BlackBoxModelBase):
+    def __init__(self, model_name: str, api_keys: str, base_url: str, generation_config=None):
+        """
+        Initializes the OpenAI model with necessary parameters.
+        :param str model_name: The name of the model to use.
+        :param str api_keys: API keys for accessing the OpenAI service.
+        :param str template_name: The name of the conversation template, defaults to 'chatgpt'.
+        :param dict generation_config: Configuration settings for generation, defaults to an empty dictionary.
+        :param str|URL base_url: The base URL for the OpenAI API, defaults to None.
+        """
+        self.client = OpenAI(api_key=api_keys, base_url=base_url)
+        self.model_name = model_name
+        self.conversation = get_conv_template('chatgpt')
+        self.generation_config = generation_config if generation_config is not None else {}
+
+        self.base_url = base_url
+
+    def set_system_message(self, system_message: str):
+        """
+        Sets a system message for the conversation.
+        :param str system_message: The system message to set.
+        """
+        self.conversation.system_message = system_message
+
+    def generate(self, messages, clear_old_history=True, **kwargs):
+        """
+        Generates a response based on messages that include conversation history.
+        :param list[str]|str messages: A list of messages or a single message string.
+                                       User and assistant messages should alternate.
+        :param bool clear_old_history: If True, clears the old conversation history before adding new messages.
+        :return str: The response generated by the OpenAI model based on the conversation history.
+        """
+        if clear_old_history:
+            self.conversation.messages = []
+        if isinstance(messages, str):
+            messages = [messages]
+        for index, message in enumerate(messages):
+            self.conversation.append_message(self.conversation.roles[index % 2], message)
+        messages = self.conversation.to_openai_api_messages()
+        response = self.client.chat.completions.create(
+            model=self.model_name,
+            messages=messages,
+            **kwargs,
+            **self.generation_config
+        )
+        return response.choices[0].message.content
+
+    def batch_generate(self, conversations, **kwargs):
+        """
+        Generates responses for multiple conversations in a batch.
+        :param list[list[str]]|list[str] conversations: A list of conversations, each as a list of messages.
+        :return list[str]: A list of responses for each conversation.
+        """
+        responses = []
+        for conversation in conversations:
+            if isinstance(conversation, str):
+                warnings.warn(
+                    'For batch generation based on several conversations, provide a list[str] for each conversation. '
+                    'Using list[list[str]] will avoid this warning.')
+            responses.append(self.generate(conversation, **kwargs))
+        return responses

easyjailbreak/models/wenxinyiyan_model.py

@@ -0,0 +1,106 @@
+"""
+Wenxinyiyan Class
+============================================
+This class provides methods to interact with Baidu's Wenxin Workshop API for generating responses using an attack model.
+It includes methods for obtaining an access token and for sending requests to the API.
+https://cloud.baidu.com/?from=console
+"""
+from typing import List
+import requests
+import json
+import warnings
+from .model_base import BlackBoxModelBase
+
+class WenxinyiyanModel(BlackBoxModelBase):
+    r"""
+    A class for interacting with Baidu's Wenxin Workshop API.
+
+    This class allows users to generate text responses from Baidu's AI system
+    by providing a simple interface to the Wenxin Workshop API. It manages authentication
+    and request sending.
+    """
+    def __init__(self, API_KEY, SECRET_KEY):
+        """
+        Initializes the Wenxinyiyan instance with necessary credentials.
+        :param str API_KEY: The API key for Baidu's service.
+        :param str SECRET_KEY: The secret key for Baidu's service.
+        """
+        self.url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token="
+        self.API_KEY = API_KEY
+        self.SECRET_KEY = SECRET_KEY
+
+    @staticmethod
+    def get_access_token(API_KEY, SECRET_KEY):
+        """
+        使用 AK，SK 生成鉴权签名（Access Token）
+        :return: access_token，或是None(如果错误)
+        """
+        url = "https://aip.baidubce.com/oauth/2.0/token"
+        params = {"grant_type": "client_credentials", "client_id": API_KEY, "client_secret": SECRET_KEY}
+        return str(requests.post(url, params=params).json().get("access_token"))
+
+    def __call__(self, text_input):
+        url = self.url + self.get_access_token(self.API_KEY, self.SECRET_KEY)
+        payload = json.dumps({
+            "messages": [
+                {
+                    "role": "user",
+                    "content": text_input
+                }
+            ]
+        })
+        headers = {
+            'Content-Type': 'application/json'
+        }
+
+        response = requests.request("POST", url, headers=headers, data=payload)
+
+        return response.json()['result']
+
+#
+    def generate(self, messages: 'List[str],str')->str:
+
+        r"""
+        Generate a response based on messages that include conversation history.
+
+        :param ~List[str] messages: A list containing several messages.
+                                    The user and assistant messages should appear in turns.
+        :return: the response from the wenxinyiyan model based on a conversation history
+
+        Example:
+        messages = [
+         "你好",
+         "你好！有什么我可以帮助你的吗？请随时提出你的问题或需要帮助的内容，我会尽力提供准确和有用的答案。",
+         "我想知道明天天气",]
+        response = generate(messages)
+        """
+        # 判断message是str
+        if isinstance(messages, str):
+            messages = [messages]
+        url = self.url + self.get_access_token(self.API_KEY, self.SECRET_KEY)
+        processed_messages = []
+        roles = ('user', 'assistant')
+        for index, message in enumerate(messages):
+            processed_messages.append({
+                "role": roles[index % 2],
+                "content": message
+            })
+        payload = json.dumps({
+            "messages": processed_messages
+        })
+        headers = {
+            'Content-Type': 'application/json'
+        }
+
+        response = requests.request("POST", url, headers=headers, data=payload)
+        return response.json()['result']
+
+    def batch_generate(self, conversations, **kwargs):
+        responses = []
+        for conversation in conversations:
+            if isinstance(conversation, str):
+                warnings.warn('If you want the model to generate batches based on several conversations, '
+                              'please construct a list[str] for each conversation, or they will be divided into individual sentences. '
+                              'Switch input type of batch_generate() to list[list[str]] to avoid this warning.')
+            responses.append(self.generate(conversation))
+        return responses