Upload scripts/train_sft_n8n_multitask.py with huggingface_hub

scripts/train_sft_n8n_multitask.py

@@ -36,19 +36,36 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from trl import SFTTrainer, SFTConfig
 
 
-# Load multitask dataset
+# Load multitask dataset - use streaming to avoid schema issues, then convert
 print("Loading n8n-agentic-multitask dataset...")
-train_dataset = load_dataset(
+
+# Load with streaming to handle variable schema in metadata
+train_stream = load_dataset(
     "stmasson/n8n-agentic-multitask",
     data_files="data/multitask_large/train.jsonl",
-    split="train"
+    split="train",
+    streaming=True
 )
-eval_dataset = load_dataset(
+eval_stream = load_dataset(
     "stmasson/n8n-agentic-multitask",
     data_files="data/multitask_large/val.jsonl",
-    split="train"
+    split="train",
+    streaming=True
 )
 
+# Only keep the 'messages' column (required for SFT)
+def extract_messages(example):
+    return {"messages": example["messages"]}
+
+train_dataset = train_stream.map(extract_messages, remove_columns=["task_type", "metadata"])
+eval_dataset = eval_stream.map(extract_messages, remove_columns=["task_type", "metadata"])
+
+# Convert streaming to regular dataset (materializes in memory)
+from datasets import Dataset
+print("Converting streaming dataset to memory...")
+train_dataset = Dataset.from_generator(lambda: (x for x in train_dataset))
+eval_dataset = Dataset.from_generator(lambda: (x for x in eval_dataset))
+
 print(f"Train: {len(train_dataset)} examples")
 print(f"Eval: {len(eval_dataset)} examples")