Upload train_survival.py with huggingface_hub

train_survival.py

@@ -39,21 +39,19 @@ model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     quantization_config=bnb_config,
     device_map="auto",
-    use_cache=False
+    use_cache=False,
+    torch_dtype=torch.float16 # Explicitly set float16
 )
 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
@@ -76,10 +74,11 @@ training_args = SFTConfig(
     logging_steps=10,
     push_to_hub=True,
     hub_model_id=OUTPUT_MODEL_ID,
-    fp16=True,
+    fp16=True, # Force FP16
+    bf16=False, # Disable BF16 explicitly
     packing=False,
     max_length=1024,
-    dataset_text_field="text" # Now this field exists and is correct
+    dataset_text_field="text"
 )
 
 # Trainer
@@ -89,7 +88,6 @@ trainer = SFTTrainer(
     peft_config=peft_config,
     args=training_args,
     processing_class=tokenizer,
-    # Removed formatting_func argument
 )
 
 print("Starting training...")
@@ -97,4 +95,4 @@ trainer.train()
 
 print("Pushing to hub...")
 trainer.push_to_hub()
-print("Done!")
+print("Done!")