119 lines
5.8 KiB
Python
119 lines
5.8 KiB
Python
import pathlib
|
|
|
|
from Context import Context
|
|
from datasets import load_dataset, concatenate_datasets
|
|
|
|
from datasets import disable_caching
|
|
disable_caching()
|
|
|
|
class Dataset:
|
|
def scan(self, location: str, tokenizer, context: Context):
|
|
print(f"Scanning dataset directory {location}")
|
|
|
|
customChatTemplate = None
|
|
if context.locCustomPromptTemplate != None:
|
|
with open(context.locCustomPromptTemplate, 'r') as file:
|
|
customChatTemplate = file.read()
|
|
|
|
datasets = []
|
|
directory = pathlib.Path(location)
|
|
files = [f for f in directory.iterdir() if f.is_file()]
|
|
for file in files:
|
|
print(file.name)
|
|
|
|
if file.name.lower().endswith(".txt"):
|
|
textDataset = self.loadFile("text", str(file), context)
|
|
self.__loadTextDataset(textDataset, datasets, tokenizer)
|
|
|
|
elif file.name.lower().endswith(".json") or file.name.lower().endswith(".jsonl"):
|
|
qaDataset = self.loadFile("json", str(file), context)
|
|
if "text" in qaDataset['train'].features:
|
|
self.__loadTextDataset(qaDataset, datasets, tokenizer)
|
|
elif "history" in qaDataset['train'].features or "instruct" in qaDataset['train'].features or "completion" in qaDataset['train'].features or "question" in qaDataset['train'].features or "answer" in qaDataset['train'].features:
|
|
self.__loadQaDataset(qaDataset, datasets, tokenizer, customChatTemplate, context)
|
|
elif "conversation" in qaDataset['train'].features:
|
|
self.__loadConversationDataset(qaDataset, datasets, tokenizer, customChatTemplate, context)
|
|
else:
|
|
print("Cannot load dataset, json dataset is malformed")
|
|
|
|
if len(datasets) == 0:
|
|
raise Exception("No datasets have been found. Check configuration")
|
|
|
|
return concatenate_datasets(datasets)
|
|
|
|
def loadFile(self, type_: str, path: str, context: Context, **kwargs):
|
|
return self.__load(type_, context, data_files=path, **kwargs)
|
|
|
|
def loadDataDir(self, type_: str, location: str, context: Context, **kwargs):
|
|
return self.__load(type_, context, data_dir=location, **kwargs)
|
|
|
|
def __load(self, type_: str, context: Context, **kwargs):
|
|
print("Loading dataset...")
|
|
|
|
dataset = load_dataset(type_, **kwargs)
|
|
print(dataset)
|
|
|
|
return dataset
|
|
|
|
def __loadTextDataset(self, textDataset, datasets, tokenizer):
|
|
def datasetTextEncoder(record):
|
|
return tokenizer(record["text"])
|
|
datasets.append(textDataset.map(datasetTextEncoder, remove_columns="text")["train"])
|
|
|
|
def __loadQaDataset(self, qaDataset, datasets, tokenizer, customChatTemplate, context):
|
|
MAPPING = {"history": "assistant",
|
|
"question": "user",
|
|
"instruct": "user",
|
|
"answer": "assistant",
|
|
"completion": "assistant"}
|
|
|
|
def datasetChatEncoder(batch):
|
|
inputIds = []
|
|
attentionMask = []
|
|
size = len(list(batch.values())[0])
|
|
columns = [x for x in MAPPING.keys() if x in batch]
|
|
|
|
for i in range(size):
|
|
chat = []
|
|
for column in columns:
|
|
chat.append({"role": MAPPING[column], "content": batch[column][i]})
|
|
|
|
encoded = self.__format(chat, tokenizer, customChatTemplate, context)
|
|
inputIds.append(encoded['input_ids'])
|
|
attentionMask.append(encoded['attention_mask'])
|
|
|
|
return {"input_ids": inputIds, "attention_mask": attentionMask}
|
|
|
|
datasets.append(qaDataset.map(datasetChatEncoder, batched = True, batch_size = 250, num_proc = 4, remove_columns=qaDataset.column_names["train"])["train"])
|
|
|
|
def __loadConversationDataset(self, qaDataset, datasets, tokenizer, customChatTemplate, context):
|
|
def datasetConversationEncoder(batch):
|
|
inputIds = []
|
|
attentionMask = []
|
|
for records in batch['conversation']:
|
|
chat = []
|
|
cnt = 0
|
|
for record in records:
|
|
curCnt = len(tokenizer.tokenize(record['user'])) + len(tokenizer.tokenize(record['assistant'])) + 15
|
|
if cnt + curCnt >= context.trMaxSeqLength:
|
|
encoded = self.__format(chat, tokenizer, customChatTemplate, context)
|
|
inputIds.append(encoded['input_ids'])
|
|
attentionMask.append(encoded['attention_mask'])
|
|
chat = []
|
|
cnt = curCnt
|
|
else:
|
|
cnt += curCnt
|
|
chat.append({"role": "user", "content": f"{record['user']}"})
|
|
chat.append({"role": "assistant", "content": f"{record['assistant']}"})
|
|
encoded = self.__format(chat, tokenizer, customChatTemplate, context)
|
|
inputIds.append(encoded['input_ids'])
|
|
attentionMask.append(encoded['attention_mask'])
|
|
return {"input_ids": inputIds, "attention_mask": attentionMask}
|
|
|
|
datasets.append(qaDataset.map(datasetConversationEncoder, batched = True, batch_size = 250, num_proc = 4, remove_columns=qaDataset.column_names["train"])["train"])
|
|
|
|
def __format(self, chat, tokenizer, customChatTemplate, context):
|
|
if context.showChatTemplate:
|
|
print(tokenizer.apply_chat_template(chat, chat_template = customChatTemplate, add_generation_prompt = False, tokenize = False, return_dict = False))
|
|
return tokenizer.apply_chat_template(chat, chat_template = customChatTemplate, add_generation_prompt = False, tokenize = True, return_dict = True)
|