Update train.py

train.py

@@ -1,9 +1,15 @@
+---
+# ✅ train.py — Trains StarCoder 7B on your dataset
+
 from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer, DataCollatorForLanguageModeling
 import os
+import sys
+
+print("🔥 Training script started...", file=sys.stderr)
 
 # === CONFIG ===
-DATASET_PATH = "python_ai_dataset.jsonl"  # Your .jsonl file
+DATASET_PATH = "python_ai_dataset.jsonl"  # Must exist in Space root
 MODEL_ID = "bigcode/starcoderbase-7b"
 OUTPUT_DIR = "train_output"
 
@@ -16,15 +22,13 @@ model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True)
 
 # === Preprocessing ===
 def tokenize(example):
-    return tokenizer(example["prompt"] + "\n" + example["completion"],
-                     truncation=True, max_length=512)
+    return tokenizer(example["prompt"] + "\n" + example["completion"], truncation=True, max_length=512)
 
 tokenized_dataset = dataset.map(tokenize, remove_columns=["prompt", "completion"])
 
-# === Data Collator ===
 data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 
-# === Training Arguments ===
+# === Training Args ===
 training_args = TrainingArguments(
     output_dir=OUTPUT_DIR,
     overwrite_output_dir=True,
@@ -37,10 +41,10 @@ training_args = TrainingArguments(
     save_total_limit=2,
     fp16=True,
     bf16=False,
-    report_to="none",  # Prevent HF integration logs
+    report_to="none",
 )
 
-# === Trainer ===
+# === Train ===
 trainer = Trainer(
     model=model,
     args=training_args,
@@ -49,9 +53,8 @@ trainer = Trainer(
     data_collator=data_collator
 )
 
-# === Start Training ===
 trainer.train()
 
-# === Save Final Model ===
+# === Save ===
 trainer.save_model(OUTPUT_DIR)
 tokenizer.save_pretrained(OUTPUT_DIR)