Update train_model.py

train_model.py

@@ -1,6 +1,6 @@
 import os
 import pandas as pd
-from huggingface_hub import login
+from huggingface_hub import login, whoami
 from datasets import Dataset
 
 from transformers import (
@@ -9,18 +9,21 @@ from transformers import (
     DataCollatorForLanguageModeling
 )
 
-# ✅ 1. Load Hugging Face token from secret
+# ✅ 1. Load Hugging Face token securely
 hf_token = os.environ.get("data3")
 if not hf_token:
-    raise ValueError("Hugging Face token not found. Make sure 'data3' is set in your Space secrets.")
+    raise ValueError("Missing HF token. Set 'data3' as a secret in your Space.")
 
-# ✅ 2. Authenticate
+# ✅ 2. Log in to Hugging Face
 login(token=hf_token)
 
-# ✅ 3. Load structured caption dataset
+# ✅ 3. Get your username dynamically
+username = whoami(token=hf_token)["name"]
+
+# ✅ 4. Load your caption data
 df = pd.read_csv("custom_captions.csv")
 
-# ✅ 4. Build training prompt
+# ✅ 5. Format training prompt
 def to_prompt(row):
     return f"""Photo: {row['photo']}
 Scenery: {row['scenery']}
@@ -29,38 +32,38 @@ Vibe to convey: {row['desired_vibe']}
 Captions user likes: "{row['liked_captions']}"
 """
 
-# ✅ 5. Combine prompt + final caption
+# ✅ 6. Create training text
 df["text"] = df.apply(lambda row: to_prompt(row) + row["caption"], axis=1)
 
-# ✅ 6. Convert to Hugging Face Dataset
+# ✅ 7. Convert to Hugging Face Dataset
 dataset = Dataset.from_pandas(df[["text"]])
 
-# ✅ 7. Load tokenizer and base model
+# ✅ 8. Load tokenizer & model
 tokenizer = AutoTokenizer.from_pretrained("gpt2")
-tokenizer.pad_token = tokenizer.eos_token  # Avoids padding issues
+tokenizer.pad_token = tokenizer.eos_token  # Avoid padding issues
 model = AutoModelForCausalLM.from_pretrained("gpt2")
 
-# ✅ 8. Tokenize training text
+# ✅ 9. Tokenize
 def tokenize(example):
     return tokenizer(example["text"], truncation=True, padding="max_length", max_length=64)
 
 tokenized = dataset.map(tokenize, batched=True)
 
-# ✅ 9. Data collator (for causal LM)
+# ✅ 10. Prepare data collator
 data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 
-# ✅ 10. Training configuration — using your model repo
+# ✅ 11. Set training arguments
 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",  # ✅ Match the repo you created manually
+    hub_model_id=f"{username}/caption-crafter-model",  # ✅ Dynamic and safe
     hub_token=hf_token
 )
 
-# ✅ 11. Trainer setup
+# ✅ 12. Trainer setup
 trainer = Trainer(
     model=model,
     args=training_args,
@@ -69,6 +72,6 @@ trainer = Trainer(
     data_collator=data_collator
 )
 
-# ✅ 12. Train and upload to Hugging Face
+# ✅ 13. Train and upload to Hugging Face Hub
 trainer.train()
 trainer.push_to_hub()