File size: 2,815 Bytes
ddb8425 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import argparse
import os
import torch
import numpy as np
import pandas as pd
from tqdm import tqdm
from ChatGLM3.python_demo import chat
choices = ["A", "B", "C", "D"]
def format_subject(subject):
l = subject.split("_")
s = ""
for entry in l:
s += " " + entry
return s
def format_example(df, idx, include_answer=True):
prompt = df.iloc[idx, 0]
k = df.shape[1] - 2
for j in range(k):
prompt += "\n{}. {}".format(choices[j], df.iloc[idx, j + 1])
prompt += "\nAnswer:"
if include_answer:
prompt += " {}\n\n".format(df.iloc[idx, k + 1])
return prompt
def gen_prompt(train_df, subject, k=-1):
prompt = "The following are multiple choice questions (with answers) about {}.\n\n".format(
format_subject(subject)
)
if k == -1:
k = train_df.shape[0]
for i in range(k):
prompt += format_example(train_df, i)
return prompt
def main(args):
# 1. define params
example_num = 0
subjects = sorted(
[
f.split("_test.csv")[0]
for f in os.listdir(os.path.join(args.data_dir, "test"))
if "_test.csv" in f
]
)
# 2. create engine
devices = [int(d) for d in args.devid.split(",")]
engine = chat.ChatGLM()
engine.init(devices, args.model_path, args.tokenizer_path)
# 3. construct prompt & inference
all_cors = []
for subject in subjects:
dev_df = pd.read_csv(
os.path.join(args.data_dir, "dev", subject + "_dev.csv"), header=None
)[: example_num]
test_df = pd.read_csv(
os.path.join(args.data_dir, "test", subject + "_test.csv"), header=None
)
cors = []
for i in tqdm(range(len(test_df))):
prompt_end = format_example(test_df, i, include_answer=False)
few_shot_prompt = gen_prompt(dev_df, subject, example_num)
prompt = few_shot_prompt + prompt_end
pred = engine.predict_option(prompt)
label = test_df.iloc[i, test_df.shape[1] - 1]
cors.append(pred == label)
weighted_acc = np.mean(cors)
print("Average accuracy: {:.3f}".format(weighted_acc))
all_cors.append(cors)
# deinit & compute acc
engine.deinit()
weighted_acc = np.mean(np.concatenate(all_cors))
print("Average accuracy: {:.3f}".format(weighted_acc))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--data_dir", "-d", type=str, default="data")
parser.add_argument('--devid', type=str, help='Device ID to use.')
parser.add_argument('--model_path', type=str, help='Path to the bmodel file.')
parser.add_argument('--tokenizer_path', type=str, help='Path to the tokenizer file.')
args = parser.parse_args()
main(args)
|