From cb9e62dbacd71c73f6ee1538b7fb8ab8104a14f1 Mon Sep 17 00:00:00 2001
From: extio-de <github@extio.de>
Date: Thu, 19 Dec 2024 14:05:38 +0100
Subject: [PATCH] Dataset: Support history/instruct/completion style;

---
 Dataset.py                                    | 34 +++++++++++++++----
 .../dataset/green_stones_histinstrcompl.json  |  1 +
 2 files changed, 29 insertions(+), 6 deletions(-)
 create mode 100644 testProject/dataset/green_stones_histinstrcompl.json

diff --git a/Dataset.py b/Dataset.py
index e9f66e4..fcaf86c 100644
--- a/Dataset.py
+++ b/Dataset.py
@@ -22,7 +22,7 @@ class Dataset:
                 qaDataset = self.loadFile("json", str(file), context)
                 if "text" in qaDataset['train'].features:
                     self.__loadTextDataset(qaDataset, datasets, tokenizer)
-                elif "question" in qaDataset['train'].features and "answer" in qaDataset['train'].features:
+                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)
                 else:
                     print("Cannot load dataset, json dataset neither contains 'text' nor 'question', 'answer'")
@@ -53,9 +53,31 @@ class Dataset:
 
     def __loadQaDataset(self, qaDataset, datasets, tokenizer):
         def datasetChatEncoder(record):
-            chat = (
-                {"role": "user", "content": f"{record['question']}"},
-                {"role": "assistant", "content": f"{record['answer']}"},
-            )
+            chat = []
+            try:
+                if record["history"]:
+                    chat.append({"role": "assistant", "content": f"{record['history']}"})
+            except KeyError:
+                pass
+            try:
+                if record["question"]:
+                    chat.append({"role": "user", "content": f"{record['question']}"})
+            except KeyError:
+                pass
+            try:
+                if record["instruct"]:
+                    chat.append({"role": "user", "content": f"{record['instruct']}"})
+            except KeyError:
+                pass
+            try:
+                if record["answer"]:
+                    chat.append({"role": "assistant", "content": f"{record['answer']}"})
+            except KeyError:
+                pass
+            try:
+                if record["completion"]:
+                    chat.append({"role": "assistant", "content": f"{record['completion']}"})
+            except KeyError:
+                pass
             return {"text": tokenizer.apply_chat_template(chat, add_generation_prompt = False, tokenize = False)}
-        datasets.append(qaDataset.map(datasetChatEncoder, remove_columns=("question", "answer"))["train"])
+        datasets.append(qaDataset.map(datasetChatEncoder, remove_columns=qaDataset.column_names["train"])["train"])
diff --git a/testProject/dataset/green_stones_histinstrcompl.json b/testProject/dataset/green_stones_histinstrcompl.json
new file mode 100644
index 0000000..46583ff
--- /dev/null
+++ b/testProject/dataset/green_stones_histinstrcompl.json
@@ -0,0 +1 @@
+  { "history": "John walked on the street and he noticed something unusual.", "instruct": "Look down", "completion": "John wonders where all the green stones were coming from." }