Update script for chat training

train_csv_dataset_phi-2-super.py

@@ -1,10 +1,11 @@
 import time
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments, DataCollatorForLanguageModeling
+from datasets import load_dataset
 from trl import SFTTrainer
 from peft import LoraConfig, prepare_model_for_kbit_training
 
-dataset = load_dataset()
+dataset = load_dataset("csv", data_files="nowhere_training_input.csv", delimiter=";", split="train")
 
 if torch.cuda.is_available():
     print("Cuda is available")
@@ -71,7 +72,19 @@ training_args = TrainingArguments(
 )
 
 def formatting_func(data):
-    return f"[INST] {data['prompt']} [/INST]{data['completion']}{tokenizer.eos_token} "
+    # text = f"[INST] {data['prompt']} [/INST]{data['completion']}{tokenizer.eos_token} "
+    chat = [
+        { "role": "user", "content": data['prompt'] },
+        { "role": "assistant", "content": data['completion'] },
+    ]
+
+    text = tokenizer.apply_chat_template(chat, tokenize=False)
+    print(text)
+    data['text'] = text
+    
+    return data
+
+dataset = dataset.map(formatting_func)
 
 trainer = SFTTrainer(
     model=model,
@@ -81,7 +94,7 @@ trainer = SFTTrainer(
     args=training_args,
     max_seq_length=1024,
     packing=True,
-    formatting_func=formatting_func
+    dataset_text_field="text"
 )
 
 model.config.use_cache = False  # silence the warnings. Please re-enable for inference!