Update train_model.py

train_model.py

@@ -9,18 +9,18 @@ from transformers import (
     DataCollatorForLanguageModeling
 )
 
-# ✅ Load token from secret named "data3"
+# ✅ 1. Load Hugging Face token from secret
 hf_token = os.environ.get("data3")
 if not hf_token:
-    raise ValueError("Hugging Face token not found. Make sure 'data3' is added as a secret in your Space.")
+    raise ValueError("Hugging Face token not found. Make sure 'data3' is set in your Space secrets.")
 
-# ✅ Authenticate
+# ✅ 2. Authenticate
 login(token=hf_token)
 
-# ✅ Load your dataset
+# ✅ 3. Load structured caption dataset
 df = pd.read_csv("custom_captions.csv")
 
-# ✅ Build training prompt
+# ✅ 4. Build training prompt
 def to_prompt(row):
     return f"""Photo: {row['photo']}
 Scenery: {row['scenery']}
@@ -29,38 +29,38 @@ Vibe to convey: {row['desired_vibe']}
 Captions user likes: "{row['liked_captions']}"
 """
 
-# ✅ Create final training text
+# ✅ 5. Combine prompt + final caption
 df["text"] = df.apply(lambda row: to_prompt(row) + row["caption"], axis=1)
 
-# ✅ Convert to Hugging Face Dataset
+# ✅ 6. Convert to Hugging Face Dataset
 dataset = Dataset.from_pandas(df[["text"]])
 
-# ✅ Load tokenizer and model
+# ✅ 7. Load tokenizer and base model
 tokenizer = AutoTokenizer.from_pretrained("gpt2")
-tokenizer.pad_token = tokenizer.eos_token  # Prevents padding issues
+tokenizer.pad_token = tokenizer.eos_token  # Avoids padding issues
 model = AutoModelForCausalLM.from_pretrained("gpt2")
 
-# ✅ Tokenization
+# ✅ 8. Tokenize training text
 def tokenize(example):
     return tokenizer(example["text"], truncation=True, padding="max_length", max_length=64)
 
 tokenized = dataset.map(tokenize, batched=True)
 
-# ✅ Data collator
+# ✅ 9. Data collator (for causal LM)
 data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 
-# ✅ TrainingArguments with your real username
+# ✅ 10. Training configuration — using your model repo
 training_args = TrainingArguments(
     output_dir="./results",
     per_device_train_batch_size=2,
     num_train_epochs=1,
     save_strategy="epoch",
     push_to_hub=True,
-    hub_model_id="yuvalon/caption-crafter-model",  # ✅ <-- This is your correct username/namespace
+    hub_model_id="yuvalon/caption-crafter-model",  # ✅ Match the repo you created manually
     hub_token=hf_token
 )
 
-# ✅ Trainer
+# ✅ 11. Trainer setup
 trainer = Trainer(
     model=model,
     args=training_args,
@@ -69,6 +69,6 @@ trainer = Trainer(
     data_collator=data_collator
 )
 
-# ✅ Train and push
+# ✅ 12. Train and upload to Hugging Face
 trainer.train()
 trainer.push_to_hub()