File size: 2,200 Bytes
d9dd063 85bcd80 d2f6ac4 d9dd063 85bcd80 a6f53e6 505b03b d9dd063 85bcd80 a6f53e6 85bcd80 d9dd063 505b03b d9dd063 d2f6ac4 d9dd063 505b03b d9dd063 85bcd80 d2f6ac4 505b03b d2f6ac4 505b03b d2f6ac4 9cc6b43 d2f6ac4 85bcd80 d2f6ac4 85bcd80 d2f6ac4 505b03b a6f53e6 d2f6ac4 505b03b d2f6ac4 505b03b d2f6ac4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | import os
import pandas as pd
from huggingface_hub import login, whoami
from datasets import Dataset
from transformers import (
AutoTokenizer, AutoModelForCausalLM,
TrainingArguments, Trainer,
DataCollatorForLanguageModeling
)
# ✅ 1. Load Hugging Face token securely
hf_token = os.environ.get("data3")
if not hf_token:
raise ValueError("Missing HF token. Set 'data3' as a secret in your Space or Colab.")
# ✅ 2. Log in to Hugging Face
login(token=hf_token)
# ✅ 3. Get your username dynamically
username = whoami(token=hf_token)["name"]
# ✅ 4. Load your caption data
df = pd.read_csv("custom_captions.csv")
# ✅ 5. Format prompt
def to_prompt(row):
return f"""Photo: {row['photo']}
Scenery: {row['scenery']}
Photo vibe: {row['photo_vibe']}
Vibe to convey: {row['desired_vibe']}
Captions user likes: "{row['liked_captions']}"
"""
# ✅ 6. Build full training text
df["text"] = df.apply(lambda row: to_prompt(row) + row["caption"], axis=1)
# ✅ 7. Convert to Hugging Face Dataset
dataset = Dataset.from_pandas(df[["text"]])
# ✅ 8. Load tokenizer and model (GPT-Neo)
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-125M")
tokenizer.pad_token = tokenizer.eos_token # Required fix for GPT-style models
model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M")
# ✅ 9. Tokenize the dataset
def tokenize(example):
return tokenizer(example["text"], truncation=True, padding="max_length", max_length=64)
tokenized = dataset.map(tokenize, batched=True)
# ✅ 10. Prepare data collator
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
# ✅ 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=f"{username}/caption-crafter-model", # Dynamic + personalized
hub_token=hf_token
)
# ✅ 12. Set up Trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized,
tokenizer=tokenizer,
data_collator=data_collator
)
# ✅ 13. Train and push to Hugging Face Hub
trainer.train()
trainer.push_to_hub()
|