Upload train_survival.py with huggingface_hub

train_survival.py

@@ -44,6 +44,18 @@ model = AutoModelForCausalLM.from_pretrained(
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 tokenizer.pad_token = tokenizer.eos_token
 
+# MANUAL FORMATTING
+# We do this manually to avoid SFTTrainer batching issues
+def format_row(example):
+    instruction = example['instruction']
+    response = example['response']
+    # Qwen/Llama chat template format
+    text = f"<|im_start|>user\n{instruction}<|im_end|>\n<|im_start|>assistant\n{response}<|im_end|>{tokenizer.eos_token}"
+    return {"text": text}
+
+# Apply formatting manually
+dataset = dataset.map(format_row)
+
 # LoRA
 peft_config = LoraConfig(
     r=16,
@@ -67,32 +79,17 @@ training_args = SFTConfig(
     fp16=True,
     packing=False,
     max_length=1024,
-    dataset_text_field="text"
+    dataset_text_field="text" # Now this field exists and is correct
 )
 
-def formatting_prompts_func(example):
-    output_texts = []
-    instructions = example['instruction']
-    responses = example['response']
-    
-    for i in range(len(instructions)):
-        if i >= len(responses): break
-        instruction = instructions[i]
-        response = responses[i]
-        if not instruction or not response: continue
-
-        text = f"<|im_start|>user\n{instruction}<|im_end|>\n<|im_start|>assistant\n{response}<|im_end|>"
-        output_texts.append(text)
-    return output_texts
-
 # Trainer
 trainer = SFTTrainer(
     model=model,
     train_dataset=dataset,
     peft_config=peft_config,
-    formatting_func=formatting_prompts_func,
     args=training_args,
-    processing_class=tokenizer, # CORRECTED: Using processing_class instead of tokenizer
+    processing_class=tokenizer,
+    # Removed formatting_func argument
 )
 
 print("Starting training...")
@@ -100,4 +97,4 @@ trainer.train()
 
 print("Pushing to hub...")
 trainer.push_to_hub()
-print("Done!")
+print("Done!")