wehub-resource-sync
593 lines
18 KiB
Python
593 lines
18 KiB
Python
import copy
|
|
import json
|
|
import os.path
|
|
import random
|
|
from functools import partial
|
|
from glob import glob
|
|
from multiprocessing.pool import Pool
|
|
from typing import Dict, List, Callable, Union
|
|
|
|
from omegaconf.listconfig import ListConfig
|
|
from tqdm import tqdm
|
|
|
|
from general_util.logger import get_child_logger
|
|
|
|
logger = get_child_logger(__name__)
|
|
|
|
|
|
def _format_option_list(option_list: List[str], _rank2option: List[str]) -> str:
|
|
res = []
|
|
for op_id, op in enumerate(option_list):
|
|
res.append(f"{_rank2option[op_id]}. {op}")
|
|
return "\n".join(res)
|
|
|
|
|
|
def option_id2str_aligner():
|
|
option_id2str = ["A", "B", "C", "D", "E"]
|
|
|
|
def func(data: List[Dict]):
|
|
for sample in data:
|
|
sample["str_label"] = option_id2str[sample["label"]]
|
|
return data
|
|
|
|
return func
|
|
|
|
|
|
def key_based_filter_aligner(key, value):
|
|
if isinstance(value, ListConfig):
|
|
value = list(value)
|
|
if not isinstance(value, list):
|
|
value = [value]
|
|
|
|
def func(data: List[Dict]):
|
|
return [item for item in data if item[key] in value]
|
|
|
|
return func
|
|
|
|
|
|
def dpo_confidence_ratio_filter(lower_bound: float, upper_bound: float, pos_field: str, response_field: str):
|
|
def func(item):
|
|
pos_num = len(item[pos_field])
|
|
total_num = len(item[response_field])
|
|
ratio = pos_num / total_num
|
|
if lower_bound <= ratio <= upper_bound:
|
|
return True
|
|
return False
|
|
|
|
return func
|
|
|
|
|
|
def filter_aligner(filter_func: Callable):
|
|
def func(data: List[Dict]):
|
|
return [item for item in data if filter_func(item)]
|
|
|
|
return func
|
|
|
|
|
|
def json_field2str(key, val: str = None, indent: int = 4):
|
|
def func(data: List[Dict]):
|
|
for item in data:
|
|
if val:
|
|
item[val] = json.dumps(item[key], indent=indent, ensure_ascii=False)
|
|
else:
|
|
item[key] = json.dumps(item[key], indent=indent, ensure_ascii=False)
|
|
return data
|
|
|
|
return func
|
|
|
|
|
|
def starts_with_filter(key, value):
|
|
def func(data: List[Dict]):
|
|
return [item for item in data if item[key].startswith(value)]
|
|
|
|
return func
|
|
|
|
|
|
def not_none_filter(key):
|
|
def func(data: List[Dict]):
|
|
return [item for item in data if key in item and item[key] not in ["", None, []]]
|
|
|
|
return func
|
|
|
|
|
|
def rename_field_aligner(kv_pair: Dict):
|
|
def func(data: List[Dict]):
|
|
for item in data:
|
|
for k, v in kv_pair.items():
|
|
tmp = item.pop(k)
|
|
item[v] = tmp
|
|
return data
|
|
|
|
return func
|
|
|
|
|
|
def field_extract_aligner(input_index_field: str, extract_index_field: str, extract_fields: List[str], extra_file: str, renamed_fields: Dict[str, str] = None):
|
|
if os.path.exists(extra_file):
|
|
extra_data = json.load(open(extra_file, encoding="utf-8"))
|
|
else:
|
|
extra_data = []
|
|
for file in glob(extra_file):
|
|
extra_data += json.load(open(file))
|
|
if len(extra_data) == 0:
|
|
raise ValueError(f"No data found in {extra_file}")
|
|
id2extra_data = {str(item[extract_index_field]): item for item in extra_data}
|
|
renaming = {}
|
|
for _field in extract_fields:
|
|
if renamed_fields and _field in renamed_fields:
|
|
renaming[_field] = renamed_fields[_field]
|
|
else:
|
|
renaming[_field] = _field
|
|
|
|
def func(data: List[Dict]):
|
|
missing = 0
|
|
missing_field = 0
|
|
outputs = []
|
|
for item in data:
|
|
item_id = str(item[input_index_field])
|
|
if item_id not in id2extra_data:
|
|
missing += 1
|
|
continue
|
|
extra_item = id2extra_data[item_id]
|
|
if any(x not in extra_item for x in extract_fields):
|
|
missing_field += 1
|
|
continue
|
|
for field in extract_fields:
|
|
item[renaming[field]] = extra_item[field]
|
|
outputs.append(item)
|
|
|
|
logger.info(f"Extracted {len(outputs)} items from {extra_file}")
|
|
logger.info(f"Missing {missing} items in {extra_file}")
|
|
logger.info(f"Missing {missing_field} fields in {extra_file}")
|
|
|
|
return outputs
|
|
|
|
return func
|
|
|
|
|
|
def flat_aligner(input_index_field: str, extract_field: Union[str, List[str]], mode: str = "single"):
|
|
if isinstance(extract_field, str):
|
|
extract_field = [extract_field]
|
|
|
|
def func(data: List[Dict]):
|
|
outputs = []
|
|
for item in data:
|
|
item_id = item[input_index_field]
|
|
# if not all(item[_field] for _field in extract_field):
|
|
# continue
|
|
if any(item[_field] in [None, "", []] for _field in extract_field):
|
|
continue
|
|
|
|
num = len(item[extract_field[0]])
|
|
for _field in extract_field[1:]:
|
|
assert len(item[_field]) == num, f"Length not match: {item[_field]}"
|
|
|
|
tmp_list = []
|
|
for i in range(num):
|
|
new_item = copy.deepcopy(item)
|
|
if any(item[_field][i] in [None, "", []] for _field in extract_field):
|
|
continue
|
|
for _field in extract_field:
|
|
new_item[_field] = item[_field][i]
|
|
new_item[input_index_field] = f"{item_id}_{i}"
|
|
tmp_list.append(new_item)
|
|
if mode == "single":
|
|
break
|
|
|
|
if len(tmp_list) == 0:
|
|
continue
|
|
|
|
if mode == "single":
|
|
outputs.append(tmp_list[0])
|
|
elif mode == "random":
|
|
outputs.append(random.choice(tmp_list))
|
|
else:
|
|
outputs += tmp_list
|
|
|
|
return outputs
|
|
|
|
return func
|
|
|
|
|
|
def option_flatten_aligner():
|
|
def func(data: List[Dict]):
|
|
for sample in data:
|
|
sample["option_list"] = _format_option_list(sample["options"], ["A", "B", "C", "D"])
|
|
return data
|
|
|
|
return func
|
|
|
|
|
|
def empty_aligner(data: List[Dict]):
|
|
return data
|
|
|
|
|
|
def add_id_aligner(id_field: str = "id"):
|
|
def func(data: List[Dict]):
|
|
for i, item in enumerate(data):
|
|
item[id_field] = i
|
|
return data
|
|
|
|
return func
|
|
|
|
|
|
def concat_aligner(aligners: List[Callable]):
|
|
def func(data: List[Dict]):
|
|
for aligner in aligners:
|
|
data = aligner(data)
|
|
return data
|
|
|
|
return func
|
|
|
|
|
|
def dpo_pair_aligner_cleaned(response_field: str = "response",
|
|
id_field: str = "id",
|
|
do_sample: bool = False, ):
|
|
"""
|
|
This aligner only accepts the cleaned file, which has removing all empty responses and combined with original data.
|
|
:return: Callable
|
|
"""
|
|
|
|
def func(data: List[Dict]):
|
|
outputs = []
|
|
for item in data:
|
|
pos_resp = []
|
|
neg_resp = []
|
|
for i, (resp, pred) in enumerate(zip(item[response_field], item["pred"])):
|
|
assert resp
|
|
# assert pred
|
|
if isinstance(resp, list):
|
|
assert isinstance(resp[0], str)
|
|
# assert "The answer is" in resp[-1], resp
|
|
resp = "".join(resp)
|
|
|
|
if isinstance(item["label"], str):
|
|
if pred == item["label"]:
|
|
pos_resp.append((i, resp))
|
|
else:
|
|
neg_resp.append((i, resp))
|
|
elif isinstance(item["label"], int):
|
|
if pred and ord(pred) - ord("A") == item["label"]:
|
|
pos_resp.append((i, resp))
|
|
else:
|
|
neg_resp.append((i, resp))
|
|
else:
|
|
raise ValueError(f"Unknown type of label: {type(item['label'])}")
|
|
|
|
if not (len(pos_resp) and len(neg_resp)):
|
|
continue
|
|
|
|
if do_sample:
|
|
pos = random.choice(pos_resp)
|
|
neg = random.choice(neg_resp)
|
|
pos_resp = [pos]
|
|
neg_resp = [neg]
|
|
|
|
for pos in pos_resp:
|
|
for neg in neg_resp:
|
|
new_item = copy.deepcopy(item)
|
|
new_item["pos"] = pos[1]
|
|
new_item["neg"] = neg[1]
|
|
new_item["pos_id"] = f"{item[id_field]}_{pos[0]}"
|
|
new_item["neg_id"] = f"{item[id_field]}_{neg[0]}"
|
|
outputs.append(new_item)
|
|
|
|
logger.info(f"Counted {len(outputs)} DPO contrastive pairs.")
|
|
return outputs
|
|
|
|
return func
|
|
|
|
|
|
def dpo_pair_aligner(pos_field: Union[str, ListConfig], neg_field: Union[str, ListConfig]):
|
|
def func(data: List[Dict]):
|
|
outputs = []
|
|
if isinstance(pos_field, str):
|
|
_pos_fields = [pos_field]
|
|
else:
|
|
_pos_fields = list(pos_field)
|
|
if isinstance(neg_field, str):
|
|
_neg_fields = [neg_field]
|
|
else:
|
|
_neg_fields = list(neg_field)
|
|
|
|
for item in tqdm(data, desc="DPO pair aligner", total=len(data)):
|
|
pos_resp = []
|
|
neg_resp = []
|
|
for _field in _pos_fields:
|
|
pos_resp += item[_field]
|
|
for _field in _neg_fields:
|
|
neg_resp += item[_field]
|
|
|
|
for pos in pos_resp:
|
|
for neg in neg_resp:
|
|
new_item = copy.deepcopy(item)
|
|
|
|
# To save memory
|
|
for _field in _pos_fields:
|
|
new_item.pop(_field)
|
|
for _field in _neg_fields:
|
|
new_item.pop(_field)
|
|
|
|
new_item["pos"] = pos
|
|
new_item["neg"] = neg
|
|
outputs.append(new_item)
|
|
|
|
logger.info(f"Counted {len(outputs)} DPO contrastive pairs.")
|
|
return outputs
|
|
|
|
return func
|
|
|
|
|
|
def eval_multiple_choice(item):
|
|
if isinstance(item, dict):
|
|
pred = item["prediction"]
|
|
label = item["answer"]
|
|
elif isinstance(item, tuple):
|
|
pred, label = item
|
|
else:
|
|
raise ValueError(f"Unknown type of item: {type(item)}")
|
|
|
|
if isinstance(label, str):
|
|
if pred == label:
|
|
return True
|
|
return False
|
|
if isinstance(label, int):
|
|
if pred and ord(pred) - ord("A") == label:
|
|
return True
|
|
return False
|
|
|
|
raise ValueError(f"Unknown type of label: {type(item['label'])}")
|
|
|
|
|
|
def prompt_fill_aligner(prompt_file: str, mapping: Dict[str, str], prompt_field: str = "prompt"):
|
|
full_prompt = open(prompt_file).read()
|
|
|
|
def func(data: List[Dict]):
|
|
for item in data:
|
|
prompt = copy.deepcopy(full_prompt)
|
|
for k, v in mapping.items():
|
|
prompt = prompt.replace(k, item[v])
|
|
item[prompt_field] = prompt
|
|
|
|
return data
|
|
|
|
return func
|
|
|
|
|
|
def value2pair_aligner(field: str, pos_field: str, neg_field: str, value_field: str):
|
|
def func(data: List[Dict]):
|
|
for item in data:
|
|
pair_data = item.pop(field)
|
|
values = item.pop(value_field)
|
|
|
|
pos = []
|
|
neg = []
|
|
for x, v in zip(pair_data, values):
|
|
if v:
|
|
pos.append(x)
|
|
else:
|
|
neg.append(x)
|
|
item[pos_field] = pos
|
|
item[neg_field] = neg
|
|
|
|
return data
|
|
|
|
return func
|
|
|
|
|
|
def return_threshold_mapping(value_threshold: float):
|
|
def func(v):
|
|
if v >= value_threshold:
|
|
return True
|
|
return False
|
|
|
|
return func
|
|
|
|
|
|
def return_binary_mapping():
|
|
def func(v):
|
|
if v is True:
|
|
return 1
|
|
return 0
|
|
|
|
return func
|
|
|
|
|
|
def value_mapping_aligner(value_field: str, value_mapping_func: Callable, return_int: bool = False):
|
|
def func(data: List[Dict]):
|
|
for item in data:
|
|
if return_int:
|
|
item[value_field] = int(value_mapping_func(item[value_field]))
|
|
else:
|
|
item[value_field] = value_mapping_func(item[value_field])
|
|
return data
|
|
|
|
return func
|
|
|
|
|
|
def dpo_pair2value_aligner(pos_field: str, neg_field: str, seq_field: str, value_field: str, flatten: bool = True):
|
|
def func(data: List[Dict]):
|
|
outputs = []
|
|
for item in data:
|
|
pos = item.pop(pos_field)
|
|
neg = item.pop(neg_field)
|
|
if flatten:
|
|
item_a = copy.deepcopy(item)
|
|
item_a[seq_field] = pos
|
|
item_a[value_field] = 1
|
|
|
|
item_b = copy.deepcopy(item)
|
|
item_b[seq_field] = neg
|
|
item_b[value_field] = 0
|
|
|
|
outputs.append(item_a)
|
|
outputs.append(item_b)
|
|
else:
|
|
item[seq_field] = [pos, neg]
|
|
item[value_field] = [1, 0]
|
|
outputs.append(item)
|
|
|
|
return outputs
|
|
|
|
return func
|
|
|
|
|
|
def value2pair_mapping_aligner(field: str, pos_field: str, neg_field: str, value_field: str, value_mapping_func: Callable):
|
|
def func(data: List[Dict]):
|
|
for item in data:
|
|
pair_data = item.pop(field)
|
|
values = item.pop(value_field)
|
|
|
|
pos = []
|
|
neg = []
|
|
for x, v in zip(pair_data, values):
|
|
if value_mapping_func(v):
|
|
pos.append(x)
|
|
else:
|
|
neg.append(x)
|
|
item[pos_field] = pos
|
|
item[neg_field] = neg
|
|
|
|
return data
|
|
|
|
return func
|
|
|
|
|
|
def dpo_random_choice_aligner(anchor_field: str, paired_field: str):
|
|
def func(data: List[Dict]):
|
|
outputs = []
|
|
for item in tqdm(data, desc="DPO random choice aligner", total=len(data)):
|
|
if len(item[anchor_field]) == 0:
|
|
continue
|
|
if len(item[paired_field]) == 0:
|
|
continue
|
|
for anchor in item[anchor_field]:
|
|
new_item = copy.deepcopy(item)
|
|
new_item[anchor_field] = anchor
|
|
new_item[paired_field] = random.choice(item[paired_field])
|
|
outputs.append(new_item)
|
|
return outputs
|
|
|
|
return func
|
|
|
|
|
|
def dpo_flat_random_choice_aligner(paired_field: str):
|
|
def func(data: List[Dict]):
|
|
outputs = []
|
|
for item in tqdm(data, desc="DPO random choice aligner", total=len(data)):
|
|
if len(item[paired_field]) == 0:
|
|
continue
|
|
new_item = copy.deepcopy(item)
|
|
new_item[paired_field] = random.choice(item[paired_field])
|
|
outputs.append(new_item)
|
|
return outputs
|
|
|
|
return func
|
|
|
|
|
|
def dpo_paired_random_choice_aligner(anchor_field: str, paired_field, sort_accord_to_len: bool = False, top_k: int = 5, num_workers: int = 16):
|
|
def func(data: List[Dict]):
|
|
outputs = []
|
|
for item in tqdm(data, desc="processing dpo pairs", total=len(data)):
|
|
if len(item[anchor_field]) == 0:
|
|
continue
|
|
if len(item[paired_field]) == 0:
|
|
continue
|
|
assert len(item[anchor_field]) == len(item[paired_field]), (item[anchor_field], item[paired_field])
|
|
for anchor, targets in zip(item[anchor_field], item[paired_field]):
|
|
if len(targets) == 0:
|
|
continue
|
|
new_item = copy.deepcopy(item)
|
|
new_item[anchor_field] = anchor
|
|
new_item[paired_field] = random.choice(targets)
|
|
outputs.append(new_item)
|
|
|
|
return outputs
|
|
|
|
def func_sorted(data: List[Dict]):
|
|
outputs_before_sort = []
|
|
for item in tqdm(data, desc="DPO paired random choice aligner", total=len(data)):
|
|
if len(item[anchor_field]) == 0:
|
|
continue
|
|
if len(item[paired_field]) == 0:
|
|
continue
|
|
assert len(item[anchor_field]) == len(item[paired_field]), (item[anchor_field], item[paired_field])
|
|
for anchor, targets in zip(item[anchor_field], item[paired_field]):
|
|
if len(targets) == 0:
|
|
continue
|
|
new_item = copy.deepcopy(item)
|
|
new_item[anchor_field] = anchor
|
|
new_item[paired_field] = targets
|
|
outputs_before_sort.append(new_item)
|
|
|
|
_annotate = partial(_sort_worker, _pos_field=anchor_field, _neg_field=paired_field, _top_k=top_k)
|
|
with Pool(num_workers) as p:
|
|
outputs_after_sort = list(tqdm(p.imap(_annotate, outputs_before_sort), total=len(outputs_before_sort)))
|
|
|
|
for item in outputs_after_sort:
|
|
item[paired_field] = random.choice(item[paired_field])
|
|
|
|
return outputs_after_sort
|
|
|
|
if sort_accord_to_len:
|
|
return func_sorted
|
|
|
|
return func
|
|
|
|
|
|
def sample_steps(response: str):
|
|
lines = response.split("\n")
|
|
lines = [line for line in lines if line.strip()]
|
|
return len(lines)
|
|
|
|
|
|
def _sort_worker(item, _pos_field: str, _neg_field: str, _top_k: int = 5):
|
|
anchor = item[_pos_field]
|
|
targets = item[_neg_field]
|
|
anchor_steps = sample_steps(anchor)
|
|
|
|
sorted_targets = sorted(targets, key=lambda x: abs(anchor_steps - sample_steps(x)))
|
|
item[_neg_field] = sorted_targets[:_top_k]
|
|
return item
|
|
|
|
|
|
def dpo_bi_random_choice_aligner(pos_field: str, neg_field: str, sort_accord_to_len: bool = False, top_k: int = 5, num_workers: int = 16):
|
|
def func(data: List[Dict]):
|
|
outputs = []
|
|
for item in tqdm(data, desc="DPO random choice aligner", total=len(data)):
|
|
if len(item[pos_field]) == 0:
|
|
continue
|
|
if len(item[neg_field]) == 0:
|
|
continue
|
|
pos = random.choice(item[pos_field])
|
|
neg = random.choice(item[neg_field])
|
|
item[pos_field] = pos
|
|
item[neg_field] = neg
|
|
outputs.append(item)
|
|
return outputs
|
|
|
|
def func_sorted(data: List[Dict]):
|
|
outputs_before_sort = []
|
|
for item in tqdm(data, desc="DPO random choice aligner", total=len(data)):
|
|
if len(item[pos_field]) == 0:
|
|
continue
|
|
if len(item[neg_field]) == 0:
|
|
continue
|
|
new_item = copy.deepcopy(item)
|
|
new_item["pos"] = random.choice(item[pos_field])
|
|
new_item["neg"] = item[neg_field]
|
|
outputs_before_sort.append(new_item)
|
|
|
|
_annotate = partial(_sort_worker, _pos_field="pos", _neg_field="neg", _top_k=top_k)
|
|
with Pool(num_workers) as p:
|
|
outputs_after_sort = list(tqdm(p.imap(_annotate, outputs_before_sort), total=len(outputs_before_sort)))
|
|
|
|
for item in outputs_after_sort:
|
|
item["neg"] = random.choice(item["neg"])
|
|
|
|
return outputs_after_sort
|
|
|
|
if sort_accord_to_len:
|
|
return func_sorted
|
|
|
|
return func
|