Update train.py

train.py

@@ -1,19 +1,40 @@
+import os
 from unsloth import FastLanguageModel
 from transformers import TrainingArguments, Trainer
+from datasets import load_dataset
+import torch
+
+# Validate environment variable
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    raise ValueError("HF_TOKEN environment variable not set")
 
 # Load quantized model
-model, tokenizer = FastLanguageModel.from_pretrained(
-    model_name="deepseek-ai/DeepSeek-V3-0324",
-    dtype=torch.bfloat16,
-    load_in_4bit=True,  # Or 2.71-bit
-    token=os.environ["HF_TOKEN"]
-)
-FastLanguageModel.for_training(model)
+try:
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name="deepseek-ai/DeepSeek-V3",
+        dtype=torch.bfloat16,
+        load_in_4bit=True,
+        token=HF_TOKEN
+    )
+    FastLanguageModel.for_training(model)
+except Exception as e:
+    raise RuntimeError(f"Failed to load model: {str(e)}")
+
+# Load and prepare dataset (example - replace with your actual dataset)
+try:
+    dataset = load_dataset("imdb")  # Example dataset
+    tokenized_dataset = dataset.map(
+        lambda x: tokenizer(x["text"], truncation=True, padding="max_length"),
+        batched=True
+    )
+except Exception as e:
+    raise RuntimeError(f"Failed to load/prepare dataset: {str(e)}")
 
 # Training arguments
 training_args = TrainingArguments(
     output_dir="/app/checkpoints",
-    per_device_train_batch_size=4,  # Adjust for A100 40GB/80GB
+    per_device_train_batch_size=4,
     per_device_eval_batch_size=4,
     num_train_epochs=2,
     learning_rate=2e-5,
@@ -23,26 +44,11 @@ training_args = TrainingArguments(
     eval_steps=500,
     logging_dir="/app/logs",
     logging_steps=100,
-    fp16=False,  # bfloat16 for A100
+    fp16=False,
+    bf16=True,
     deepspeed="/app/ds_config.json"
 )
 
-# DeepSpeed config
-with open("/app/ds_config.json", "w") as f:
-    f.write('''
-    {
-        "fp16": {"enabled": false},
-        "bf16": {"enabled": true},
-        "zero_optimization": {
-            "stage": 3,
-            "offload_optimizer": {"device": "cpu"},
-            "offload_param": {"device": "cpu"}
-        },
-        "train_batch_size": "auto",
-        "gradient_accumulation_steps": 4
-    }
-    ''')
-
 # Initialize trainer
 trainer = Trainer(
     model=model,
@@ -52,7 +58,10 @@ trainer = Trainer(
 )
 
 # Train
-trainer.train()
+try:
+    trainer.train()
+except Exception as e:
+    raise RuntimeError(f"Training failed: {str(e)}")
 
 # Save model
 model.save_pretrained("/app/fine_tuned_model")