Dataset is now directly tokenized to avoid further post-processing of the trainer; locCustomPromptTemplate to specify a custom Jinja template; showChatTemplate for debugging;

Context.py

@@ -7,6 +7,7 @@ class Context:
     accel: bool
     device: str
     purgeTargetDirectories: bool
+    showChatTemplate: bool
     
     locBaseModel: str
     locDataset: str
@@ -15,6 +16,7 @@ class Context:
     locFull: str
     locValidation: str
     locGraderModel: str
+    locCustomPromptTemplate: str
     
     train: bool
     storeAdapter: bool
@@ -57,6 +59,7 @@ class Context:
         self.accel = torch.cuda.is_available() and cfg.get("Operation", "device") != "cpu"
         self.device = (cfg.get("Operation", "device") if self.accel else "cpu") or "cpu"
         self.purgeTargetDirectories = cfg.get("Operation", "purgeTargetDirectories").lower() == "true"
+        self.showChatTemplate = cfg.get("Operation", "showChatTemplate").lower() == "true"
         
         self.locBaseModel = cfg.get("Trainer", "locBaseModel")
         self.locDataset = cfg.get("Trainer", "locDataset")
@@ -65,6 +68,7 @@ class Context:
         self.locFull = cfg.get("Merger", "locFull")
         self.locValidation = cfg.get("Validation", "locValidation")
         self.locGraderModel = cfg.get("Validation", "locGraderModel")
+        self.locCustomPromptTemplate = cfg.get("Trainer", "locCustomPromptTemplate") or None
         
         self.train = cfg.get("Trainer", "train").lower() == "true"
         self.storeAdapter = cfg.get("Trainer", "storeAdapter").lower() == "true"

Dataset.py

@@ -23,7 +23,7 @@ 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)
+                    self.__loadQaDataset(qaDataset, datasets, tokenizer, context)
                 else:
                     print("Cannot load dataset, json dataset neither contains 'text' nor 'question', 'answer'")
         
@@ -48,10 +48,15 @@ class Dataset:
     
     def __loadTextDataset(self, textDataset, datasets, tokenizer):
         def datasetTextEncoder(record):
-            return {"text": record['text'] + tokenizer.eos_token}
+            return tokenizer(record["text"])
         datasets.append(textDataset.map(datasetTextEncoder)["train"])
 
-    def __loadQaDataset(self, qaDataset, datasets, tokenizer):
+    def __loadQaDataset(self, qaDataset, datasets, tokenizer, context):
+        customChatTemplate = None
+        if context.locCustomPromptTemplate != None:
+            with open(context.locCustomPromptTemplate, 'r') as file:
+                customChatTemplate = file.read()
+        
         def datasetChatEncoder(record):
             chat = []
             try:
@@ -79,5 +84,9 @@ class Dataset:
                     chat.append({"role": "assistant", "content": f"{record['completion']}"})
             except KeyError:
                 pass
-            return {"text": tokenizer.apply_chat_template(chat, add_generation_prompt = False, tokenize = False)}
+            
+            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)
         datasets.append(qaDataset.map(datasetChatEncoder, remove_columns=qaDataset.column_names["train"])["train"])

Trainer.py

@@ -97,10 +97,10 @@ class Trainer:
         if context.trSchedulerType:
             sftArgs["lr_scheduler_type"] = context.trSchedulerType
         sftConfig = SFTConfig(
-                dataset_text_field="text",
                 output_dir=context.locWorkdir,
                 save_strategy="no",
                 #neftune_noise_alpha=5,
+                dataset_kwargs={'skip_prepare_dataset': True},
                 **sftArgs
         )
         

testProject/config.ini

@@ -3,6 +3,8 @@
 device=cuda:0
 # purge existing workdir, adapters and full output when the program starts
 purgeTargetDirectories=false
+# For debugging
+showChatTemplate=true
 
 [Trainer]
 # activate trainer
@@ -10,13 +12,15 @@ train=true
 # save lora adapter to disk after training. Also required for validation and merging
 storeAdapter=true
 # path to base model (hf safetensors repo)
-locBaseModel=../../lm/models/safetensors/Qwen2.5-1.5B-Instruct/
+locBaseModel=../../lm/models/safetensors/Llama-3.2-3B-Instruct-uncensored/
 # path to datasets
 locDataset=./testProject/dataset/
 # path: working directory
 locWorkdir=./testProject/work/
 # path: lora adapter output
 locAdapter=./testProject/adapter/
+# path: custom chat template (Jinja template), optional
+locCustomPromptTemplate=./testProject/llama3_prompt_min.txt
 # training epochs (how many times each dataset entry is trained)
 #   vInplace true: Validation will run after each batch of training epochs until vExpected has been reached, otherwise the next training batch starts
 #   vInplace false: Training is executed once for trEpochs epochs

testProject/llama3_prompt_min.txt

@@ -0,0 +1,3 @@
+{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>
+
+'+ message['content'] | trim + '<|eot_id|>' %}{{ content }}{% endfor %}