'''
AutoDAN Class
============================================
This Class achieves a jailbreak method describe in the paper below.
This part of code is based on the code from the paper.

Paper title: AUTODAN: GENERATING STEALTHY JAILBREAK PROMPTS ON ALIGNED LARGE LANGUAGE MODELS
arXiv link: https://arxiv.org/abs/2310.04451
Source repository: https://github.com/SheltonLiu-N/AutoDAN.git
'''
import os
import json
import logging
import gc
import numpy as np
import torch
import torch.nn as nn
import time
import random
from fastchat import model
import nltk
from nltk.corpus import stopwords, wordnet
from transformers import AutoModelForCausalLM
from tqdm import tqdm
from itertools import islice
from easyjailbreak.attacker import AttackerBase
from easyjailbreak.datasets import JailbreakDataset
from easyjailbreak.datasets.instance import Instance
from easyjailbreak.mutation.generation import Rephrase
from easyjailbreak.mutation.rule import CrossOver, ReplaceWordsWithSynonyms
from easyjailbreak.metrics.Evaluator import EvaluatorPatternJudge
from easyjailbreak.seed import SeedTemplate

from tqdm import tqdm

__all__ = ["AutoDAN", "autodan_PrefixManager"]


def load_conversation_template(template_name):
    r"""
    load conversation template
    """
    if template_name == 'llama2':
        template_name = 'llama-2'
    conv_template = model.get_conversation_template(template_name)
    if conv_template.name == 'zero_shot':
        conv_template.roles = tuple(['### ' + r for r in conv_template.roles])
        conv_template.sep = '\n'
    elif conv_template.name == 'llama-2':
        conv_template.sep2 = conv_template.sep2.strip()
    return conv_template


def get_developer(model_name):
    r"""
    get model developer
    """
    developer_dict = {"llama2": "Meta"}
    return developer_dict[model_name]


def generate(model: AutoModelForCausalLM, tokenizer, input_ids, assistant_role_slice, gen_config=None):
    if gen_config is None:
        gen_config = model.generation_config
        gen_config.max_new_tokens = 128
    input_ids = input_ids[:assistant_role_slice.stop].to(model.device).unsqueeze(0)
    attn_masks = torch.ones_like(input_ids).to(model.device)
    output_ids = model.generate(input_ids,
                                attention_mask=attn_masks,
                                generation_config=gen_config,
                                pad_token_id=tokenizer.pad_token_id)[0]
    return output_ids[assistant_role_slice.stop:]


def forward(*, model, input_ids, attention_mask, batch_size=32):
    logits = []
    for i in range(0, input_ids.shape[0], batch_size):
        batch_input_ids = input_ids[i:i + batch_size]
        if attention_mask is not None:
            batch_attention_mask = attention_mask[i:i + batch_size]
        else:
            batch_attention_mask = None
        logits.append(model(input_ids=batch_input_ids, attention_mask=batch_attention_mask).logits)
        gc.collect()
    del batch_input_ids, batch_attention_mask
    return torch.cat(logits, dim=0)


class AutoDAN(AttackerBase):
    r"""
    AutoDAN is a class for conducting jailbreak attacks on language models.
    AutoDAN can automatically generate stealthy jailbreak prompts by hierarchical genetic algorithm.
    """

    def __init__(
            self,
            attack_model,
            target_model,
            jailbreak_datasets: JailbreakDataset,
            save_path,
            dataset_name,
            eval_model=None,
            max_query: int = 100,
            max_jailbreak: int = 100,
            max_reject: int = 100,
            max_iteration: int = 100,
            device='cuda:0',
            num_steps: int = 10,
            sentence_level_steps: int = 5,
            word_dict_size: int = 30,
            batch_size: int = 32,
            num_elites: float = 0.2,
            crossover_rate: float = 0.5,
            mutation_rate: float = 0.01,
            num_points: int = 5,
            model_name: str = "llama2",
            low_memory: int = 0,
            pattern_dict: dict = None,
    ):
        r"""
        Initialize the AutoDAN attack instance.
        :param ~model_wrapper attack_model: The model used to generate attack prompts.
        :param ~model_wrapper target_model: The target model to be attacked.
        :param ~JailbreakDataset jailbreak_datasets: The dataset containing harmful queries.
        :param ~model_wrapper eval_model: The model used for evaluating attck effectiveness during attacks.
        :param ~int num_steps: the number of paragraph-level iteration of AutoDAN-HGA algorithm.
        :param ~int sentence_level_steps: the number of sentence-level iteration of AutoDAN-HGA algorithm.
        :param ~int word_dict_size: the word_dict size of AutoDAN-HGA algorithm.
        :param ~int batch_size: the number of candidate prompts of each query.
        :param ~float num_elites: the proportion of elites used in Genetic Algorithm.
        :param ~float crossover_rate: the probability to execute crossover mutation.
        :param ~float mutation_rate: the probability to execute rephrase mutation.
        :param ~int num_points: the number of break points used in crossover mutation.
        :param ~str model_name: the target model name.
        :param ~int low_memory: 1 if low memory else 0
        :param ~dict pattern_dict: the pattern dictionary used in EvaluatorPatternJudge.
        """

        super().__init__(attack_model, target_model, eval_model, jailbreak_datasets)

        self.attack_results = JailbreakDataset([])

        self.current_query: int = 0
        self.current_jailbreak: int = 0
        self.current_reject: int = 0
        self.current_iteration: int = 0

        self.max_query: int = max_query
        self.max_jailbreak: int = max_jailbreak
        self.max_reject: int = max_reject
        self.max_iteration: int = max_iteration

        # AutoDAN特有参数
        self.device = device
        self.num_steps: int = num_steps
        self.sentence_level_steps: int = sentence_level_steps
        self.word_dict_size: int = word_dict_size
        self.batch_size: int = batch_size
        self.num_elites: float = num_elites
        self.crossover_rate: float = crossover_rate
        self.num_points: int = num_points
        self.mutation_rate: float = mutation_rate
        self.model_name: str = model_name
        self.low_memory: int = low_memory

        self.evaluator = EvaluatorPatternJudge(pattern_dict=pattern_dict)

        self.crossover_mutation = CrossOver(attr_name='jailbreak_prompt', num_points=self.num_points, seed_pool=None)
        self.replace_words_with_synonyms_mutation = ReplaceWordsWithSynonyms(attr_name="jailbreak_prompt", word_dict={})
        self.rephrase_mutation = Rephrase(model=self.attack_model)

        self.conv_template = load_conversation_template(model_name)

        # original candidate prompts
        self.reference = SeedTemplate().new_seeds(seeds_num=self.batch_size, prompt_usage='attack',
                                                  method_list=["AutoDAN-a"])
        for o in range(len(self.reference)):
            self.reference[o] = self.reference[o].replace('ChatGPT', self.model_name)
            self.reference[o] = self.reference[o].replace('chatGPT', self.model_name)
            self.reference[o] = self.reference[o].replace('ModelKeeper', get_developer(self.model_name))
            self.reference[o] = self.reference[o] + ' [PROMPT]:'

        self.save_path = save_path
        self.dataset_name = dataset_name

    def roulette_wheel_selection(self, data_list, score_list, num_selected):
        r"""
        apply roulette_wheel_selection on data_list
        """
        selection_probs = np.exp(score_list - np.max(score_list))
        selection_probs = selection_probs / selection_probs.sum()
        selected_indices = np.random.choice(len(data_list), size=num_selected, p=selection_probs, replace=True)
        selected_data = [data_list[i] for i in selected_indices]
        return selected_data

    def replace_with_synonyms(self, sentence, num=10):
        r"""
        replace words in sentence with synonyms
        """
        T = {"llama2", "meta", "vicuna", "lmsys", "guanaco", "theblokeai", "wizardlm", "mpt-chat",
             "mosaicml", "mpt-instruct", "falcon", "tii", "chatgpt", "modelkeeper", "prompt"}
        stop_words = set(stopwords.words('english'))
        words = nltk.word_tokenize(sentence)
        uncommon_words = [word for word in words if word.lower() not in stop_words and word.lower() not in T]
        selected_words = random.sample(uncommon_words, min(num, len(uncommon_words)))
        for word in selected_words:
            synonyms = wordnet.synsets(word)
            if synonyms and synonyms[0].lemmas():
                synonym = synonyms[0].lemmas()[0].name()
                sentence = sentence.replace(word, synonym, 1)
        return sentence

    def construct_momentum_word_dictionary(self, word_dict, individuals, score_list):
        r"""
        calculate momentum with score_list to maintain a momentum word_dict
        """
        word_scores = {}
        for individual, score in zip(individuals, score_list):
            T = {"llama2", "meta", "vicuna", "lmsys", "guanaco", "theblokeai", "wizardlm", "mpt-chat",
                 "mosaicml", "mpt-instruct", "falcon", "tii", "chatgpt", "modelkeeper", "prompt"}
            stop_words = set(stopwords.words('english'))
            words = nltk.word_tokenize(individual)
            uncommon_words = [word for word in words if word.lower() not in stop_words and word.lower() not in T]
            for word in uncommon_words:
                if word in word_scores.keys():
                    word_scores[word].append(score)
                else:
                    word_scores[word] = []
                    word_scores[word].append(score)
            for word, scores in word_scores.items():
                avg_score = sum(scores) * 1.0 / len(scores)
                if word in word_dict.keys():
                    word_dict[word] = (word_dict[word] + avg_score) / 2
                else:
                    word_dict[word] = avg_score
            sorted_word_dict = dict(sorted(word_dict.items(), key=lambda x: x[1], reverse=True))
            return dict(islice(sorted_word_dict.items(), self.word_dict_size))

    def get_score_autodan(self, conv_template, instruction, target, model, device, test_controls=None, crit=None):
        r"""
        Convert all test_controls to token ids and find the max length
        """
        input_ids_list = []
        target_slices = []
        for item in test_controls:
            prefix_manager = autodan_PrefixManager(tokenizer=self.target_model.tokenizer,
                                                   conv_template=conv_template,
                                                   instruction=instruction,
                                                   target=target,
                                                   adv_string=item)
            input_ids = prefix_manager.get_input_ids(adv_string=item).to(device)
            input_ids_list.append(input_ids)
            target_slices.append(prefix_manager._target_slice)

        # Pad all token ids to the max length
        pad_tok = 0
        for ids in input_ids_list:
            while pad_tok in ids:
                pad_tok += 1

        # Find the maximum length of input_ids in the list
        max_input_length = max([ids.size(0) for ids in input_ids_list])

        # Pad each input_ids tensor to the maximum length
        padded_input_ids_list = []
        for ids in input_ids_list:
            pad_length = max_input_length - ids.size(0)
            padded_ids = torch.cat([ids, torch.full((pad_length,), pad_tok, device=device)], dim=0)
            padded_input_ids_list.append(padded_ids)

        # Stack the padded input_ids tensors
        input_ids_tensor = torch.stack(padded_input_ids_list, dim=0)

        attn_mask = (input_ids_tensor != pad_tok).type(input_ids_tensor.dtype).to(device)

        # Forward pass and compute loss
        logits = forward(model=model, input_ids=input_ids_tensor, attention_mask=attn_mask,
                         batch_size=len(test_controls))
        losses = []
        for idx, target_slice in enumerate(target_slices):
            loss_slice = slice(target_slice.start - 1, target_slice.stop - 1)
            logits_slice = logits[idx, loss_slice, :].unsqueeze(0).transpose(1, 2)
            targets = input_ids_tensor[idx, target_slice].unsqueeze(0)
            loss = crit(logits_slice, targets)
            losses.append(loss)

        del input_ids_list, target_slices, input_ids_tensor, attn_mask
        gc.collect()
        return torch.stack(losses)

    def get_score_autodan_low_memory(self, conv_template, instruction, target, model, device, test_controls=None,
                                     crit=None):
        r"""
        Convert all test_controls to token ids and find the max length when memory is low
        """
        losses = []
        for item in test_controls:
            prefix_manager = autodan_PrefixManager(tokenizer=self.target_model.tokenizer,
                                                   conv_template=conv_template,
                                                   instruction=instruction,
                                                   target=target,
                                                   adv_string=item)
            input_ids = prefix_manager.get_input_ids(adv_string=item).to(device)
            input_ids_tensor = torch.stack([input_ids], dim=0)

            # Forward pass and compute loss
            logits = forward(model=model, input_ids=input_ids_tensor, attention_mask=None,
                             batch_size=len(test_controls))

            target_slice = prefix_manager._target_slice
            loss_slice = slice(target_slice.start - 1, target_slice.stop - 1)
            logits_slice = logits[0, loss_slice, :].unsqueeze(0).transpose(1, 2)
            targets = input_ids_tensor[0, target_slice].unsqueeze(0)
            loss = crit(logits_slice, targets)
            losses.append(loss)

        del input_ids_tensor
        gc.collect()
        return torch.stack(losses)

    def evaluate_candidate_prompts(self, sample: Instance, prefix_manager):
        r"""
        Calculate current candidate prompts scores of sample, get the currently best prompt and the corresponding response.
        """
        if self.low_memory == 1:
            losses = self.get_score_autodan_low_memory(
                conv_template=self.conv_template, instruction=sample.query, target=sample.reference_responses[0],
                model=self.target_model,
                device=self.device,
                test_controls=sample.candidate_prompts,
                crit=nn.CrossEntropyLoss(reduction='mean')
            )
        else:
            losses = self.get_score_autodan(
                conv_template=self.conv_template, instruction=sample.query, target=sample.reference_responses[0],
                model=self.target_model,
                device=self.device,
                test_controls=sample.candidate_prompts,
                crit=nn.CrossEntropyLoss(reduction='mean')
            )
        score_list = losses.cpu().numpy().tolist()

        best_new_adv_prefix_id = losses.argmin()
        best_new_adv_prefix = sample.candidate_prompts[best_new_adv_prefix_id]

        current_loss = losses[best_new_adv_prefix_id]

        adv_prefix = best_new_adv_prefix

        output_ids = generate(
            model=self.target_model.model,
            tokenizer=self.target_model.tokenizer,
            input_ids=prefix_manager.get_input_ids(adv_string=adv_prefix).to(self.device),
            assistant_role_slice=prefix_manager._assistant_role_slice,
            gen_config=None
        )
        response = self.target_model.tokenizer.decode(output_ids).strip()

        return score_list, current_loss, adv_prefix, response

    def update(self, Dataset: JailbreakDataset):
        r"""
        update jailbreak state
        """
        self.current_iteration += 1
        for instance in Dataset:
            self.current_jailbreak += instance.num_jailbreak
            self.current_query += instance.num_query
            self.current_reject += instance.num_reject

    def log(self):
        r"""
        Report the attack results.
        """
        logging.info("Jailbreak report:")
        logging.info(f"Total queries: {self.current_query}")
        logging.info(f"Total jailbreak: {self.current_jailbreak}")
        logging.info(f"Total reject: {self.current_reject}")
        logging.info(f"Total iteration: {self.current_iteration}")

    def attack(self):
        r"""
        Main loop for the attack process, iterate through jailbreak_datasets.
        """
        logging.info("Jailbreak started!")
        try:
            with open(self.save_path, 'w') as f:
                for instance in tqdm(self.jailbreak_datasets, desc="processing instance"):

                    if self.dataset_name == "trustllm":
                        self.target_model.set_system_message(instance.system_message)

                    new_instance = self.single_attack(instance)[0]
                    self.attack_results.add(new_instance)

                    line = new_instance.to_dict()
                    f.write(json.dumps(line, ensure_ascii=False) + '\n')

            self.update(self.attack_results)
        except KeyboardInterrupt:
            logging.info("Jailbreak interrupted by user!")
        self.log()
        logging.info("Jailbreak finished!")
        return self.attack_results

    def single_attack(self, instance: Instance):
        r"""
        Perform the AutoDAN-HGA algorithm on a single query.
        """
        best_prompt = ""
        user_prompt = instance.query
        target = instance.reference_responses[0]
        prefix_manager = autodan_PrefixManager(tokenizer=self.target_model.tokenizer,
                                               conv_template=self.conv_template,
                                               instruction=user_prompt,
                                               target=target,
                                               adv_string=self.reference[0])

        new_adv_prefixes = self.reference
        instance.candidate_prompts = new_adv_prefixes

        # 1. Initialize population with LLM-based Diversification
        for i in range(len(instance.candidate_prompts)):
            if random.random() < self.mutation_rate:
                instance.candidate_prompts[i] = self.rephrase_mutation.rephrase(instance.candidate_prompts[i])

        word_dict = {}
        # GENETIC ALGORITHM
        # Paragraph-level Iterations
        for j in range(self.num_steps):
            with torch.no_grad():
                epoch_start_time = time.time()

                # 2. Evaluate the fitness score of each individual in population
                score_list, current_loss, adv_prefix, response = self.evaluate_candidate_prompts(instance,
                                                                                                 prefix_manager)

                # 3. Evaluate jailbreak success or not
                instance.target_responses.append(response)
                self.evaluator(JailbreakDataset([instance]))
                is_success = instance.eval_results[-1]

                if is_success == 1:
                    epoch_end_time = time.time()
                    epoch_cost_time = round(epoch_end_time - epoch_start_time, 2)
                    print(
                        "################################\n"
                        f"Current Epoch: {j}/{self.num_steps}\n"
                        f"Passed:{is_success}\n"
                        f"Loss:{current_loss.item()}\n"
                        f"Epoch Cost:{epoch_cost_time}\n"
                        f"Current prefix:\n{adv_prefix}\n"
                        f"Current Response:\n{response}\n"
                        "################################\n")
                    gc.collect()
                    torch.cuda.empty_cache()
                    best_prompt = adv_prefix
                    break

                # 4. Sort the score_list and get corresponding control_prefixes
                score_list = [-x for x in score_list]
                sorted_indices = sorted(range(len(score_list)), key=lambda k: score_list[k], reverse=True)
                sorted_control_prefixes = [new_adv_prefixes[i] for i in sorted_indices]
                sorted_socre_list = [score_list[i] for i in sorted_indices]

                # 5. Select the elites
                num_elites = int(self.batch_size * self.num_elites)
                elites = sorted_control_prefixes[:num_elites]

                # 6. Use roulette wheel selection for the remaining positions
                parents_list = self.roulette_wheel_selection(sorted_control_prefixes[num_elites:],
                                                             sorted_socre_list[num_elites:],
                                                             self.batch_size - num_elites)
                instance.candidate_prompts = parents_list

                # 7. Apply crossover and mutation to the selected parents
                mutated_prompts = []
                mutation_dataset = JailbreakDataset([])
                for p in parents_list:
                    mutation_dataset.add(Instance(jailbreak_prompt=p))
                for i in range(0, len(parents_list), 2):
                    parent1 = mutation_dataset[i]
                    parent2 = mutation_dataset[i + 1] if (i + 1) < len(parents_list) else mutation_dataset[0]
                    if random.random() < self.crossover_rate:
                        dataset = self.crossover_mutation(JailbreakDataset([parent1]), other_instance=parent2)
                        child1 = dataset[0]
                        child2 = dataset[1]
                        mutated_prompts.append(child1.jailbreak_prompt)
                        mutated_prompts.append(child2.jailbreak_prompt)
                    else:
                        mutated_prompts.append(parent1.jailbreak_prompt)
                        mutated_prompts.append(parent2.jailbreak_prompt)
                for i in range(len(mutated_prompts)):
                    if random.random() < self.mutation_rate:
                        mutated_prompts[i] = self.rephrase_mutation.rephrase(mutated_prompts[i])

                # 8. Combine elites with the mutated offspring
                next_generation = elites + mutated_prompts
                assert len(next_generation) == self.batch_size
                instance.candidate_prompts = next_generation

                # HIERARCHICAL GENETIC ALGORITHM
                # Sentence-level Iterations
                for s in range(self.sentence_level_steps):
                    # 9. Evaluate the fitness score of each individual in population
                    score_list, current_loss, adv_prefix, response = self.evaluate_candidate_prompts(instance,
                                                                                                     prefix_manager)

                    # 10. Evaluate jailbreak success or not
                    instance.target_responses.append(response)
                    self.evaluator(JailbreakDataset([instance]))
                    is_success = instance.eval_results[-1]

                    if is_success == 1:
                        break

                    # 11. Calculate momentum word score and Update sentences in each prompt
                    word_dict = self.construct_momentum_word_dictionary(word_dict, instance.candidate_prompts,
                                                                        score_list)
                    self.replace_words_with_synonyms_mutation.update(word_dict)
                    mutation_dataset = JailbreakDataset([])
                    for p in instance.candidate_prompts:
                        mutation_dataset.add(Instance(jailbreak_prompt=p))
                    dataset = self.replace_words_with_synonyms_mutation(mutation_dataset)

                    mutated_prompts = []
                    for d in dataset:
                        mutated_prompts.append(d.jailbreak_prompt)
                    instance.candidate_prompts = mutated_prompts

                # 12. Evaluate the fitness score of each individual in population
                score_list, current_loss, adv_prefix, response = self.evaluate_candidate_prompts(instance,
                                                                                                 prefix_manager)

                # 13. Evaluate jailbreak success or not
                instance.target_responses.append(response)
                self.evaluator(JailbreakDataset([instance]))
                is_success = instance.eval_results[-1]

                if is_success == 1:
                    epoch_end_time = time.time()
                    epoch_cost_time = round(epoch_end_time - epoch_start_time, 2)
                    print(
                        "################################\n"
                        f"Current Epoch: {j}/{self.num_steps}\n"
                        f"Passed:{is_success}\n"
                        f"Loss:{current_loss.item()}\n"
                        f"Epoch Cost:{epoch_cost_time}\n"
                        f"Current prefix:\n{adv_prefix}\n"
                        f"Current Response:\n{response}\n"
                        "################################\n")
                    gc.collect()
                    torch.cuda.empty_cache()
                    best_prompt = adv_prefix
                    break

                epoch_end_time = time.time()
                epoch_cost_time = round(epoch_end_time - epoch_start_time, 2)
                print(
                    "################################\n"
                    f"Current Epoch: {j}/{self.num_steps}\n"
                    f"Passed:{is_success}\n"
                    f"Loss:{current_loss.item()}\n"
                    f"Epoch Cost:{epoch_cost_time}\n"
                    f"Current prefix:\n{adv_prefix}\n"
                    f"Current Response:\n{response}\n"
                    "################################\n")
                gc.collect()
                torch.cuda.empty_cache()
                best_prompt = adv_prefix

        new_instance = instance.copy()
        new_instance.parents.append(instance)
        instance.children.append(new_instance)
        new_instance.jailbreak_prompt = best_prompt + '{query}'
        return JailbreakDataset([new_instance])


class autodan_PrefixManager:
    def __init__(self, *, tokenizer, conv_template, instruction, target, adv_string):
        r"""
        :param ~str instruction: the harmful query.
        :param ~str target: the target response for the query.
        :param ~str adv_string: the jailbreak prompt.
        """
        self.tokenizer = tokenizer
        self.conv_template = conv_template
        self.instruction = instruction
        self.target = target
        self.adv_string = adv_string

    def get_prompt(self, adv_string=None):

        if adv_string is not None:
            self.adv_string = adv_string

        self.conv_template.append_message(self.conv_template.roles[0], f"{self.adv_string} {self.instruction} ")
        self.conv_template.append_message(self.conv_template.roles[1], f"{self.target}")
        prompt = self.conv_template.get_prompt()

        encoding = self.tokenizer(prompt)
        toks = encoding.input_ids

        if self.conv_template.name == 'llama-2':
            self.conv_template.messages = []

            self.conv_template.append_message(self.conv_template.roles[0], None)
            toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
            self._user_role_slice = slice(None, len(toks))

            self.conv_template.update_last_message(f"{self.instruction}")
            toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
            self._goal_slice = slice(self._user_role_slice.stop, max(self._user_role_slice.stop, len(toks)))

            separator = ' ' if self.instruction else ''
            self.conv_template.update_last_message(f"{self.adv_string}{separator}{self.instruction}")
            toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
            self._control_slice = slice(self._goal_slice.stop, len(toks))

            self.conv_template.append_message(self.conv_template.roles[1], None)
            toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
            self._assistant_role_slice = slice(self._control_slice.stop, len(toks))

            self.conv_template.update_last_message(f"{self.target}")
            toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
            self._target_slice = slice(self._assistant_role_slice.stop, len(toks) - 2)
            self._loss_slice = slice(self._assistant_role_slice.stop - 1, len(toks) - 3)

        else:
            python_tokenizer = False or self.conv_template.name == 'oasst_pythia'
            try:
                encoding.char_to_token(len(prompt) - 1)
            except:
                python_tokenizer = True

            if python_tokenizer:
                # This is specific to the vicuna and pythia tokenizer and conversation prompt.
                # It will not work with other tokenizers or prompts.
                self.conv_template.messages = []

                self.conv_template.append_message(self.conv_template.roles[0], None)
                toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
                self._user_role_slice = slice(None, len(toks))

                self.conv_template.update_last_message(f"{self.instruction}")
                toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
                self._goal_slice = slice(self._user_role_slice.stop, max(self._user_role_slice.stop, len(toks) - 1))

                separator = ' ' if self.instruction else ''
                self.conv_template.update_last_message(f"{self.adv_string}{separator}{self.instruction}")
                toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
                self._control_slice = slice(self._goal_slice.stop, len(toks) - 1)

                self.conv_template.append_message(self.conv_template.roles[1], None)
                toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
                self._assistant_role_slice = slice(self._control_slice.stop, len(toks))

                self.conv_template.update_last_message(f"{self.target}")
                toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
                self._target_slice = slice(self._assistant_role_slice.stop, len(toks) - 1)
                self._loss_slice = slice(self._assistant_role_slice.stop - 1, len(toks) - 2)
            else:
                self._system_slice = slice(
                    None,
                    encoding.char_to_token(len(self.conv_template.system))
                )
                self._user_role_slice = slice(
                    encoding.char_to_token(prompt.find(self.conv_template.roles[0])),
                    encoding.char_to_token(
                        prompt.find(self.conv_template.roles[0]) + len(self.conv_template.roles[0]) + 1)
                )
                self._goal_slice = slice(
                    encoding.char_to_token(prompt.find(self.instruction)),
                    encoding.char_to_token(prompt.find(self.instruction) + len(self.instruction))
                )
                self._control_slice = slice(
                    encoding.char_to_token(prompt.find(self.adv_string)),
                    encoding.char_to_token(prompt.find(self.adv_string) + len(self.adv_string))
                )
                self._assistant_role_slice = slice(
                    encoding.char_to_token(prompt.find(self.conv_template.roles[1])),
                    encoding.char_to_token(
                        prompt.find(self.conv_template.roles[1]) + len(self.conv_template.roles[1]) + 1)
                )
                self._target_slice = slice(
                    encoding.char_to_token(prompt.find(self.target)),
                    encoding.char_to_token(prompt.find(self.target) + len(self.target))
                )
                self._loss_slice = slice(
                    encoding.char_to_token(prompt.find(self.target)) - 1,
                    encoding.char_to_token(prompt.find(self.target) + len(self.target)) - 1
                )

        self.conv_template.messages = []

        return prompt

    def get_input_ids(self, adv_string=None):
        prompt = self.get_prompt(adv_string=adv_string)
        toks = self.tokenizer(prompt).input_ids
        input_ids = torch.tensor(toks[:self._target_slice.stop])
        return input_ids