NLPBench: Evaluating NLP-Related Problem-solving Ability in Large Language Models
import json
import numpy as np
from os import listdir
from os.path import isfile, join

from utils import ApiManager

cates = [
    "Language Modeling and Syntax and Parsing",
    "Pragmatics and Discourse and Dialogue and Applications",
    "Semantics and Logic",
    "Information Retrieval and Topic Modeling",
    "Artificial Intelligence",
    "Other Topics"
]


class Evaluation:
    def __init__(self, seed, open_sourced):
        self.open_sourced = open_sourced
        self.seed = seed
        osd = '_os' if open_sourced else ''
        self.osd = osd

        self.no_ctx_files = [f for f in listdir(f'res/{seed}{osd}/no_ctx/') if
                             isfile(join(f'res/{seed}{osd}/no_ctx/', f))]
        self.ctx_files = [f for f in listdir(f'res/{seed}{osd}/ctx/') if isfile(join(f'res/{seed}{osd}/ctx/', f))]
        self.ctx_data = json.load(open('data/w_ctx.json', 'r'))
        self.no_ctx_data = json.load(open('data/wo_ctx.json', 'r'))

    def _init_result(self):
        return {
            'gpt-3.5-turbo': {},
            'gpt-4': {},
            'text-bison-001': {}
        } if self.open_sourced is False else {
            'Llama-2-70b-chat-hf': {},
            'Llama-2-13b-chat-hf': {},
        }

    def _file_format(self, f, ctx_path='no_ctx'):
        args = f.split('_')
        llm_name = f.split('_')[1].replace('.json', '')
        shot_type = f.split('_')[0]
        prompt_r = ''
        if len(args) == 3:
            prompt_r = f.split('_')[2].split('.')[0]
        if len(args) == 4:
            prompt_r = f.split('_')[2] + '_' + f.split('_')[3].split('.')[0]
        llm_res = json.load(open(f'res/{self.seed}{self.osd}/{ctx_path}/{f}', 'r'))
        if prompt_r != '':
            prompt_r = '_' + prompt_r

        return llm_name, shot_type, prompt_r, llm_res

    def _path_format(self, suffix):
        no_ctx_save_path = f'res/{self.seed}{self.osd}/res_no_ctx_{suffix}.json'
        ctx_save_path = f'res/{self.seed}{self.osd}/res_ctx_{suffix}.json'

        return no_ctx_save_path, ctx_save_path

    # We use accuracy to evaluate multiple choice questions
    def evaluate_mc(self):
        ### w/o context
        res = self._init_result()
        no_ctx_save_path, ctx_save_path = self._path_format('mc')

        for f in self.no_ctx_files:
            llm_name, shot_type, prompt_r, llm_res = self._file_format(f)
            corr = {c: 0 for c in cates}
            cnt = {c: 0 for c in cates}
            for i, d in enumerate(llm_res):
                if d['type'] == 0:
                    cate = self.no_ctx_data[i]['category']
                    cate_cnt = cnt.get(cate, 0)
                    cnt[cate] = cate_cnt + 1
                    if d['retrived_answer'] is not None:
                        ans = set(d['answer'])
                        llm_ans = d['retrived_answer'].replace(' ', '').replace("'", '').replace('"', '').split(',')
                        try:
                            llm_ans = [int(a) for a in llm_ans]
                        except Exception:
                            cate_ans = corr.get(cate, 0)
                            corr[cate] = cate_ans
                        llm_ans = set(llm_ans)
                        if ans == llm_ans:
                            cate_ans = corr.get(cate, 0)
                            corr[cate] = cate_ans + 1
                        else:
                            cate_ans = corr.get(cate, 0)
                            corr[cate] = cate_ans

            acc = {k: v / cnt[k] if cnt[k] != 0 else 1 for k, v in corr.items()}
            res['count'] = cnt
            res['total_count'] = sum(cnt.values())
            overall_acc = np.sum(np.array(list(acc.values())) * np.array(list(cnt.values()))) / sum(cnt.values())
            res[llm_name][f'{shot_type}{prompt_r}'] = {
                'acc': acc,
                'overall_acc': overall_acc
            }
        json.dump(res, open(no_ctx_save_path, 'w'), indent=4)

        ### w/ context
        res = self._init_result()
        for f in self.ctx_files:
            llm_name, shot_type, prompt_r, llm_res = self._file_format(f, 'ctx')
            corr = {c: 0 for c in cates}
            cnt = {c: 0 for c in cates}
            for i, ds in enumerate(llm_res):
                for j, d in enumerate(ds['questions']):
                    if ds['type'][j] == 0:
                        cate = self.ctx_data[i]['category'][j]
                        cate_cnt = cnt.get(cate, 0)
                        cnt[cate] = cate_cnt + 1
                        if ds['retrived_answer'][j] is not None:
                            ans = set(ds['answers'][j])
                            llm_ans = ds['retrived_answer'][j].replace(' ', '').replace("'", '').replace('"', '').split(
                                ',')
                            try:
                                llm_ans = [int(a) for a in llm_ans]
                            except Exception:
                                cate_ans = corr.get(cate, 0)
                                corr[cate] = cate_ans
                                continue
                            llm_ans = set(llm_ans)
                            if ans == llm_ans:
                                cate_ans = corr.get(cate, 0)
                                corr[cate] = cate_ans + 1
                            else:
                                cate_ans = corr.get(cate, 0)
                                corr[cate] = cate_ans

            acc = {k: v / cnt[k] if cnt[k] != 0 else 1 for k, v in corr.items()}
            res['count'] = cnt
            res['total_count'] = sum(cnt.values())
            overall_acc = np.sum(np.array(list(acc.values())) * np.array(list(cnt.values()))) / sum(cnt.values())
            res[llm_name][f'{shot_type}{prompt_r}'] = {
                'acc': acc,
                'overall_acc': overall_acc
            }
        json.dump(res, open(ctx_save_path, 'w'), indent=4)

    def evaluate_sc(self):
        ### w/o context
        res = self._init_result()
        no_ctx_save_path, ctx_save_path = self._path_format('mc-sc')
        files = [
            'few-shot-sc_gpt-3.5-turbo.json',
            'few-shot-sc_gpt-3.5-turbo_cot.json',
            'few-shot-sc_gpt-4.json',
            'few-shot-sc_gpt-4_cot.json',
            'few-shot-sc_text-bison-001.json',
            'few-shot-sc_text-bison-001_cot.json',
            'zero-shot-sc_gpt-3.5-turbo.json',
            'zero-shot-sc_gpt-3.5-turbo_cot.json',
            'zero-shot-sc_gpt-4.json',
            'zero-shot-sc_gpt-4_cot.json',
            'zero-shot-sc_text-bison-001.json',
            'zero-shot-sc_text-bison-001_cot.json',
        ]
        for f in files:
            llm_name, shot_type, prompt_r, llm_res = self._file_format(f)
            corr = {c: 0 for c in cates}
            cnt = {c: 0 for c in cates}
            for i, d in enumerate(llm_res):
                if d['type'] == 0:
                    cate = self.no_ctx_data[i]['category']
                    cate_cnt = cnt.get(cate, 0)
                    cnt[cate] = cate_cnt + 1
                    if d['llm_answer'] is not None:
                        ans = set(d['answer'])
                        llm_ans = d['llm_answer'].replace(' ', '').replace("'", '').replace('"', '').split(',')
                        try:
                            llm_ans = [int(a) for a in llm_ans]
                        except Exception:
                            print('model:', f, '\nerr: ', llm_ans)
                            cate_ans = corr.get(cate, 0)
                            corr[cate] = cate_ans
                        llm_ans = set(llm_ans)
                        if ans == llm_ans:
                            cate_ans = corr.get(cate, 0)
                            corr[cate] = cate_ans + 1
                        else:
                            cate_ans = corr.get(cate, 0)
                            corr[cate] = cate_ans

            acc = {k: v / cnt[k] if cnt[k] != 0 else 1 for k, v in corr.items()}
            res['count'] = cnt
            res['total_count'] = sum(cnt.values())
            overall_acc = np.sum(np.array(list(acc.values())) * np.array(list(cnt.values()))) / sum(cnt.values())
            res[llm_name][f'{shot_type}{prompt_r}'] = {
                'acc': acc,
                'overall_acc': overall_acc
            }
        json.dump(res, open(no_ctx_save_path, 'w'), indent=4)

        ### w/ context
        res = self._init_result()
        for f in files:
            llm_name, shot_type, prompt_r, llm_res = self._file_format(f, 'ctx')
            corr = {c: 0 for c in cates}
            cnt = {c: 0 for c in cates}
            for i, ds in enumerate(llm_res):
                for j, d in enumerate(ds['questions']):
                    if ds['type'][j] == 0:
                        cate = self.ctx_data[i]['category'][j]
                        cate_cnt = cnt.get(cate, 0)
                        cnt[cate] = cate_cnt + 1
                        if ds['llm_answer'][j] is not None:
                            ans = set(ds['answers'][j])
                            llm_ans = ds['llm_answer'][j].replace(' ', '').replace("'", '').replace('"', '').split(
                                ',')
                            try:
                                llm_ans = [int(a) for a in llm_ans]
                            except Exception:
                                print('model:', f, '\nerr: ', llm_ans)
                                cate_ans = corr.get(cate, 0)
                                corr[cate] = cate_ans
                                continue
                            llm_ans = set(llm_ans)
                            if ans == llm_ans:
                                cate_ans = corr.get(cate, 0)
                                corr[cate] = cate_ans + 1
                            else:
                                cate_ans = corr.get(cate, 0)
                                corr[cate] = cate_ans

            acc = {k: v / cnt[k] if cnt[k] != 0 else 1 for k, v in corr.items()}
            res['count'] = cnt
            res['total_count'] = sum(cnt.values())
            overall_acc = np.sum(np.array(list(acc.values())) * np.array(list(cnt.values()))) / sum(cnt.values())
            res[llm_name][f'{shot_type}{prompt_r}'] = {
                'acc': acc,
                'overall_acc': overall_acc
            }
        json.dump(res, open(ctx_save_path, 'w'), indent=4)

    # We use ROUGE-L, CIDEr (for unique answer) to evaluate short answer questions with unique answer
    def evaluate_sa_unique(self):
        from pycocoevalcap.rouge.rouge import Rouge
        from pycocoevalcap.cider.cider import Cider

        res = self._init_result()
        no_ctx_save_path, ctx_save_path = self._path_format('sa_unique')

        for f in self.no_ctx_files:
            llm_name, shot_type, prompt_r, llm_res = self._file_format(f)

            llm_ans_dict = {}
            ans_dict = {}
            overall_llm_ans = {}
            overall_ans = {}
            cnt = {}
            for i, d in enumerate(llm_res):
                if d['type'] == 1 and self.no_ctx_data[i].get('unique_ans', None) == 1:
                    cate = self.no_ctx_data[i]['category']
                    cate_cnt = cnt.get(cate, 0)
                    cnt[cate] = cate_cnt + 1

                    ans = d['answer']
                    llm_ans = d['llm_answer']

                    if llm_ans == "" or llm_ans is None:
                        llm_ans = "No answer provided."

                    tmp1 = llm_ans_dict.get(cate, {})
                    tmp1[i] = [llm_ans]
                    llm_ans_dict[cate] = tmp1

                    tmp2 = ans_dict.get(cate, {})
                    tmp2[i] = [ans]
                    ans_dict[cate] = tmp2

                    overall_llm_ans[i] = [llm_ans]
                    overall_ans[i] = [ans]

            scores = {}
            for k in cnt.keys():
                rouge = Rouge().compute_score(llm_ans_dict[k], ans_dict[k])[0]
                cider = Cider().compute_score(llm_ans_dict[k], ans_dict[k])[0]
                scores[k] = {
                    'ROUGE-L': rouge,
                    'CIDEr': cider,
                }
            res['count'] = cnt
            res['total_count'] = sum(cnt.values())
            res[llm_name][f'{shot_type}{prompt_r}'] = {
                'scores': scores,
                'avg_score': {
                    'ROUGE-L': Rouge().compute_score(overall_llm_ans, overall_ans)[0],
                    'CIDEr': Cider().compute_score(overall_llm_ans, overall_ans)[0],
                }
            }
        json.dump(res, open(no_ctx_save_path, 'w'), indent=4)

        ### w/ context
        res = self._init_result()
        for f in self.ctx_files:
            llm_name, shot_type, prompt_r, llm_res = self._file_format(f, 'ctx')
            llm_ans_dict = {}
            ans_dict = {}
            overall_llm_ans = {}
            overall_ans = {}
            cnt = {}
            for i, ds in enumerate(llm_res):
                for j, d in enumerate(ds['questions']):
                    if ds['type'][j] == 1 and self.ctx_data[i].get('unique_ans', [0 for _ in range(j + 1)])[j] == 1:
                        cate = self.ctx_data[i]['category'][j]
                        cate_cnt = cnt.get(cate, 0)
                        cnt[cate] = cate_cnt + 1

                        ans = ds['answers'][j]
                        llm_ans = ds['llm_answer'][j]

                        if llm_ans == "" or llm_ans is None:
                            llm_ans = "No answer provided."

                        tmp1 = llm_ans_dict.get(cate, {})
                        tmp1[f'{i}_{j}'] = [llm_ans]
                        llm_ans_dict[cate] = tmp1

                        tmp2 = ans_dict.get(cate, {})
                        tmp2[f'{i}_{j}'] = [ans]
                        ans_dict[cate] = tmp2

                        overall_llm_ans[f'{i}_{j}'] = [llm_ans]
                        overall_ans[f'{i}_{j}'] = [ans]

            scores = {}
            for k in cnt.keys():
                rouge = Rouge().compute_score(llm_ans_dict[k], ans_dict[k])[0]
                cider = Cider().compute_score(llm_ans_dict[k], ans_dict[k])[0]
                scores[k] = {
                    'ROUGE-L': rouge,
                    'CIDEr': cider,
                }

            res['count'] = cnt
            res['total_count'] = sum(cnt.values())
            res[llm_name][f'{shot_type}{prompt_r}'] = {
                'scores': scores,
                'avg_score': {
                    'ROUGE-L': Rouge().compute_score(overall_llm_ans, overall_ans)[0],
                    'CIDEr': Cider().compute_score(overall_llm_ans, overall_ans)[0],
                }
            }
        json.dump(res, open(ctx_save_path, 'w'), indent=4)

    # We use GPT-4 to evaluate short answer questions
    def evaluate_sa(self):

        res = self._init_result()
        no_ctx_save_path, ctx_save_path = self._path_format('sa')
        api = ApiManager(
            model_name='gpt-4',
            seed=41,
            default_max_tokens=20,
            temperature=0,
        )
        prompt_SA_EVAL = '''You are a NLP professional assistant, your work is to evaluate whether the student's answer is correct for the given short answer question. 
        A teacher answer is also provided, your evaluation should based on the teacher answer.
        If the student is correct, return 1, else return 0.
        Your response should ONLY contain 0 or 1.

        Short answer question:
        "{q}"
        
        Teacher answer:
        "{eg_ans}"

        Student answer (evaluate this answer):
        "{llm_ans}"

        Your response:
        '''

        for f in self.no_ctx_files:
            llm_name, shot_type, prompt_r, llm_res = self._file_format(f)
            corr = {c: 0 for c in cates}
            cnt = {c: 0 for c in cates}
            for i, d in enumerate(llm_res):
                if d['type'] == 1:
                    print(f'Processing Q.{i}')
                    cate = self.no_ctx_data[i]['category']
                    cate_cnt = cnt.get(cate, 0)
                    cnt[cate] = cate_cnt + 1

                    q = d['question']
                    ans = d['answer']
                    llm_ans = d['llm_answer']
                    if llm_ans == "" or llm_ans is None:
                        llm_ans = "No answer provided."

                    score = api(
                        [{'role': 'user', 'content': prompt_SA_EVAL.format(q=q, eg_ans=ans, llm_ans=llm_ans)}]
                    )
                    try:
                        score = int(score)
                    except ValueError:
                        score = 0
                    if score == 1:
                        cate_ans = corr.get(cate, 0)
                        corr[cate] = cate_ans + 1
                    else:
                        cate_ans = corr.get(cate, 0)
                        corr[cate] = cate_ans
                    d['retrived_answer'] = score

            acc = {k: v / cnt[k] if cnt[k] != 0 else 1 for k, v in corr.items()}
            res['count'] = cnt
            res['total_count'] = sum(cnt.values())
            overall_acc = np.sum(np.array(list(acc.values())) * np.array(list(cnt.values()))) / sum(cnt.values())
            res[llm_name][f'{shot_type}{prompt_r}'] = {
                'acc': acc,
                'overall_acc': overall_acc
            }
            json.dump(llm_res, open(f'res/{self.seed}{self.osd}/no_ctx/{f}', 'w'), indent=4)
        json.dump(res, open(no_ctx_save_path, 'w'), indent=4)

        ### w/ context
        res = self._init_result()
        for f in self.ctx_files:
            llm_name, shot_type, prompt_r, llm_res = self._file_format(f, 'ctx')
            corr = {c: 0 for c in cates}
            cnt = {c: 0 for c in cates}
            for i, ds in enumerate(llm_res):
                tmp = []
                for j, d in enumerate(ds['questions']):
                    if ds['type'][j] == 1:
                        print(f'Processing Q.{i}_{j}')
                        cate = self.ctx_data[i]['category'][j]
                        cate_cnt = cnt.get(cate, 0)
                        cnt[cate] = cate_cnt + 1

                        q = ds['questions'][j]
                        ans = ds['answers'][j]
                        llm_ans = ds['llm_answer'][j]
                        if llm_ans == "" or llm_ans is None:
                            llm_ans = "No answer provided."

                        score = api(
                            [{'role': 'user', 'content': prompt_SA_EVAL.format(q=q, eg_ans=ans, llm_ans=llm_ans)}]
                        )
                        try:
                            score = int(score)
                        except ValueError:
                            print('err: ', score)
                            score = 0
                        if score == 1:
                            cate_ans = corr.get(cate, 0)
                            corr[cate] = cate_ans + 1
                        else:
                            cate_ans = corr.get(cate, 0)
                            corr[cate] = cate_ans
                        tmp.append(score)
                ds['retrived_answer'] = tmp

            acc = {k: v / cnt[k] if cnt[k] != 0 else 1 for k, v in corr.items()}
            res['count'] = cnt
            res['total_count'] = sum(cnt.values())
            overall_acc = np.sum(np.array(list(acc.values())) * np.array(list(cnt.values()))) / sum(cnt.values())
            res[llm_name][f'{shot_type}{prompt_r}'] = {
                'acc': acc,
                'overall_acc': overall_acc
            }
            json.dump(llm_res, open(f'res/{self.seed}{self.osd}/ctx/{f}', 'w'), indent=4)
        json.dump(res, open(ctx_save_path, 'w'), indent=4)


if __name__ == '__main__':
    eval_mg = Evaluation(41, False)
    eval_mg_oc = Evaluation(41, True)

    eval_mg.evaluate_mc()
    eval_mg.evaluate_sa()
    eval_mg.evaluate_sc()
    eval_mg.evaluate_sa_unique()

    eval_mg_oc.evaluate_mc()
    eval_mg_oc.evaluate_sa()
    eval_mg_oc.evaluate_sc()
    eval_mg_oc.evaluate_sa_unique()

evaluate.py
LinxinS97/NLPBench
from .sys_prompt import SYS_PROMPT_MAPPING
from .qa_prompt import PROMPT_MAPPING, SHORTANS_ZS_COT_ST1, CTX
from .tot_prompt import SOLUTION_PROMPT, STATE_PROMPT, VOTE_PROMPT
prompts/__init__.py
LinxinS97/NLPBench
MULTICHOICE_STD_ZS = '''
Answer the final multiple choice question. 
Your output must be only numbers spliting by comma (e.g., 0,1,...) with no descriptions.

Example Input:
ChatGPT is created by which of the following companies?
0: Google
1: Meta
2: Microsoft
4: Amazon
3: OpenAI

Example Output:
3

Example Input:
This is the input question, choose the correct answer.
0: Correct answer
1: Option 2
2: Correct answer
3: Option 4

Example Output:
0,2

Example Input:
GPT-4 is created by OpenAI.
0: True
1: False

Example Output:
0

Input (You need to answer this question): 
{input}

Output: 
'''


MULTICHOICE_STD_ZS_COT = '''
Answer the final multiple choice question. Your output must be only numbers spliting by comma (e.g., 0,1,...) with no descriptions.

Example Input:
ChatGPT is created by which of the following companies?
0: Google
1: Meta
2: Microsoft
4: Amazon
3: OpenAI

Example Thought:
ChatGPT is a large-scale transformer-based language model, created by OpenAI at 2022.

Example Output:
3

Example Input:
This is the input question, choose the correct answer.
0: Correct answer
1: Option 2
2: Correct answer
3: Option 4

Example Thought:
This is a multiple choice question, the "correct answer" appears at index 0 and 2.

Example Output:
0,2

Example Input:
GPT-4 is created by OpenAI.
0: True
1: False

Example Thought:
GPT-4 is created by OpenAI at 2022.

Example Output:
0

Input (You need to answer this question): 
{input}

Output: 
'''


MULTICHOICE_STD_FS = '''
Answer the final multiple choice question. Your output must be only numbers spliting by comma (e.g., 0,1,...) with no descriptions.

Example Input:
What is the main challenge(s) of NLP?
0: Handling Ambiguity of Sentences
1: Handling Tokenization
2: Handling POS-Tagging
3: All of the mentioned

Example Output:
0

Example Input:
What is the field of Natural Language Processing (NLP)>
0: Computer Science
1: Artificial Intelligence
2: Linguistics
3: All of the mentioned

Example Output:
3

Example Input:
Choose from the following areas where NLP can be useful.
0: Automatic Text Summarization
1: Automatic Question-Answering Systems
2: Information Retrieval
3: All of the mentioned

Example Output:
3

Input (You need to answer this question): 
{input}

Output: 
'''


MULTICHOICE_STD_FS_COT = '''
Answer the final multiple choice question. Your output must be only numbers spliting by comma (e.g., 0,1,...) with no descriptions.

Example Input:
What is the main challenge(s) of NLP?
0: Handling Ambiguity of Sentences
1: Handling Tokenization
2: Handling POS-Tagging
3: All of the mentioned

Example Thought:
There are enormous ambiguity exists when processing natural language.

Example Output:
0

Example Input:
What is Machine Translation?
0: Converts one human language to another
1: Converts human language to machine language
2: Converts any human language to English
3: Converts Machine language to human language

Example Thought:
The best known example of machine translation is google translator, which help people to translate one language to another.

Example Output:
0

Example Input:
What is Coreference Resolution?
0: Anaphora Resolution
1: Given a sentence or larger chunk of text, determine which words (“mentions”) refer to the same objects (“entities”)
2: All of the mentioned
3: None of the mentioned

Example Thought:
Anaphora resolution is a specific type of coreference resolution.

Example Output:
1

Input (You need to answer this question): 
{input}

Output: 
'''


SHORTANS_ZS = '''
Answer the following short answer question. Your answer should be no more then 150 words.

Input (You need to answer this question): 
{input}

Output: 
'''


SHORTANS_ZS_COT_ST1 = '''
Answer the following short answer question. Your answer should be no more then 150 words.

Input (You need to answer this question): 
{input}

Let's think step by step! Output your thought of the question first:
'''


SHORTANS_ZS_COT = '''
Answer the following short answer question. Your answer should be no more then 150 words.

Input (You need to answer this question): 
{input}

Your thought:
{thought}

Output: 
'''


SHORTANS_FS = '''
Answer the following short answer question. Your answer should be no more then 150 words.

Example Input:
Order the following syntactic features in decreasing order by salience (according to the Lappin/Leass algorithm for anaphora resolution): direct object (accusative), indirect object, subject, recency.

Example Output:
From strongest to weakest:
1. Recency
2. Subject
3. Direct object
4. Indirect object

Example Input: 
List any two real-life applications of Natural Language Processing.

Example Output:
1. Google Translate.
2. ChatGPT.

Example Input:
What are stop words?

Example Output:
Stop words are said to be useless data for a search engine. Words such as articles, prepositions, etc. are considered stop words. 
There are stop words such as was, were, is, am, the, a, an, how, why, and many more.

Input (You need to answer this question): 
{input}

Output: 
'''


SHORTANS_FS_COT = '''
Answer the following short answer question. Your answer should be no more then 150 words.

Example Input:
Order the following syntactic features in decreasing order by salience (according to the Lappin/Leass algorithm for anaphora resolution): direct object (accusative), indirect object, subject, recency.

Example Thought:
The following is the order of the given syntactic features in decreasing order by salience (according to the Lappin/Leass algorithm for anaphora resolution):
Direct object (accusative): The least salient antecedent is the direct object of the sentence.
Indirect object: The third most salient antecedent is the indirect object of the sentence.
Subject: The next most salient antecedent is the subject of the sentence.
Recency: The most salient antecedent is the noun phrase that occurs most recently in the discourse.

Example Output:
From strongest to weakest:
1. Recency
2. Subject
3. Direct object
4. Indirect object

Example Input: 
List any two real-life applications of Natural Language Processing.

Example Thought:
Natural Language Processing (NLP) has a wide range of applications across various industries due to its ability to understand and generate human language.
There are a lot of applications of NLP in deep translation (like google translate), chatbot (like ChatGPT) and many more.

Example Output:
1. Google Translate.
2. ChatGPT.

Example Input:
What are stop words?

Example Thought:
Stop words are said to be useless data for a search engine. 
There are stop words such as was, were, is, am, the, a, an, how, why, and many more.

Example Output:
Words such as articles, prepositions, etc. are considered stop words. 

Input (You need to answer this question): 
{input}
'''


MATH_ZS = '''
Answer the following math question. Your answer should be a number, a list of numbers, or a LaTeX expression.

Example Input:
1 + 1

Example Output:
2

Example Input:
$\\frac{{1}}{{2}} + \\frac{{1}}{{3}}$

Example Output:
\\frac{{5}}{{6}}

Example Input:
$f(x) = 4x^2 + 3y$
Solve the $\\frac{{\\partial f(x)}}{{\\partial x}}$

Example Output:
8x

Input (You need to answer this question):
{input}

Output:
'''


MATH_ZS_COT = '''
Answer the following math question. Your answer should be a number, a list of numbers, or a LaTeX expression.

Example Input:
1 + 1

Example Thought:
1 + 1 = 2

Example Output:
2

Example Input:
$\\frac{{1}}{{2}} + \\frac{{1}}{{3}}$

Example Thought:
$\\frac{{1}}{{2}} + \\frac{{1}}{{3}} = \\frac{{5}}{{6}}$

Example Output:
\\frac{{5}}{{6}}

Example Input:
$f(x) = 4x^2 + 3y$
Solve the $\\frac{{\\partial f(x)}}{{\\partial x}}$

Example Thought:
$\\frac{{\\partial f(x)}}{{\\partial x}} = 2\\times 4x + 0$

Example Output:
8x

Example Input (You need to answer this question):
{input}

Output:
'''


MATH_FS = '''
Answer the following math question. Your answer should be a number, a list of numbers, or a LaTeX expression.

Example Input:
$\\frac{{1}}{{2}} + \\frac{{1}}{{3}}$

Example Output:
\\frac{{5}}{{6}}

Example Input:
$f(x) = 4x^2 + 3y$
Solve the $\\frac{{\\partial f(x)}}{{\\partial x}}$

Example Output:
8x

Example Input:
Consider the following bilingual (Spanish-English) corpus.
gato blanco
white cat

el gato
the cat

Considering only the following three alignment types:
1. || {{1,2}}
2. X {{2,1}}

Start with a uniform distribution for $t(white|gato)$, $t(white|blanco)$, $t(cat|gato)$, $t(cat|blanco)$, $t(cat|el)$, $t(the|el)$, and $t(the|gato)$. 
Show the values of them after two iterations of the EM algorithm.

Example Output:
\\frac{{1}}{{6}},\\frac{{2}}{{3}},\\frac{{2}}{{3}},\\frac{{1}}{{3}},\\frac{{1}}{{3}},\\frac{{2}}{{3}},\\frac{{1}}{{6}}

Input:
{input}

Output (You need to answer this question):
'''


MATH_FS_COT = '''
Answer the following math question. Your answer should be a number, a list of numbers, or a LaTeX expression.

Example Input:
$\\frac{{1}}{{2}} + \\frac{{1}}{{3}}$

Example Thought:
1 + 1 = 2

Example Output:
\\frac{{5}}{{6}}

Example Input:
$f(x) = 4x^2 + 3y$
Solve the $\\frac{{\\partial f(x)}}{{\\partial x}}$

Example Thought:
$\\frac{{\\partial f(x)}}{{\\partial x}} = 2\\times 4x + 0$

Example Output:
8x

Example Input:
Consider the following bilingual (Spanish-English) corpus.
gato blanco
white cat

el gato
the cat

Considering only the following three alignment types:
1. || {{1,2}}
2. X {{2,1}}

Start with a uniform distribution for $t(white|gato)$, $t(white|blanco)$, $t(cat|gato)$, $t(cat|blanco)$, $t(cat|el)$, $t(the|el)$, and $t(the|gato)$. 
Show the values of them after two iterations of the EM algorithm.

Example Thought:
Initialization:
$t(white|gato)=\\frac{{1}}{{2}}$
$t(white|blanco)=\\frac{{1}}{{2}}$
$t(cat|gato)=\\frac{{1}}{{3}}$
$t(cat|blanco)=\\frac{{1}}{{3}}$
$t(cat|el)=\\frac{{1}}{{3}}$
$t(the|el)=\\frac{{1}}{{2}}$
$t(the|gato)=\\frac{{1}}{{2}}$

$\\textbf{{First Iteration:}}$

gato blanco
white cat

For alignment type X:
$p(a,f|e)=\\frac{{1}}{{2}}\\times \\frac{{1}}{{3}}=\\frac{{1}}{{6}}$
$p(a|e,f)=\\frac{{1}}{{2}}$

For alignment type ||:
$p(a, f|e) = \\frac{{1}}{{2}}\\times\\frac{{1}}{{3}} = \\frac{{1}}{{6}}$
$p(a|e, f) = \\frac{{1}}{{2}}$

el gato
the cat

For alignment type X:
$p(a,f|e)=\\frac{{1}}{{2}}\\times \\frac{{1}}{{3}}=\\frac{{1}}{{6}}$
$p(a|e,f)=\\frac{{1}}{{2}}$

For alignment type ||:
$p(a, f|e) = \\frac{{1}}{{2}}\\times\\frac{{1}}{{3}} = \\frac{{1}}{{6}}$
$p(a|e, f) = \\frac{{1}}{{2}}$

Fractional Counts:
$t(white|gato)=\\frac{{1}}{{2}}$
$t(white|blanco)=\\frac{{1}}{{2}}$
$t(cat|gato)=\\frac{{1}}{{2}}+\\frac{{1}}{{2}}=1$
$t(cat|blanco)=\\frac{{1}}{{2}}$
$t(cat|el)=\\frac{{1}}{{2}}$
$t(the|el)=\\frac{{1}}{{2}}$
$t(the|gato)=\\frac{{1}}{{2}}$

Normalize to get updated parameters:
$t(white|gato)=\\frac{{1}}{{4}}$
$t(white|blanco)=\\frac{{1}}{{2}}$
$t(cat|gato)=\\frac{{1}}{{2}}$
$t(cat|blanco)=\\frac{{1}}{{2}}$
$t(cat|el)=\\frac{{1}}{{2}}$
$t(the|el)=\\frac{{1}}{{2}}$
$t(the|gato)=\\frac{{1}}{{4}}$

$\\textbf{{Second Iteration:}}$
gato blanco\nwhite cat

For alignment type X:
$p(a,f|e)=\\frac{{1}}{{2}}\\times \\frac{{1}}{{2}}=\\frac{{1}}{{4}}$
$p(a|e,f)=\\frac{{2}}{{3}}$

For alignment type ||:
$p(a, f|e) = \\frac{{1}}{{4}}\\times\\frac{{1}}{{2}} = \\frac{{1}}{{8}}$
$p(a|e, f) = \\frac{{1}}{{3}}$

el gato
the cat

For alignment type X:
$p(a,f|e)=\\frac{{1}}{{4}}\\times \\frac{{1}}{{2}}=\\frac{{1}}{{8}}$
$p(a|e,f)=\\frac{{1}}{{3}}$

For alignment type ||:
$p(a, f|e) = \\frac{{1}}{{2}}\\times\\frac{{1}}{{2}} = \\frac{{1}}{{4}}$
$p(a|e, f) = \\frac{{2}}{{3}}$

Fractional Counts:
$t(white|gato)=\\frac{{1}}{{3}}$
$t(white|blanco)=\\frac{{2}}{{3}}$
$t(cat|gato)=\\frac{{2}}{{3}}+\\frac{{2}}{{3}}=\\frac{{4}}{{3}}$
$t(cat|blanco)=\\frac{{1}}{{3}}$
$t(cat|el)=\\frac{{1}}{{3}}$
$t(the|el)=\\frac{{2}}{{3}}$
$t(the|gato)=\\frac{{1}}{{3}}$

Finally, you can normalize to get updated parameters.

Example Output:
\\frac{{1}}{{6}},\\frac{{2}}{{3}},\\frac{{2}}{{3}},\\frac{{1}}{{3}},\\frac{{1}}{{3}},\\frac{{2}}{{3}},\\frac{{1}}{{6}}

Input (You need to answer this question):
{input}

Output:
'''


CTX = '''The context of the following questions is:
{context}
'''

PROMPT_MAPPING = {
    0: {
        'zero-shot': MULTICHOICE_STD_ZS,
        'few-shot': MULTICHOICE_STD_FS,
        'zero-shot_cot': MULTICHOICE_STD_ZS_COT,
        'few-shot_cot': MULTICHOICE_STD_FS_COT,
        'zero-shot_tot': MULTICHOICE_STD_ZS_COT,
        'few-shot_tot': MULTICHOICE_STD_FS_COT
    },
    1: {
        'zero-shot': SHORTANS_ZS,
        'few-shot': SHORTANS_FS,
        'zero-shot_cot': SHORTANS_ZS_COT,
        'few-shot_cot': SHORTANS_FS_COT,
        'zero-shot_tot': SHORTANS_ZS_COT,
        'few-shot_tot': SHORTANS_FS_COT
    },
    2: {
        'zero-shot': MATH_ZS,
        'few-shot': MATH_FS,
        'zero-shot_cot': MATH_ZS_COT,
        'few-shot_cot': MATH_FS_COT,
        'zero-shot_tot': MATH_ZS_COT,
        'few-shot_tot': MATH_FS_COT
    }
}

prompts/qa_prompt.py
LinxinS97/NLPBench
SYS_PROMPT_MULTICHOICE = """Answer the multiple choice questions in the field of Natural Language Processing (NLP).
You should select a single (or a couple of) answer(s) from the given options. 
"""

SYS_PROMPT_SHORTANS = """Answer the short answer questions in the field of Natural Language Processing (NLP).
You should provide a concise short answer for the question.
All the math symbols in your answer must be converted to LaTeX format (e.g., \\pi, \\sqrt{{2}}). 
"""

SYS_PROMPT_MATH = """Answer the mathematics questions in the field of Natural Language Processing (NLP).
You should provide a (or a couple of) number(s) or LaTeX expression(s). 
When the answer is a fraction, please use \\frac{{}}{{}} to express it (e.g., \\frac{{1}}{{2}}).
When the answer is a vector or matrix, please use \\begin{{bmatrix}} \\end{{bmatrix}} to express it (e.g., \\begin{{bmatrix}} 1 & 2 \\\\ 3 & 4 \\end{{bmatrix}}).
You should not trun \\pi, \\sqrt{{2}}... to a decimals, but use the original format.
"""

SYS_PROMPT_MAPPING = {
    0: SYS_PROMPT_MULTICHOICE,
    1: SYS_PROMPT_SHORTANS,
    2: SYS_PROMPT_MATH
}

prompts/sys_prompt.py
LinxinS97/NLPBench
SOLUTION_PROMPT = """You're an Tree-of-Thoughts, an superintelligent AI model devoted to help Human by any means necessary. 
Your purpose is to generate a series of solutions to comply with the user's instructions, you must generate solutions on the basis of determining the most reliable solution in the shortest amount of time, while taking rejected solutions into account and learning from them. 
Considering the reasoning provided:\n\n
###'{state_text}'\n\n###
Devise the best possible solution for the task: {initial_prompt}, Here are evaluated solutions that were rejected: 
###{rejected_solutions}###, 
complete the {initial_prompt} without making the same mistakes you did with the evaluated rejected solutions. 
Be simple. Be direct. Provide intuitive solutions as soon as you think of them."""

STATE_PROMPT = """Given a question: 
'{initial_prompt}', value the following past solutions as a float number between 0 and 1.\n
If the past solution is not directly and concretely in achieving the goal, give it a low score.
Past solutions:\n\n
'{state_text}'
"""

VOTE_PROMPT = """Given the following states of reasoning, vote for the best state utilizing an scalar value 1-10:\n{states_text}\n\nVote, on the probability of this state of reasoning achieveing {initial_prompt} and become very pessimistic very NOTHING ELSE"""
prompts/tot_prompt.py
LinxinS97/NLPBench
from solve import Solver
from sacred import Experiment

ex = Experiment("LLM")


@ex.config
def config():
    model_name = 'gpt-3.5-turbo'
    config = 'zero-shot'
    ctx = False


@ex.automain
def run(
        model_name,
        seed,
        shot_type,
        prompt_r,
        sys,
        ctx,
        max_tokens,
        self_consistency,
        self_reflection
):
    if ctx:
        data_path = "data/w_ctx.json"
    else:
        data_path = "data/wo_ctx.json"

    solver = Solver(
        data_path=data_path,
        model_name=model_name,
        seed=seed,
        shot_type=shot_type,
        prompt_r=prompt_r,
        sys=sys,
        ctx=ctx,
        max_tokens=max_tokens,
        self_consistency=self_consistency,
        self_reflection=self_reflection,
    )

    solver.run()

run.py
LinxinS97/NLPBench
import os
import json
import traceback
from typing import Optional
from collections import Counter
from utils import ApiManager, PromptGenerator, tot_solver


class Solver:
    def __init__(
            self,
            data_path: str,
            model_name: str,
            seed: int,
            shot_type: str,
            prompt_r: str,
            sys: bool,
            max_tokens: int,
            ctx: bool,
            self_consistency: bool = False,
            self_reflection: bool = False,
    ):
        self.data_path = data_path
        self.model_name_save = model_name.split('/')[-1]
        ctx_path = 'ctx' if ctx else 'no_ctx'
        self.save_path = (f'res/{str(seed)}/{ctx_path}/{shot_type}'
                          f'{"-sc" if self_consistency else ""}'
                          f'_{self.model_name_save}'
                          f'{("_" + prompt_r) if prompt_r is not None else ""}'
                          f'{"_sys" if sys else ""}.json')

        root = self.save_path.rsplit('/', 1)[0]
        if not os.path.exists(root):
            os.makedirs(root)

        self.ctx = ctx

        ### Init api manager and prompt generator
        self.llm = ApiManager(
            model_name=model_name,
            seed=seed,
            default_max_tokens=max_tokens,
        )

        self.prompt = PromptGenerator(
            api=self.llm,
            system_prompt=sys,
            shot_type=shot_type,
            prompt_reinforcement=prompt_r,
        )

        self.self_consistency = self_consistency
        self.self_reflection = self_reflection

    @staticmethod
    def load_cache(save_path: str):
        skip_cnt = 0
        res_cache = []
        if os.path.exists(save_path):
            res_cache = json.load(open(save_path, 'r'))
            skip_cnt = len(res_cache)

        return res_cache, skip_cnt

    @staticmethod
    def majority_voting(resp_list: list):
        res = [r.replace(' ', '').replace("'", '').replace('"', '') for r in resp_list]
        res = Counter(res).most_common(1)[0][0]
        return res

    def solve_ctx(self):
        ### Load cache (if palm, else init result)
        if self.llm.api_type == 'palm':
            res, skip_cnt = self.load_cache(self.save_path)
        else:
            res = []
            skip_cnt = 0

        for data in json.load(open(self.data_path, 'r')):
            if skip_cnt > 0:
                skip_cnt -= 1
                continue
            history = None
            resps = []
            for idx, q in enumerate(data['questions']):
                q_type = data['type'][idx]
                q_opt = [] if data['options'][idx] is None else data['options'][idx]
                ctx = None if data['context'] == '-1' else data['context']
                messages = self.prompt.generate(q=q, q_type=q_type, q_opt=q_opt, ctx=ctx, history=history)

                if self.self_consistency is True and q_type == 0:
                    try:
                        resp = self.llm(messages, choices=3)
                        if resp is None:
                            resp = ""
                    except Exception as e:
                        print(e)
                        resp = ""
                    if resp != "":
                        resp = self.majority_voting(resp)
                else:
                    try:
                        resp = self.llm(messages)
                        if resp is None:
                            resp = ""
                    except Exception as e:
                        print(e)
                        resp = ""

                if self.llm.api_type == 'palm':
                    history = messages
                    history['prompt'] += '\n' + resp

                elif self.llm.api_type == 'oai':
                    history = messages
                    history[-1]['content'] += resp
                resps.append(resp)
                print(f"\n{self.model_name_save}: {resp}")

            res.append({
                **data,
                'prompt': history,
                'llm_answer': resps
            })
            if self.llm.api_type == 'palm':
                json.dump(res, open(self.save_path, 'w'), indent=4)

        json.dump(res, open(self.save_path, 'w'), indent=4)

    def solve_no_ctx(self):
        ### Load cache (if palm, else init result)
        if self.llm.api_type == 'palm':
            res, skip_cnt = self.load_cache(self.save_path)
        else:
            res = []
            skip_cnt = 0

        for data in json.load(open(self.data_path, 'r')):
            if skip_cnt > 0:
                skip_cnt -= 1
                continue
            q = data['question']
            q_type = data['type']
            q_opt = data.get('options', [])

            messages = self.prompt.generate(q=q, q_type=q_type, q_opt=q_opt)
            if q_type == 0 and self.self_consistency is True:
                try:
                    resp = self.llm(messages, choices=3)
                except Exception as e:
                    traceback.print_exc()
                    messages = []
                    resp = ""
                if resp != "":
                    resp = self.majority_voting(resp)
            else:
                try:
                    resp = self.llm(messages)
                except Exception as e:
                    traceback.print_exc()
                    messages = []
                    resp = ""
            print(f"\n {self.model_name_save}: {resp}")
            res.append({
                **data,
                'prompt': messages,
                'llm_answer': resp
            })
            if self.llm.api_type == 'palm':
                json.dump(res, open(self.save_path, 'w'), indent=4)

        json.dump(res, open(self.save_path, 'w'), indent=4)

    def solve_tot_no_ctx(
            self,
            evaluation_strategy: Optional[str] = "vote",  # value or vote
            num_thoughts: Optional[int] = 1,
            max_steps: Optional[int] = 3,
            max_states: Optional[int] = 4,
            pruning_threshold: Optional[float] = 0.5,
    ):
        ### Load cache (if palm, else init result)
        if self.llm.api_type == 'palm':
            res, skip_cnt = self.load_cache(self.save_path)
        else:
            res = []
            skip_cnt = 0

        for data in json.load(open(self.data_path, 'r')):
            if skip_cnt > 0:
                skip_cnt -= 1
                continue
            try:
                resp, messages = tot_solver(
                    generator=self.prompt,
                    api=self.llm,
                    q=data['question'],
                    q_type=data['type'],
                    q_opt=data.get('options', []),
                    evaluation_strategy=evaluation_strategy,
                    num_thoughts=num_thoughts,
                    max_steps=max_steps,
                    max_states=max_states,
                    pruning_threshold=pruning_threshold,
                )
            except Exception as e:
                traceback.print_exc()
                messages = []
                resp = ""

            print(f"\n {self.model_name_save}: {resp}")
            res.append({
                **data,
                'messages': messages,
                'llm_answer': resp
            })
            if self.llm.api_type == 'palm':
                out_str = json.dumps(res[-1])
                with open(self.save_path, 'a') as f:
                    f.write(out_str + '\n')

        json.dump(res, open(self.save_path, 'w'), indent=4)

    def solve_ctx_tot(self):
        # TODO: implement
        pass

    def run(self):
        if self.ctx:
            self.solve_ctx()
        else:
            if self.prompt.prompt_reinf == 'tot':
                self.solve_tot_no_ctx()
            else:
                self.solve_no_ctx()

solve/solver.py
LinxinS97/NLPBench
from .arg_parser import parser
from .utils import PromptGenerator, ApiManager
from .tot_solver import tot_solver
utils/__init__.py
LinxinS97/NLPBench
import argparse


parser = argparse.ArgumentParser()
parser.add_argument("--config", type=str, default='zero-shot')
utils/arg_parser.py
LinxinS97/NLPBench
import logging
import concurrent.futures
import os
import json
import numpy as np
from typing import List, Dict, Any, Optional, Union
from utils import ApiManager, PromptGenerator
from prompts import SOLUTION_PROMPT, STATE_PROMPT, VOTE_PROMPT


logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)



class ToTModel:

    api: ApiManager
    evaluation_strategy: Optional[str] = "value"

    def __init__(self, api: ApiManager, evaluation_strategy: Optional[str] = "value"):
        self.api = api
        self.evaluation_strategy = evaluation_strategy


    def message_handler(self, message: str):
        if self.api.api_type == 'palm':
            return {
                'prompt': message,
            }
        elif self.api.api_type == 'oai':
            return [{
                'role': 'user',
                'content': message,
            }]


    def generate_thoughts(self, 
                          state: Union[str, list], 
                          k: int, 
                          initial_prompt: Union[str, dict], 
                          rejected_solutions: Optional[str]=None):
        
        if type(state) == str:
            state_text = state
        else:
            state_text = '\n'.join(state)

        prompt = SOLUTION_PROMPT.format(state_text=state_text,
                                        initial_prompt=initial_prompt, 
                                        rejected_solutions=rejected_solutions)
        
        prompt = self.message_handler(prompt)
        thoughts = [self.api(prompt) for _ in range(k)]
        return thoughts

        
    def generate_solution(self, 
                          initial_prompt: str, 
                          state: Union[str, list], 
                          rejected_solutions: Optional[str]=None):
        try:
            if type(state) == str:
                state_text = state
            else:
                state_text = '\n'.join(state)
            
            prompt = SOLUTION_PROMPT.format(state_text=state_text,
                                            initial_prompt=initial_prompt, 
                                            rejected_solutions=rejected_solutions)
            prompt = self.message_handler(prompt)
            answer = self.api(prompt)
            return answer
        
        except Exception as e:
            logger.error(f"Error in generate_solutions: {e}")
            return None


    def evaluate_states(self, 
                        states: List, 
                        initial_prompt: str):
        if not states:
            return {}

        if self.evaluation_strategy == 'value':
            state_values = {}
            for state in states:
                if type(state) == str:
                    state_text = state
                else:
                    state_text = '\n'.join(state)

                prompt = STATE_PROMPT.format(initial_prompt=initial_prompt, state_text=state_text)
                prompt = self.message_handler(prompt)
                try:
                    value_text = self.api(prompt, 10)
                    value = float(value_text)
                    print(f"Evaluated Thought Value: {value}")
                except ValueError:
                    value = 0  # Assign a default value if the conversion fails
                state_values[state] = value
            return state_values

        elif self.evaluation_strategy == 'vote':
            states_text = '\n'.join([' '.join(state) for state in states])

            prompt = VOTE_PROMPT.format(states_text=states_text, initial_prompt=initial_prompt)
            prompt = self.message_handler(prompt)
            best_state_text = self.api(prompt, 50)

            best_state = tuple(best_state_text.split())

            return {state: 1 if state == best_state else 0 for state in states}
        
        else:
            raise ValueError("Invalid evaluation strategy. Choose 'value' or 'vote'.")



class TreeofThoughts:

    model: ToTModel
    best_state = None
    best_value = float("-inf")
    history = [] #added line initalize history
    tree: Dict[str, Dict[str, Union[float, Dict[str, Any]]]] = {
        "nodes": {},
    }

    def __init__(self, model: ToTModel):
        self.model = model

    def save_tree_to_json(self, file_name):
        os.makedirs(os.path.dirname(file_name), exist_ok=True)
        with open(file_name, 'w') as json_file:
            json.dump(self.tree, json_file, indent=4)

    def logNewState(self, state, evaluation):
        if not (type(state) == str):
            state = " | ".join(state)
        if state in self.tree['nodes']:
            self.tree['nodes'][state]['thoughts'].append(evaluation)
        else:
            self.tree['nodes'][state] = {'thoughts': [evaluation]}

    def adjust_pruning_threshold_precentile(self, evaluated_thoughts, percentile):
        values = np.array(list(evaluated_thoughts.values()))
        if values.size == 0:
            return 0 
        return max(np.percentile(values, percentile), 0.1)
    

    def adjust_pruning_threshold_moving_average(self, evaluated_thoughts, window_size):
        values = list(evaluated_thoughts.values())
        if len(values) < window_size:
            return np.mean(values) if values else 0
        else:
            return max(np.mean(values[-window_size:]), 0.1)



class TreeofThoughtsBFS(TreeofThoughts):
    def solve(
        self, 
        initial_prompt, 
        num_thoughts, 
        max_steps, 
        max_states, 
        pruning_threshold=0.5
    ):
        current_states = [initial_prompt]
        state_values = {}
        dynamic_pruning_threshold = pruning_threshold

        try:
            with concurrent.futures.ThreadPoolExecutor() as executor:
                for _ in range(1, max_steps + 1):
                    selected_states = []
                    for state in current_states:
                        thoughts = self.model.generate_thoughts(state, num_thoughts, initial_prompt)
                        futures = [executor.submit(self.model.evaluate_states, {thought: 0}, initial_prompt) for thought in thoughts]
                        concurrent.futures.wait(futures)
                        evaluated_thoughts = {thought: list(fut.result().values()) for thought, fut in zip(thoughts, futures)}

                        if evaluated_thoughts:  # only adjust if you have evaluated thoughts
                            dynamic_pruning_threshold = self.adjust_pruning_threshold_moving_average(evaluated_thoughts, 5)

                        for thought, value in evaluated_thoughts.items():
                            flattened_state = (state, thought) if isinstance(state, str) else (*state, thought)
                            selected_states.append((flattened_state, value))

                        selected_states.sort(key=lambda x: x[1], reverse=True)
                        selected_states = selected_states[:max_states]  # Select only the top states

                        for state, value in selected_states:
                            if value >= dynamic_pruning_threshold:
                                state_values[state] = value
                                self.logNewState(state, value)
                                logger.debug(f"State Values: {state_values}")

            if state_values:
                highest_rated_solution = max(state_values.items(), key=lambda x: x[1])
                highest_rated_state = highest_rated_solution[0]
                solution = self.model.generate_solution(initial_prompt, highest_rated_state)

                return solution, highest_rated_state

            else:
                return None

        except Exception as e:
            logger.error(f"Error in tot_bfs: {e}")
            return None



def tot_solver(
    generator: PromptGenerator,
    api: ApiManager,
    q: str,
    q_type: int,
    q_opt: Optional[List[str]] = [],
    evaluation_strategy: Optional[str] = "vote",  # value or vote
    num_thoughts: Optional[int] = 1,
    max_steps: Optional[int] = 3,
    max_states: Optional[int] = 4,
    pruning_threshold: Optional[float] = 0.5,
):
    tot_model = ToTModel(api, evaluation_strategy)
    tot = TreeofThoughtsBFS(tot_model)
    initial_prompt = generator.generate(q, q_type, q_opt)

    if generator.api_type == 'palm':
        input_prompt = initial_prompt['prompt']
    else:
        input_prompt = initial_prompt[-1]['content']

    resp = tot.solve(
        initial_prompt=input_prompt,
        num_thoughts=num_thoughts,
        pruning_threshold=pruning_threshold,
        max_steps=max_steps,
        max_states=max_states
    )

    return resp


utils/tot_solver.py
LinxinS97/NLPBench
import os
import time

import google.generativeai as palm
from flaml import autogen
from typing import List, Optional, Union
from prompts import SYS_PROMPT_MAPPING, PROMPT_MAPPING, SHORTANS_ZS_COT_ST1


palm.configure(api_key=os.environ.get("PALM_API_KEY"))


palm_llm_config = {
    "models/text-bison-001": {
        "model": "models/text-bison-001",
    },
}


oai_llm_config = {
    "meta-llama/Llama-2-7b-chat-hf": {
        "model": "meta-llama/Llama-2-7b-chat-hf",
        "api_key": "empty",
        "api_base": "http://127.0.0.1:8000/v1",
    },
    "meta-llama/Llama-2-13b-chat-hf": {
        "model": "meta-llama/Llama-2-13b-chat-hf",
        "api_key": "empty",
        "api_base": "http://127.0.0.1:8000/v1",
    },
    "meta-llama/Llama-2-70b-chat-hf": {
        "model": "meta-llama/Llama-2-70b-chat-hf",
        "api_key": "empty",
        "api_base": "http://127.0.0.1:8000/v1",
    },
     "sambanovasystems/BLOOMChat-176B-v1": {
        "model":  "sambanovasystems/BLOOMChat-176B-v1",
        "api_key": "empty",
        "api_base": "http://127.0.0.1:8000/v1",
    },
    "gpt-3.5-turbo": {
        "model": "gpt-3.5-turbo",
        "api_key": os.environ.get("OPENAI_API_KEY"),
        "api_base": "https://api.openai.com/v1",
        "api_version": None,
    },
    "gpt-4": {
        "model": "gpt-4",
        "api_key": os.environ.get("OPENAI_API_KEY"),
        "api_base": "https://api.openai.com/v1",
        "api_version": None,
    },
}


class Completion(autogen.ChatCompletion):
    request_timeout = 300
    retry_time = 20


class ApiManager:
    model_name: str
    seed: int
    default_max_tokens: int

    def __init__(
            self,
            model_name: str,
            seed: int,
            temperature: Optional[float] = 1.0,
            default_max_tokens: Optional[int] = 945
    ):
        self.model_name = model_name
        self.seed = seed
        self.max_tokens = default_max_tokens
        self.temperature = temperature

    @property
    def api_type(self):
        return 'palm' if self.model_name in palm_llm_config.keys() else 'oai'

    def get_api_type(self):
        return self.api_type

    def palm_api(self, messages: str, choices: int = 1):

        config = {
            "max_output_tokens": self.max_tokens,
            "temperature": self.temperature,
            **messages,  # prompt
            **palm_llm_config[self.model_name]
        }
        while True:
            try:
                if choices == 1:
                    result = palm.generate_text(**config).result
                else:
                    result = [palm.generate_text(**config).result for _ in range(choices)]
            except Exception:
                print('Retrying in 20s...')
                time.sleep(20)
                continue
            break
        return result

    def oai_api(self, messages: list, choices: int = 1):
        basic_config = {
            "api_type": "open_ai",
            "max_tokens": self.max_tokens,
            "temperature": self.temperature,
            "seed": self.seed,  # NOTE: the dialog will be cached for the same seed
            "n": choices,
            **oai_llm_config[self.model_name]
        }

        # https://microsoft.github.io/FLAML/docs/Use-Cases/Autogen#basic-concept
        completion = Completion.create(
            messages=messages,
            **basic_config
        )
        if choices == 1:
            return completion.choices[0].message.content
        else:
            return [c.message.content for c in completion.choices]

    def __call__(self, 
                 messages: Union[list, str],
                 temperature: Optional[float] = None,
                 max_tokens: Optional[int] = None,
                 choices: Optional[int] = 1):
        if max_tokens is not None:
            self.max_tokens = max_tokens
        if temperature is not None:
            self.temperature = temperature
        if self.model_name in palm_llm_config.keys():
            resp = self.palm_api(messages, choices)
            return resp

        elif self.model_name in oai_llm_config.keys():
            resp = self.oai_api(messages, choices)
            return resp

        else:
            raise NotImplementedError


class PromptGenerator:
    """
    Generate the prompt for the given question and example.

    Args:
        api: the api manager
        sys_prompt: system prompt
        shot_type: question type, 'zero-shot' or 'few-shot'
        prompt_reinf: whether to use the prompt reinforcement (none, cot or tot)
        self_consistency: whether to use the self-consistency
        self_reflection: whether to use the self-reflection
    """

    api: ApiManager
    sys_prompt: Optional[bool]
    shot_type: Optional[str]
    prompt_reinf: Optional[str]
    return_symbol: Optional[str] = '\n'

    def __init__(
            self,
            api: ApiManager,
            system_prompt: Optional[bool] = False,
            shot_type: Optional[str] = 'zero-shot',
            prompt_reinforcement: Optional[str] = None,
    ):
        self.api = api
        self.sys_prompt = system_prompt
        self.shot_type = shot_type
        self.prompt_reinf = prompt_reinforcement

    @property
    def api_type(self):
        return self.api.get_api_type()

    def oai_prompt_generator(
            self,
            q: str,
            q_type: int,
            q_opt: Optional[List[str]] = None,
            ctx: Optional[str] = None,
            history: Optional[List[dict]] = None,
    ) -> List[dict]:
        """
        Generate the prompt for the given question and example (Open AI version).

        Args:
            args: the user arguments
            q: question content
            q_type: question type
            q_opt: question options (multiple choice only)

        Return:
            messages: the generated prompt
        """
        messages = [] if history is None else history
        prompt_reinf = f'_{self.prompt_reinf}' if self.prompt_reinf is not None else ''
        prompt_type = f'{self.shot_type}{prompt_reinf}'

        if self.sys_prompt and history is None:
            messages += [{"role": "system", "content": SYS_PROMPT_MAPPING[q_type]}]

        if ctx is not None and history is None:
            messages += [{"role": "user", "content": ctx}]

        if q_type == 0:
            q = f"{q}\n{self.return_symbol.join([f'{i}: {opt}' for i, opt in enumerate(q_opt)])}"

        if q_type == 1 and prompt_type == 'zero-shot_cot':
            cot_st1 = [{"role": "user", "content": SHORTANS_ZS_COT_ST1.format(input=q)}]
            if self.sys_prompt:
                cot_st1 = [{"role": "system", "content": SYS_PROMPT_MAPPING[q_type]}] + cot_st1
            try:
                t = self.api(cot_st1, max_tokens=128)
            except Exception as e:
                print(e)
                t = ''
            messages += [{"role": "user", "content": PROMPT_MAPPING[q_type][prompt_type].format(input=q, thought=t)}]

            return messages
        try:
            messages += [{"role": "user", "content": PROMPT_MAPPING[q_type][prompt_type].format(input=q)}]
        except Exception as e:
            print(e)

        return messages

    def palm_prompt_generator(
            self,
            q: str,
            q_type: int,
            q_opt: Optional[List[str]] = [],
            ctx: Optional[str] = None,
            history: Optional[List[dict]] = None,
    ) -> dict:
        """
        Generate the prompt for the given question and example (PaLM version).

        Args:
            args: the user arguments
            q: question content
            q_type: question type
            q_opt: question options (multiple choice only)

        Return:
            messages: the generated prompt
        """
        messages = {'prompt': ''} if history is None else history
        prompt_reinf = f'_{self.prompt_reinf}' if self.prompt_reinf is not None else ''
        prompt_type = f'{self.shot_type}{prompt_reinf}'
        sp = SYS_PROMPT_MAPPING[q_type] + self.return_symbol

        if self.sys_prompt and history is None:
            messages['prompt'] += sp

        if ctx is not None and history is None:
            messages['prompt'] += ctx + self.return_symbol

        if q_type == 0:
            q = f"{q}\n{self.return_symbol.join([f'{i}: {opt}' for i, opt in enumerate(q_opt)])}"

        if q_type == 1 and prompt_type == 'zero-shot_cot':
            cot_st1 = {'prompt': SHORTANS_ZS_COT_ST1.format(input=q)}
            if self.sys_prompt:
                cot_st1['prompt'] = sp + cot_st1['prompt']
            try:
                t = self.api(cot_st1, max_tokens=128)
            except Exception as e:
                print(e)
                t = ''
            messages['prompt'] += PROMPT_MAPPING[q_type][prompt_type].format(input=q, thought=t)

            return messages
        try:
            messages['prompt'] += PROMPT_MAPPING[q_type][prompt_type].format(input=q)
        except Exception as e:
            print(e)
        return messages

    def generate(
            self,
            q: str = "",
            q_type: int = 0,
            q_opt: Optional[List[str]] = None,
            ctx: Optional[str] = None,
            history: Optional[List[dict]] = None,
    ):
        if self.api_type == 'oai':
            messages = self.oai_prompt_generator(q, q_type, q_opt, ctx, history)
        else:
            messages = self.palm_prompt_generator(q, q_type, q_opt, ctx, history)

        return messages

utils/utils.py
LinxinS97/NLPBench