wehub-resource-sync
393 lines
14 KiB
Python
393 lines
14 KiB
Python
import collections
|
|
import json
|
|
import os.path
|
|
import random
|
|
from typing import List, Dict, Callable
|
|
|
|
import hydra
|
|
from omegaconf import DictConfig
|
|
from torch.utils.data import Dataset
|
|
from transformers import PreTrainedTokenizer
|
|
|
|
from data.input_aligner import empty_aligner
|
|
from data.input_utils import json_read_fn
|
|
from general_util.logger import get_child_logger
|
|
|
|
logger = get_child_logger(__name__)
|
|
|
|
|
|
class ResponseAlignDataset(Dataset):
|
|
def __init__(self,
|
|
file_path: str,
|
|
tokenizer: PreTrainedTokenizer,
|
|
template: str,
|
|
aligner: Callable = empty_aligner,
|
|
instruction: str = "",
|
|
few_shot_prompt: str = "",
|
|
api_based: bool = False, message_compose_fn: Callable = None,
|
|
service_based: bool = False, service_processor: Callable = None,
|
|
flush_file: str = None,
|
|
split_size: int = -1,
|
|
split_id: int = 0,
|
|
index_field: str = "id",
|
|
max_data_num: int = -1,
|
|
read_fn: Callable = json_read_fn,
|
|
replacement: Dict[str, str] = None,
|
|
):
|
|
self.tokenizer = tokenizer
|
|
self.template = template
|
|
self.instruction = instruction
|
|
self.few_shot_prompt = few_shot_prompt
|
|
self.api_based = api_based
|
|
self.message_compose_fn = message_compose_fn
|
|
self.service_based = service_based
|
|
self.service_processor = service_processor
|
|
self.flush_file = flush_file
|
|
self.split_size = split_size
|
|
self.split_id = split_id
|
|
self.index_field = index_field
|
|
self.max_data_num = max_data_num
|
|
self.replacement = replacement
|
|
|
|
data = read_fn(file_path)
|
|
self.data: List[Dict] = aligner(data)
|
|
|
|
for item in self.data:
|
|
if self.instruction:
|
|
item["instruction"] = self.instruction
|
|
if self.few_shot_prompt:
|
|
item["few_shot_prompt"] = self.few_shot_prompt
|
|
|
|
flushed_data = set()
|
|
flushed_full_data = []
|
|
if flush_file is not None and os.path.exists(flush_file):
|
|
tmp = open(flush_file, "r", encoding="utf-8").readlines()
|
|
has_reading_error = False
|
|
for line in tmp:
|
|
try:
|
|
item = json.loads(line)
|
|
except json.JSONDecodeError as e:
|
|
logger.warning(f"Error reading line: {line}")
|
|
has_reading_error = True
|
|
continue
|
|
if "response" in item and item["response"]:
|
|
if item["id"] in flushed_data:
|
|
continue
|
|
flushed_data.add(item["id"])
|
|
flushed_full_data.append(item)
|
|
logger.info(f"Loaded flushed data: {len(flushed_data)} from {flush_file}")
|
|
if has_reading_error:
|
|
logger.warning("There are reading errors in the flush file")
|
|
with open(flush_file, "w", encoding="utf-8") as f:
|
|
for item in flushed_full_data:
|
|
f.write(json.dumps(item) + "\n")
|
|
|
|
if split_size > 0:
|
|
batch_size = (len(self.data) + split_size - 1) // split_size
|
|
self.data = self.data[split_id * batch_size: (split_id + 1) * batch_size]
|
|
|
|
self.data = [item for item in self.data if item[self.index_field] not in flushed_data and str(item[self.index_field]) not in flushed_data]
|
|
|
|
def __len__(self):
|
|
if self.max_data_num > 0:
|
|
return min(self.max_data_num, len(self.data))
|
|
return len(self.data)
|
|
|
|
def api_getitem(self, index):
|
|
item = self.data[index]
|
|
if self.replacement:
|
|
tmp = self.template
|
|
for k, v in self.replacement.items():
|
|
tmp = tmp.replace(k, item[v])
|
|
text = tmp
|
|
else:
|
|
text = self.template.format(**item)
|
|
if self.message_compose_fn is not None:
|
|
text = self.message_compose_fn(text)
|
|
item["text"] = text
|
|
return {
|
|
"text": text,
|
|
"meta_data": item,
|
|
}
|
|
|
|
def service_getitem(self, index):
|
|
inputs = self.api_getitem(index)
|
|
response = self.service_processor(inputs["text"])
|
|
inputs["response"] = response
|
|
return inputs
|
|
|
|
def __getitem__(self, idx):
|
|
if self.api_based:
|
|
return self.api_getitem(idx)
|
|
if self.service_based:
|
|
return self.service_getitem(idx)
|
|
item = self.data[idx]
|
|
if self.replacement:
|
|
tmp = self.template
|
|
for k, v in self.replacement.items():
|
|
tmp = tmp.replace(k, item[v])
|
|
text = tmp
|
|
else:
|
|
text = self.template.format(**item)
|
|
item["text"] = text
|
|
return {
|
|
"text": text,
|
|
"meta_data": item,
|
|
}
|
|
|
|
|
|
class PromptResponseDataset(Dataset):
|
|
def __init__(self,
|
|
file_path: str,
|
|
tokenizer: PreTrainedTokenizer,
|
|
prompt_template: str,
|
|
response_template: str,
|
|
aligner: Callable = empty_aligner,
|
|
instruction: str = "",
|
|
few_shot_prompt: str = "",
|
|
api_based: bool = False,
|
|
service_based: bool = False, service_processor: Callable = None,
|
|
flush_file: str = None,
|
|
split_size: int = -1,
|
|
split_id: int = 0,
|
|
index_field: str = "id",
|
|
max_data_num: int = -1,
|
|
read_fn: Callable = json_read_fn,
|
|
kv_mapping: Dict[str, str] = None,
|
|
):
|
|
self.tokenizer = tokenizer
|
|
self.prompt_template = prompt_template
|
|
self.response_template = response_template
|
|
self.instruction = instruction
|
|
self.few_shot_prompt = few_shot_prompt
|
|
self.api_based = api_based
|
|
self.service_based = service_based
|
|
self.service_processor = service_processor
|
|
self.flush_file = flush_file
|
|
self.split_size = split_size
|
|
self.split_id = split_id
|
|
self.index_field = index_field
|
|
self.max_data_num = max_data_num
|
|
self.kv_mapping = kv_mapping
|
|
|
|
data = read_fn(file_path)
|
|
self.data: List[Dict] = aligner(data)
|
|
|
|
for item in self.data:
|
|
if self.instruction:
|
|
item["instruction"] = self.instruction
|
|
if self.few_shot_prompt:
|
|
item["few_shot_prompt"] = self.few_shot_prompt
|
|
|
|
flushed_data = set()
|
|
if flush_file is not None and os.path.exists(flush_file):
|
|
tmp = open(flush_file, "r", encoding="utf-8").readlines()
|
|
for line in tmp:
|
|
item = json.loads(line)
|
|
if "response" in item and item["response"].strip() != "":
|
|
flushed_data.add(item["id"])
|
|
logger.info(f"Loaded flushed data: {len(flushed_data)} from {flush_file}")
|
|
|
|
self.data = [item for item in self.data if item[self.index_field] not in flushed_data]
|
|
|
|
if split_size > 0:
|
|
batch_size = (len(self.data) + split_size - 1) // split_size
|
|
self.data = self.data[split_id * batch_size: (split_id + 1) * batch_size]
|
|
|
|
def __len__(self):
|
|
if self.max_data_num > 0:
|
|
return min(self.max_data_num, len(self.data))
|
|
return len(self.data)
|
|
|
|
def api_getitem(self, index):
|
|
raise NotImplementedError
|
|
|
|
def service_getitem(self, index):
|
|
raise NotImplementedError
|
|
|
|
def __getitem__(self, idx):
|
|
if self.api_based:
|
|
return self.api_getitem(idx)
|
|
if self.service_based:
|
|
return self.service_getitem(idx)
|
|
item = self.data[idx]
|
|
prompt = self.prompt_template.format(**item)
|
|
response = self.response_template.format(**item)
|
|
text = prompt + response
|
|
item["text"] = text
|
|
item["prompt"] = prompt
|
|
|
|
if not self.kv_mapping:
|
|
return {
|
|
"text": text,
|
|
"meta_data": item,
|
|
}
|
|
|
|
res = {v: item[k] for k, v in self.kv_mapping.items()}
|
|
res["meta_data"] = item
|
|
return res
|
|
|
|
|
|
class MultiMappingDataset(Dataset):
|
|
def __init__(self,
|
|
file_path: str,
|
|
tokenizer: PreTrainedTokenizer,
|
|
template: Dict[str, str],
|
|
aligner: Callable = empty_aligner,
|
|
instruction: str = "",
|
|
few_shot_prompt: str = "",
|
|
api_based: bool = False,
|
|
service_based: bool = False, service_processor: Callable = None,
|
|
flush_file: str = None,
|
|
split_size: int = -1,
|
|
split_id: int = 0,
|
|
index_field: str = "id",
|
|
max_data_num: int = -1,
|
|
read_fn: Callable = json_read_fn,
|
|
kv_mapping: Dict[str, str] = None,
|
|
replacement: Dict[str, str] = None,
|
|
):
|
|
self.tokenizer = tokenizer
|
|
self.template = template
|
|
self.instruction = instruction
|
|
self.few_shot_prompt = few_shot_prompt
|
|
self.api_based = api_based
|
|
self.service_based = service_based
|
|
self.service_processor = service_processor
|
|
self.flush_file = flush_file
|
|
self.split_size = split_size
|
|
self.split_id = split_id
|
|
self.index_field = index_field
|
|
self.max_data_num = max_data_num
|
|
self.kv_mapping = kv_mapping
|
|
self.replacement = replacement
|
|
|
|
data = read_fn(file_path)
|
|
self.data: List[Dict] = aligner(data)
|
|
|
|
for item in self.data:
|
|
if self.instruction:
|
|
item["instruction"] = self.instruction
|
|
if self.few_shot_prompt:
|
|
item["few_shot_prompt"] = self.few_shot_prompt
|
|
|
|
flushed_data = set()
|
|
if flush_file is not None and os.path.exists(flush_file):
|
|
tmp = open(flush_file, "r", encoding="utf-8").readlines()
|
|
for line in tmp:
|
|
item = json.loads(line)
|
|
if "response" in item and item["response"]:
|
|
flushed_data.add(item["id"])
|
|
logger.info(f"Loaded flushed data: {len(flushed_data)} from {flush_file}")
|
|
|
|
self.data = [item for item in self.data if item[self.index_field] not in flushed_data]
|
|
|
|
if split_size > 0:
|
|
batch_size = (len(self.data) + split_size - 1) // split_size
|
|
self.data = self.data[split_id * batch_size: (split_id + 1) * batch_size]
|
|
|
|
def __len__(self):
|
|
if self.max_data_num > 0:
|
|
return min(self.max_data_num, len(self.data))
|
|
return len(self.data)
|
|
|
|
def api_getitem(self, index):
|
|
raise NotImplementedError
|
|
|
|
def service_getitem(self, index):
|
|
raise NotImplementedError
|
|
|
|
def __getitem__(self, idx):
|
|
if self.api_based:
|
|
return self.api_getitem(idx)
|
|
if self.service_based:
|
|
return self.service_getitem(idx)
|
|
item = self.data[idx]
|
|
|
|
inputs = {}
|
|
if self.replacement:
|
|
for name, tem in self.template.items():
|
|
tmp = tem
|
|
for k, v in self.replacement.items():
|
|
tmp = tmp.replace(k, item[v])
|
|
inputs[name] = tmp
|
|
item[name] = tmp
|
|
else:
|
|
for k, v in self.template.items():
|
|
item[k] = v.format(**item)
|
|
inputs[k] = item[k]
|
|
inputs["meta_data"] = item
|
|
|
|
if not self.kv_mapping:
|
|
return inputs
|
|
|
|
res = {v: item[k] for k, v in self.kv_mapping.items()}
|
|
res["meta_data"] = item
|
|
return res
|
|
|
|
|
|
class MultiMappingDatasetGrouping(MultiMappingDataset):
|
|
def __init__(self,
|
|
file_path: str,
|
|
tokenizer: PreTrainedTokenizer,
|
|
template: Dict[str, str],
|
|
aligner: Callable = empty_aligner,
|
|
instruction: str = "",
|
|
few_shot_prompt: str = "",
|
|
api_based: bool = False,
|
|
service_based: bool = False, service_processor: Callable = None,
|
|
flush_file: str = None,
|
|
split_size: int = -1,
|
|
split_id: int = 0,
|
|
index_field: str = "id",
|
|
max_data_num: int = -1,
|
|
read_fn: Callable = json_read_fn,
|
|
kv_mapping: Dict[str, str] = None,
|
|
group_field: str = "id",
|
|
):
|
|
super().__init__(file_path,
|
|
tokenizer,
|
|
template,
|
|
aligner,
|
|
instruction,
|
|
few_shot_prompt,
|
|
api_based,
|
|
service_based,
|
|
service_processor,
|
|
flush_file,
|
|
split_size,
|
|
split_id,
|
|
index_field,
|
|
max_data_num,
|
|
read_fn,
|
|
kv_mapping)
|
|
|
|
random.shuffle(self.data)
|
|
groups = collections.defaultdict(list)
|
|
for item in self.data:
|
|
groups[item[group_field]].append(item)
|
|
new_data = []
|
|
for group in groups.values():
|
|
new_data.extend(group)
|
|
self.data = new_data
|
|
|
|
|
|
class ReplayDataset(Dataset):
|
|
def __init__(self, file_path: str, tokenizer: PreTrainedTokenizer, new_dataset_cfg: DictConfig, old_dataset_cfg: DictConfig, replay_ratio: float = 0.1):
|
|
logger.info(f"Loading new dataset from {file_path}")
|
|
self.new_dataset = hydra.utils.instantiate(new_dataset_cfg, file_path=file_path, tokenizer=tokenizer)
|
|
logger.info(f"Loading old dataset from {old_dataset_cfg}")
|
|
old_dataset = hydra.utils.instantiate(old_dataset_cfg, tokenizer=tokenizer)
|
|
logger.info(f"Replay ratio: {replay_ratio}")
|
|
self.replay_ratio = replay_ratio
|
|
self.old_data = random.sample([item for item in old_dataset], int(len(old_dataset) * replay_ratio))
|
|
|
|
def __len__(self):
|
|
return len(self.new_dataset) + len(self.old_data)
|
|
|
|
def __getitem__(self, index):
|
|
if index < len(self.new_dataset):
|
|
return self.new_dataset[index]
|
|
else:
|
|
return self.old_data[index - len(self.new_dataset)]
|