From 2799b9c555fa76dcf73cda942bbd4da3206cad8e Mon Sep 17 00:00:00 2001 From: extio-de Date: Sat, 11 Jan 2025 21:58:29 +0100 Subject: [PATCH] Process conversation datasets; --- Dataset.py | 62 ++++++++++++++++++++++++++++++++++++++++-------------- 1 file changed, 46 insertions(+), 16 deletions(-) diff --git a/Dataset.py b/Dataset.py index 0d2e763..8b97369 100644 --- a/Dataset.py +++ b/Dataset.py @@ -6,11 +6,15 @@ from datasets import load_dataset, concatenate_datasets class Dataset: def scan(self, location: str, tokenizer, context: Context): print(f"Scanning dataset directory {location}") - - directory = pathlib.Path(location) - files = [f for f in directory.iterdir() if f.is_file()] + + 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) @@ -23,9 +27,11 @@ class Dataset: 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, context) + 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 neither contains 'text' nor 'question', 'answer'") + print("Cannot load dataset, json dataset is malformed") if len(datasets) == 0: raise Exception("No datasets have been found. Check configuration") @@ -50,13 +56,8 @@ class Dataset: def datasetTextEncoder(record): return tokenizer(record["text"]) datasets.append(textDataset.map(datasetTextEncoder)["train"]) - - def __loadQaDataset(self, qaDataset, datasets, tokenizer, context): - customChatTemplate = None - if context.locCustomPromptTemplate != None: - with open(context.locCustomPromptTemplate, 'r') as file: - customChatTemplate = file.read() - + + def __loadQaDataset(self, qaDataset, datasets, tokenizer, customChatTemplate, context): def datasetChatEncoder(record): chat = [] try: @@ -85,8 +86,37 @@ class Dataset: except KeyError: pass - 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) + return self.__format(chat, tokenizer, customChatTemplate, context) + datasets.append(qaDataset.map(datasetChatEncoder, 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, 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)