script/ask.py

import argparse
import copy
import json
import os
import random
import sys
import time
import typing
import warnings

import openai
import tiktoken
from loguru import logger
from tenacity import Retrying, _utils, retry_if_not_exception_type
from tenacity.stop import stop_base
from tenacity.wait import wait_base

sys.path.append("..")

from model.crs_model import CRSModel

warnings.filterwarnings("ignore")


def get_exist_dialog_set():
    exist_id_set = set()
    for file in os.listdir(save_dir):
        file_id = os.path.splitext(file)[0]
        exist_id_set.add(file_id)
    return exist_id_set


def my_before_sleep(retry_state):
    logger.debug(
        f"Retrying: attempt {retry_state.attempt_number} ended with: {retry_state.outcome}, spend {retry_state.seconds_since_start} in total"
    )


class my_wait_exponential(wait_base):
    def __init__(
        self,
        multiplier: typing.Union[int, float] = 1,
        max: _utils.time_unit_type = _utils.MAX_WAIT,  # noqa
        exp_base: typing.Union[int, float] = 2,
        min: _utils.time_unit_type = 0,  # noqa
    ) -> None:
        self.multiplier = multiplier
        self.min = _utils.to_seconds(min)
        self.max = _utils.to_seconds(max)
        self.exp_base = exp_base

    def __call__(self, retry_state: "RetryCallState") -> float:
        if retry_state.outcome == openai.error.Timeout:
            return 0

        try:
            exp = self.exp_base ** (retry_state.attempt_number - 1)
            result = self.multiplier * exp
        except OverflowError:
            return self.max
        return max(max(0, self.min), min(result, self.max))


class my_stop_after_attempt(stop_base):
    """Stop when the previous attempt >= max_attempt."""

    def __init__(self, max_attempt_number: int) -> None:
        self.max_attempt_number = max_attempt_number

    def __call__(self, retry_state: "RetryCallState") -> bool:
        if retry_state.outcome == openai.error.Timeout:
            retry_state.attempt_number -= 1
        return retry_state.attempt_number >= self.max_attempt_number


def annotate_completion(prompt, logit_bias=None):
    if logit_bias is None:
        logit_bias = {}

    request_timeout = 20
    for attempt in Retrying(
        reraise=True,
        retry=retry_if_not_exception_type(
            (
                openai.error.InvalidRequestError,
                openai.error.AuthenticationError,
            )
        ),
        wait=my_wait_exponential(min=1, max=60),
        stop=(my_stop_after_attempt(8)),
    ):
        with attempt:
            response = openai.Completion.create(
                model="text-davinci-003",
                prompt=prompt,
                temperature=0,
                max_tokens=128,
                stop="Recommender",
                logit_bias=logit_bias,
                request_timeout=request_timeout,
            )["choices"][0]["text"]
        request_timeout = min(300, request_timeout * 2)

    return response


def get_instruction(dataset):
    if dataset == "redial_eval":
        item_with_year = True
        init_ask_instruction = """To recommend me items that I will accept, you can choose one of the following options.
A: ask my preference for genre
B: ask my preference for actor
C: ask my preference for director
D: I can directly give recommendations
Please enter the option character. Please only response a character."""
        ask_instruction = """To recommend me items that I will accept, you can choose one of the following options.
A: ask my preference for genre
B: ask my preference for actor
C: ask my preference for director
D: I can directly give recommendations
You have selected {}, do not repeat them. Please enter the option character."""
        option2attr = {
            "A": "genre",
            "B": "star",
            "C": "director",
            "D": "recommend",
        }
        option2temaplte = {
            "A": "Which genre do you like?",
            "B": "Which star do you like?",
            "C": "Which director do you like?",
        }
    elif dataset == "opendialkg_eval":
        item_with_year = False
        init_ask_instruction = """To recommend me items that I will accept, you can choose one of the following options.
A: ask my preference for genre
B: ask my preference for actor
C: ask my preference for director
D: ask my preference for writer
E: I can directly give recommendations
Please enter the option character. Please only response a character."""
        ask_instruction = """To recommend me items that I will accept, you can choose one of the following options.
A: ask my preference for genre
B: ask my preference for actor
C: ask my preference for director
D: ask my preference for writer
E: I can directly give recommendations
You have selected {}, do not repeat them. Please enter the option character."""
        option2attr = {
            "A": "genre",
            "B": "actor",
            "C": "director",
            "D": "writer",
            "E": "recommend",
        }
        option2temaplte = {
            "A": "Which genre do you like?",
            "B": "Which actor do you like?",
            "C": "Which director do you like?",
            "D": "Which writer do you like?",
        }
    else:
        raise Exception("do not support this dataset")

    if item_with_year is True:
        rec_instruction = "Please give me 10 recommendations according to my preference (Format: no. title (year if exists). No other things except the movie list in your response)."
    else:
        rec_instruction = "Please give me 10 recommendations according to my preference (Format: no. title. No other things except the item list in your response). You can recommend mentioned items in our dialog."

    return (
        init_ask_instruction,
        ask_instruction,
        rec_instruction,
        option2attr,
        option2temaplte,
    )


def get_model_args(model_name):
    if model_name == "kbrd":
        args_dict = {
            "debug": args.debug,
            "kg_dataset": args.kg_dataset,
            "hidden_size": args.hidden_size,
            "entity_hidden_size": args.entity_hidden_size,
            "num_bases": args.num_bases,
            "rec_model": args.rec_model,
            "conv_model": args.conv_model,
            "context_max_length": args.context_max_length,
            "entity_max_length": args.entity_max_length,
            "tokenizer_path": args.tokenizer_path,
            "encoder_layers": args.encoder_layers,
            "decoder_layers": args.decoder_layers,
            "text_hidden_size": args.text_hidden_size,
            "attn_head": args.attn_head,
            "resp_max_length": args.resp_max_length,
            "seed": args.seed,
        }
    elif model_name == "barcor":
        args_dict = {
            "debug": args.debug,
            "kg_dataset": args.kg_dataset,
            "rec_model": args.rec_model,
            "conv_model": args.conv_model,
            "context_max_length": args.context_max_length,
            "resp_max_length": args.resp_max_length,
            "tokenizer_path": args.tokenizer_path,
            "seed": args.seed,
        }
    elif model_name == "unicrs":
        args_dict = {
            "debug": args.debug,
            "seed": args.seed,
            "kg_dataset": args.kg_dataset,
            "tokenizer_path": args.tokenizer_path,
            "context_max_length": args.context_max_length,
            "entity_max_length": args.entity_max_length,
            "resp_max_length": args.resp_max_length,
            "text_tokenizer_path": args.text_tokenizer_path,
            "rec_model": args.rec_model,
            "conv_model": args.conv_model,
            "model": args.model,
            "num_bases": args.num_bases,
            "text_encoder": args.text_encoder,
        }
    elif model_name == "chatgpt":
        args_dict = {
            "seed": args.seed,
            "debug": args.debug,
            "kg_dataset": args.kg_dataset,
        }

    return args_dict


if __name__ == "__main__":
    local_time = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
    warnings.filterwarnings("ignore")

    parser = argparse.ArgumentParser()
    parser.add_argument("--api_key")
    parser.add_argument(
        "--dataset", type=str, choices=["redial_eval", "opendialkg_eval"]
    )
    parser.add_argument("--turn_num", type=int, default=5)
    parser.add_argument(
        "--crs_model",
        type=str,
        choices=["kbrd", "barcor", "unicrs", "chatgpt"],
    )

    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--debug", action="store_true")
    parser.add_argument("--kg_dataset", type=str, choices=["redial", "opendialkg"])

    # model_detailed
    parser.add_argument("--hidden_size", type=int)
    parser.add_argument("--entity_hidden_size", type=int)
    parser.add_argument("--num_bases", type=int, default=8)
    parser.add_argument("--context_max_length", type=int)
    parser.add_argument("--entity_max_length", type=int)

    # model
    parser.add_argument("--rec_model", type=str)
    parser.add_argument("--conv_model", type=str)

    # conv
    parser.add_argument("--tokenizer_path", type=str)
    parser.add_argument("--encoder_layers", type=int)
    parser.add_argument("--decoder_layers", type=int)
    parser.add_argument("--text_hidden_size", type=int)
    parser.add_argument("--attn_head", type=int)
    parser.add_argument("--resp_max_length", type=int)

    # prompt
    parser.add_argument("--model", type=str)
    parser.add_argument("--text_tokenizer_path", type=str)
    parser.add_argument("--text_encoder", type=str)

    args = parser.parse_args()
    openai.api_key = args.api_key
    save_dir = f"../save_{args.turn_num}/ask/{args.crs_model}/{args.dataset}"
    os.makedirs(save_dir, exist_ok=True)

    random.seed(args.seed)

    # recommender
    recommendation_template = "I would recommend the following items:\n\n{}"

    # recommender
    model_args = get_model_args(args.crs_model)
    recommender = CRSModel(crs_model=args.crs_model, **model_args)

    # seeker
    (
        init_ask_instruction,
        ask_instruction,
        rec_instruction,
        option2attr,
        option2template,
    ) = get_instruction(args.dataset)
    options = list(option2attr.keys())

    # scorer
    persuasiveness_template = """Does the explanation make you want to accept the recommendation? Please give your score.
If mention one of [{}], give 2.
Else if you think recommended items are worse than [{}], give 0.
Else if you think recommended items are comparable to [{}] according to the explanation, give 1.
Else if you think recommended items are better than [{}] according to the explanation, give 2.
Only answer the score number."""
    encoding = tiktoken.encoding_for_model("text-davinci-003")
    logit_bias = {encoding.encode(str(score))[0]: 10 for score in range(3)}

    with open(f"../data/{args.kg_dataset}/entity2id.json", "r", encoding="utf-8") as f:
        entity2id = json.load(f)
    id2entity = {}
    for k, v in entity2id.items():
        id2entity[int(v)] = k
    entity_list = list(entity2id.keys())

    name2id = {}

    with open(f"../data/{args.kg_dataset}/id2info.json", "r", encoding="utf-8") as f:
        id2info = json.load(f)

    for k, v in id2info.items():
        name2id[v["name"]] = k

    dialog_id2data = {}
    with open(
        f"../data/{args.dataset}/test_data_processed.jsonl", encoding="utf-8"
    ) as f:
        lines = f.readlines()
        for line in lines:
            line = json.loads(line)
            dialog_id = str(line["dialog_id"]) + "_" + str(line["turn_id"])
            dialog_id2data[dialog_id] = line

    dialog_id_set = set(dialog_id2data.keys()) - get_exist_dialog_set()
    while len(dialog_id_set) > 0:
        print(len(dialog_id_set))
        dialog_id = random.choice(tuple(dialog_id_set))

        data = dialog_id2data[dialog_id]
        conv_dict = copy.deepcopy(data)  # for model
        goal_item_list = [f'"{item}"' for item in conv_dict["rec"]]
        goal_item_str = ", ".join(goal_item_list)
        rec_labels = [name2id[rec] for rec in data["rec"]]

        context_dict = []  # for save
        for i, text in enumerate(conv_dict["context"]):
            if len(text) == 0:
                continue
            if i % 2 == 0:
                role_str = "user"
            else:
                role_str = "assistant"
            context_dict.append({"role": role_str, "content": text})

        # dialog state
        rec_success = False
        asked_options = []
        option2index = {"A": 0, "B": 1, "C": 2, "D": 3, "E": 4}
        if args.kg_dataset == "redial":
            state = [0, 0, 0, 0]
        elif args.kg_dataset == "opendialkg":
            state = [0, 0, 0, 0, 0]

        for i in range(0, args.turn_num):
            # seeker
            # choose option

            if args.crs_model == "chatgpt":
                conv_dict["context"].append(init_ask_instruction)

            # recommender
            # options (list of str): available options, generate one of them
            gen_inputs, recommender_text = recommender.get_conv(conv_dict)
            if args.crs_model != "chatgpt":
                recommender_choose = recommender.get_choice(gen_inputs, options, state)
            else:
                recommender_choose = recommender.get_choice(
                    gen_inputs, options, state, conv_dict
                )
            selected_option = recommender_choose

            if selected_option == options[-1]:  # choose to rec
                # recommender
                rec_items, rec_truth = recommender.get_rec(conv_dict)
                rec_pred = rec_items[0]

                rec_items_str = ""
                for j, rec_item in enumerate(rec_pred[:50]):
                    rec_items_str += f"{i + 1}: {id2entity[rec_item]}\n"
                recommender_text = recommendation_template.format(rec_items_str)

                # judge whether success
                for rec_label in rec_truth:
                    if rec_label in rec_pred:
                        rec_success = True
                        break

                context_dict.append(
                    {
                        "role": "assistant",
                        "content": recommender_text,
                        "rec_items": rec_pred,
                        "rec_success": rec_success,
                        "option": selected_option,
                    }
                )
                conv_dict["context"].append(recommender_text)

                # seeker
                if rec_success is True:
                    seeker_text = "That's perfect, thank you!"
                else:
                    seeker_text = "I don't like them."

                context_dict.append({"role": "user", "content": seeker_text})
                conv_dict["context"].append(seeker_text)

            else:  # choose to ask
                recommender_text = option2template[selected_option]
                context_dict.append(
                    {
                        "role": "assistant",
                        "content": recommender_text,
                        "option": selected_option,
                    }
                )
                conv_dict["context"].append(recommender_text)

                # seeker
                ask_attr = option2attr[selected_option]

                # update state
                state[option2index[selected_option]] = -1e5

                ans_attr_list = []
                for label_id in rec_labels:
                    if str(label_id) in id2info and ask_attr in id2info[str(label_id)]:
                        ans_attr_list.extend(id2info[str(label_id)][ask_attr])
                if len(ans_attr_list) > 0:
                    seeker_text = ", ".join(list(set(ans_attr_list)))
                else:
                    seeker_text = "Sorry, no information about this, please choose another option."

                context_dict.append(
                    {
                        "role": "user",
                        "content": seeker_text,
                        "entity": ans_attr_list,
                    }
                )
                conv_dict["context"].append(seeker_text)
                conv_dict["entity"] += ans_attr_list

            if rec_success is True:
                break

        # score persuasiveness
        # seeker_prompt = ''
        # for turn_dict in context_dict:
        #     if turn_dict['role'] == 'user':
        #         role_str = 'Seeker'
        #     else:
        #         role_str = 'Recommender'
        #     seeker_prompt += f'{role_str}: {turn_dict["content"]}\n'
        # persuasiveness_str = persuasiveness_template.format(goal_item_str, goal_item_str, goal_item_str,
        #                                                          goal_item_str)
        # prompt_str_for_persuasiveness = seeker_prompt + persuasiveness_str
        # prompt_str_for_persuasiveness += '\nSeeker:'

        # persuasiveness_score = annotate_completion(prompt_str_for_persuasiveness, logit_bias).strip()

        # save
        conv_dict["context"] = context_dict
        data["simulator_dialog"] = conv_dict

        with open(f"{save_dir}/{dialog_id}.json", "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)

        dialog_id_set -= get_exist_dialog_set()