Spaces:

YUVALON
/

TrialApp2

Runtime error

TrialApp2 / train_model.py

Update train_model.py

505b03b verified over 1 year ago

2.2 kB

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()

	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()