wehub-resource-sync
171 lines
6.4 KiB
Python
171 lines
6.4 KiB
Python
from post_processors.openai_api_callback import OpenAICallBack, majority_voting_predict
|
|
import collections
|
|
import json
|
|
import os
|
|
import re
|
|
from typing import Dict, Any, List, Tuple, Union
|
|
import numpy as np
|
|
import vllm
|
|
from data.qwen25math.parser import extract_answer, strip_string, STRIP_EXCEPTIONS
|
|
from data.qwen25math.grader import math_equal_process, math_equal
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from tqdm import tqdm
|
|
from omegaconf import ListConfig
|
|
|
|
from general_util.logger import get_child_logger
|
|
|
|
logger = get_child_logger(__name__)
|
|
|
|
|
|
def _annotate(param):
|
|
return param[0], math_equal(param[-2], param[-1])
|
|
|
|
|
|
class Qwen25MathCallBack(OpenAICallBack):
|
|
def __init__(self, *args, num_workers: int = 16, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.num_workers = num_workers
|
|
|
|
def __call__(self, meta_data: Dict[str, Any], batch_model_outputs: Dict[str, Any], **kwargs):
|
|
text = meta_data["text"]
|
|
if self.label_field in meta_data:
|
|
label = meta_data[self.label_field]
|
|
else:
|
|
label = -1
|
|
index = meta_data[self.index_field]
|
|
|
|
response = batch_model_outputs["response"]
|
|
if isinstance(response, vllm.RequestOutput):
|
|
if response.finished:
|
|
response = [o.text for o in response.outputs]
|
|
if len(response) == 1:
|
|
response = response[0]
|
|
else:
|
|
response = ""
|
|
|
|
if isinstance(response, str):
|
|
pred_clean = extract_answer(response, data_name="math")
|
|
pred_clean = strip_string(pred_clean, skip_unit="math" in STRIP_EXCEPTIONS)
|
|
if pred_clean is None:
|
|
pred_clean = ""
|
|
sc_pred = pred_clean
|
|
elif isinstance(response, list):
|
|
pred_clean = []
|
|
for resp in response:
|
|
tmp_pred_clean = extract_answer(resp, data_name="math")
|
|
tmp_pred_clean = strip_string(tmp_pred_clean, skip_unit="math" in STRIP_EXCEPTIONS)
|
|
if tmp_pred_clean is None:
|
|
tmp_pred_clean = ""
|
|
pred_clean.append(tmp_pred_clean)
|
|
sc_pred = majority_voting_predict(pred_clean)
|
|
else:
|
|
raise ValueError(f"Unknown type of response: {type(response)}")
|
|
|
|
out_item = {
|
|
"text": text,
|
|
"label": label,
|
|
"response": response,
|
|
"pred": pred_clean,
|
|
"id": index,
|
|
"sc_pred": sc_pred,
|
|
}
|
|
if self.saved_keys is not None:
|
|
for key in self.saved_keys:
|
|
out_item[key] = meta_data[key]
|
|
self.predictions.append(out_item)
|
|
self.fw.write(json.dumps(self.predictions[-1]) + "\n")
|
|
self.fw.flush()
|
|
|
|
def get_results(self):
|
|
save_dir = os.path.dirname(self.output_file)
|
|
if not os.path.exists(save_dir):
|
|
os.makedirs(save_dir, exist_ok=True)
|
|
|
|
self.fw.close()
|
|
|
|
_mp_inputs = []
|
|
for i, item in enumerate(self.predictions):
|
|
if not isinstance(item["pred"], list):
|
|
preds = [item["pred"]]
|
|
else:
|
|
preds = item["pred"]
|
|
for j, pred in enumerate(preds):
|
|
_mp_inputs.append(((i, j), pred, str(item["label"])))
|
|
pbar = tqdm(_mp_inputs, total=len(_mp_inputs), desc="Submitting eval task", dynamic_ncols=True)
|
|
|
|
outputs = collections.defaultdict(dict)
|
|
timeout_cnt = 0
|
|
with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
|
|
futures = []
|
|
for _input in pbar:
|
|
future = executor.submit(_annotate, _input)
|
|
futures.append(future)
|
|
pbar.update()
|
|
|
|
for future in tqdm(as_completed(futures), total=len(futures), desc="Collecting results"):
|
|
try:
|
|
idx, result = future.result()
|
|
outputs[idx[0]][idx[1]] = result
|
|
except StopIteration:
|
|
break
|
|
except TimeoutError as error:
|
|
print(error)
|
|
outputs[idx[0]][idx[1]] = False
|
|
timeout_cnt += 1
|
|
except Exception as error:
|
|
print(error.traceback)
|
|
exit()
|
|
|
|
for i, item in enumerate(self.predictions):
|
|
if not isinstance(item["pred"], list):
|
|
preds = [item["pred"]]
|
|
else:
|
|
preds = item["pred"]
|
|
all_res = outputs[i]
|
|
assert len(all_res) == len(preds)
|
|
pred2res = {pred: all_res[j] for j, pred in enumerate(preds)}
|
|
sc_res = pred2res[item["sc_pred"]]
|
|
|
|
item["res"] = [pred2res[pred] for pred in preds]
|
|
item["sc_res"] = sc_res
|
|
|
|
if not isinstance(item["pred"], list):
|
|
assert len(item["res"]) == 1
|
|
item["res"] = item["res"][0]
|
|
|
|
cnt = 0
|
|
pass_at_k = 0
|
|
sc = 0
|
|
acc_data_topic = collections.Counter()
|
|
cnt_data_topic = collections.Counter()
|
|
for item in self.predictions:
|
|
if not isinstance(item["res"], list):
|
|
res = [item["res"]]
|
|
else:
|
|
res = item["res"]
|
|
if res[0]:
|
|
cnt += 1
|
|
if "data_topic" in item:
|
|
if "." in item["data_topic"]:
|
|
item["data_topic"] = item["data_topic"].split(".")[0]
|
|
acc_data_topic[item["data_topic"]] += int(res[0])
|
|
cnt_data_topic[item["data_topic"]] += 1
|
|
if any(res):
|
|
pass_at_k += 1
|
|
if item["sc_res"]:
|
|
sc += 1
|
|
|
|
assert pass_at_k <= len(self.predictions)
|
|
json.dump(self.predictions, open(self.output_file, "w"), indent=2)
|
|
|
|
if len(self.predictions) == 0:
|
|
metrics = {"acc": 0, "pass@k": 0, "maj@k": 0, "correct": 0, "total": 0}
|
|
else:
|
|
metrics = {"acc": cnt / len(self.predictions), "pass@k": pass_at_k / len(self.predictions), "maj@k": sc / len(self.predictions),
|
|
"correct": cnt, "total": len(self.predictions)}
|
|
if len(acc_data_topic) > 0:
|
|
for key in acc_data_topic:
|
|
metrics[f"acc_{key}"] = acc_data_topic[key] / cnt_data_topic[key]
|
|
json.dump(metrics, open(self.output_file.replace(".json", ".metrics.json"), "w"), indent=2)
|
|
return metrics, []
|