Task2/promp_get.py

def get_prompt(ask_name, shot_type):
    import pandas as pd
    import random
    if ask_name == 'yes_no':
        # 设置判断数据路径
        data_path = './data/data.xlsx'
        sheet_train = 'Yes or No Train'
        sheet_test = 'Yes or No Test'   
        # 读取数据
        data_train = pd.read_excel(data_path, sheet_name=sheet_train)
        data_test = pd.read_excel(data_path, sheet_name=sheet_test)
        # 设置问题和文本
        question_train = data_train['Question']
        text_train = data_train['Text']
        answer_train = data_train['Answer']
        question_test = data_test['Question']
        text_test = data_test['Text']
        answer_test = data_test['Answer']
        # 设置prompt
        if shot_type == 'zero_shot':
            prompt = '''
请根据给定的文本回答问题，
给定文本：“''' + text_train + '''”'''+'\n'+'问题：“''' + question_train + '''”
请用Yes或No直接回答。
            '''
        elif shot_type == 'one_shot':
            # 随机设置一个例子
            random_index = random.randint(0, len(data_test)-1)
            example = '文本：' + str(text_test[random_index]) +'\n'+ '问题：' + str(question_test[random_index]) +'\n'+ '答案：' + str(answer_test[random_index])
            prompt = '''
请根据给定的文本回答问题，
            ''' + '\n' + '''例子如下：''' + '\n' + example + '''
            ''' + '\n' + '''给定文本：“''' + text_train + '''”'''+'\n'+'''问题：“''' + question_train + '”' + '''
请用Yes或No直接回答。
            ''' 
        elif shot_type == 'two_shot':
            # 随机设置两个例子
            random_index1 = random.randint(0, len(data_test)-1)
            random_index2 = random.randint(0, len(data_test)-1)
            while random_index2 == random_index1:
                random_index2 = random.randint(0, len(data_test)-1)
            example1 = '文本：' + str(text_test[random_index1]) + '\n'+ '问题：' + str(question_test[random_index1]) + '\n'+ '答案：' + str(answer_test[random_index1])
            example2 = '文本：' + str(text_test[random_index2]) + '\n'+ '问题：' + str(question_test[random_index2]) + '\n'+ '答案：' + str(answer_test[random_index2])
            prompt = '''
请根据给定的文本回答问题，
            ''' + '\n' + '''例子一如下：''' + '\n' + example1 + '''
            ''' + '\n' + '''例子二如下：''' + '\n' + example2 + '''
            ''' + '\n' + '''给定文本：“''' + text_train + '''”'''+'\n'+'''问题：“''' + question_train + '”' + '''
请用Yes或No直接回答。
            '''
        elif shot_type == 'three_shot':
            # 随机设置三个例子
            random_index1 = random.randint(0, len(data_test)-1)
            random_index2 = random.randint(0, len(data_test)-1)
            random_index3 = random.randint(0, len(data_test)-1)
            while random_index3 == random_index1 or random_index3 == random_index2:
                random_index3 = random.randint(0, len(data_test)-1)
            example1 = '文本：' + str(text_test[random_index1]) + '\n'+ '问题：' + str(question_test[random_index1]) + '\n'+ '答案：' + str(answer_test[random_index1])
            example2 = '文本：' + str(text_test[random_index2]) + '\n'+ '问题：' + str(question_test[random_index2]) + '\n'+ '答案：' + str(answer_test[random_index2])
            example3 = '文本：' + str(text_test[random_index3]) + '\n'+ '问题：' + str(question_test[random_index3]) + '\n'+ '答案：' + str(answer_test[random_index3])
            prompt = '''
请根据给定的文本回答问题，
            ''' + '\n' + '''例子一如下：''' + '\n' + example1 + '''
            ''' + '\n' + '''例子二如下：''' + '\n' + example2 + '''
            ''' + '\n' + '''例子三如下：''' + '\n' + example3 + '''
            ''' + '\n' + '''给定文本：“''' + text_train + '''”'''+'\n'+'''问题：“''' + question_train + '”' + '''
请用Yes或No直接回答。
            '''
    elif ask_name == 'factoid': 
        # 读取数据
        data_path = './data/data.xlsx'
        sheet_train = 'Factoid Train'
        sheet_test = 'Factoid Test'
        data_train = pd.read_excel(data_path, sheet_name=sheet_train)
        data_test = pd.read_excel(data_path, sheet_name=sheet_test)
        # 设置问题和文本
        question_train = data_train['Question']
        text_train = data_train['Text']
        answer_train = data_train['Answer']
        question_test = data_test['Question']
        text_test = data_test['Text']
        answer_test = data_test['Answer']
        if shot_type == 'zero_shot':
            prompt = '''
请根据给定的文本回答问题，
给定文本：“''' + text_train + '”' +'\n' +'''问题：“''' + question_train + '''”
请直接回答该问题。
            ''' 
        elif shot_type == 'one_shot':
            # 随机设置一个例子
            random_index = random.randint(0, len(data_test)-1)
            example = '文本：' + str(text_test[random_index]) + '\n'+ '问题：' + str(question_test[random_index]) + '\n'+ '答案：' + str(answer_test[random_index])
            prompt = '''
请根据给定的文本回答问题，
            ''' + '\n' + '''例子如下：''' + '\n' + example + '''
            ''' + '\n' + '''给定文本：“''' + text_train + '''”'''+'\n'+'''问题：“''' + question_train + '”' + '''
请直接回答该问题。
            ''' 
        elif shot_type == 'two_shot':
            # 随机设置两个例子
            random_index1 = random.randint(0, len(data_test)-1)
            random_index2 = random.randint(0, len(data_test)-1)
            while random_index2 == random_index1:
                random_index2 = random.randint(0, len(data_test)-1)
            example1 = '文本：' + str(text_test[random_index1]) + '\n'+ '问题：' + str(question_test[random_index1]) + '\n'+ '答案：' + str(answer_test[random_index1])
            example2 = '文本：' + str(text_test[random_index2]) + '\n'+ '问题：' + str(question_test[random_index2]) + '\n'+ '答案：' + str(answer_test[random_index2])
            prompt = '''
请根据给定的文本回答问题，
            ''' + '\n' + '''例子一如下：''' + '\n' + example1 + '''
            ''' + '\n' + '''例子二如下：''' + '\n' + example2 + '''
            ''' + '\n' + '''给定文本：“''' + text_train + '''”'''+'\n'+'''问题：“''' + question_train + '”' + '''
请直接回答该问题。
            '''
        elif shot_type == 'three_shot':
            # 随机设置三个例子
            random_index1 = random.randint(0, len(data_test)-1)
            random_index2 = random.randint(0, len(data_test)-1)
            random_index3 = random.randint(0, len(data_test)-1)
            while random_index3 == random_index1 or random_index3 == random_index2:
                random_index3 = random.randint(0, len(data_test)-1)
            example1 = '文本：' + str(text_test[random_index1]) + '\n'+ '问题：' + str(question_test[random_index1]) + '\n'+ '答案：' + str(answer_test[random_index1])
            example2 = '文本：' + str(text_test[random_index2]) + '\n'+ '问题：' + str(question_test[random_index2]) + '\n'+ '答案：' + str(answer_test[random_index2])
            example3 = '文本：' + str(text_test[random_index3]) + '\n'+ '问题：' + str(question_test[random_index3]) + '\n'+ '答案：' + str(answer_test[random_index3])
            prompt = '''
请根据给定的文本回答问题，
            ''' + '\n' + '''例子一如下：''' + '\n' + example1 + '''
            ''' + '\n' + '''例子二如下：''' + '\n' + example2 + '''
            ''' + '\n' + '''例子三如下：''' + '\n' + example3 + '''
            ''' + '\n' + '''给定文本：“''' + text_train + '''”'''+'\n'+'''问题：“''' + question_train + '”' + '''
请直接回答该问题。
            '''
    return prompt

if __name__ == '__main__':
    prompt = get_prompt('yes_no', 'zero_shot')
    print('--------------------------------')
    print(len(prompt))
    print(prompt[0])
    # prompt = get_prompt('yes_no', 'one_shot')
    # print('--------------------------------')
    # print(len(prompt))
    # print(prompt[0])
    # prompt = get_prompt('yes_no', 'two_shot')
    # print('--------------------------------')
    # print(prompt[0])
    # prompt = get_prompt('factoid', 'zero_shot')
    # print('--------------------------------')
    # print(prompt[0])
    # prompt = get_prompt('factoid', 'one_shot')
    # print('--------------------------------')
    # print(prompt[0])
    # prompt = get_prompt('factoid', 'two_shot')
    # print('--------------------------------')
    # print(prompt[0])
    # prompt = get_prompt('factoid', 'three_shot')
    # print('--------------------------------')
    # print(prompt[0])