File size: 4,606 Bytes
93c5df6 | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | """
Training script for Restaurant Aspect Classifier.
Fine-tunes DistilBERT on Yelp reviews with rule-based aspect labeling.
"""
from datasets import load_dataset
from transformers import (
AutoModelForSequenceClassification,
AutoTokenizer,
Trainer,
TrainingArguments,
)
# Aspect categories
ASPECTS = ["FOOD", "SERVICE", "HYGIENE", "PARKING", "CLEANLINESS"]
def label_aspects(text: str, label: int) -> list[float]:
"""
Create aspect labels from review text and overall sentiment.
Args:
text: Review text
label: Overall sentiment (0=negative, 1=positive)
Returns:
List of 5 scores (0.0-1.0) for each aspect
"""
base_score = float(label)
scores = [base_score] * 5 # Initialize all aspects with base score
text_lower = text.lower()
# HYGIENE (index 2)
if any(word in text_lower for word in ["dirty", "filthy", "gross", "disgusting"]):
scores[2] = 0.0
elif any(word in text_lower for word in ["clean", "spotless", "sanitary"]):
scores[2] = 1.0
# FOOD (index 0)
if any(word in text_lower for word in ["delicious", "tasty", "amazing", "excellent"]):
scores[0] = 1.0
elif any(word in text_lower for word in ["terrible", "awful", "bland", "disgusting"]):
scores[0] = 0.0
# SERVICE (index 1)
if any(word in text_lower for word in ["friendly", "attentive", "helpful", "great service"]):
scores[1] = 1.0
elif any(word in text_lower for word in ["rude", "slow", "terrible service", "unfriendly"]):
scores[1] = 0.0
# PARKING (index 3)
if "parking" in text_lower:
if label == 0 or any(word in text_lower for word in ["no parking", "parking nightmare"]):
scores[3] = 0.0
else:
scores[3] = 1.0
# CLEANLINESS (index 4) - similar to hygiene
if any(word in text_lower for word in ["dirty", "messy", "unkempt"]):
scores[4] = 0.0
elif any(word in text_lower for word in ["clean", "tidy", "well-maintained"]):
scores[4] = 1.0
return scores
def main():
"""Train the model."""
print("🚀 Starting Restaurant Aspect Classifier Training\n")
# Load dataset
print("📥 Loading Yelp dataset (1500 samples)...")
dataset = load_dataset("yelp_polarity", split="train[:1500]", trust_remote_code=True)
print(f"✅ Loaded {len(dataset)} reviews\n")
# Create aspect labels
print("🏷️ Creating aspect labels from reviews...")
dataset = dataset.map(
lambda x: {"aspects": label_aspects(x["text"], x["label"])},
desc="Labeling aspects",
)
print("✅ Aspect labels created\n")
# Load tokenizer and model
print("🤖 Loading DistilBERT model and tokenizer...")
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained(
"distilbert-base-uncased",
num_labels=5,
problem_type="multi_label_classification",
)
print("✅ Model loaded\n")
# Tokenize dataset
print("✂️ Tokenizing text...")
def tokenize_function(examples):
return tokenizer(
examples["text"],
truncation=True,
padding="max_length",
max_length=256,
)
tokenized_dataset = dataset.map(
tokenize_function,
batched=True,
desc="Tokenizing",
)
# Prepare dataset for training
tokenized_dataset = tokenized_dataset.rename_column("aspects", "labels")
tokenized_dataset.set_format(
type="torch",
columns=["input_ids", "attention_mask", "labels"],
)
print("✅ Tokenization complete\n")
# Training arguments
training_args = TrainingArguments(
output_dir="./model",
num_train_epochs=2,
per_device_train_batch_size=8,
warmup_steps=50,
logging_steps=10,
save_strategy="epoch",
save_total_limit=1,
report_to="none",
push_to_hub=False,
)
# Create trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_dataset,
)
# Train
print("🎯 Training model (this will take ~45 minutes)...\n")
trainer.train()
print("\n✅ Training complete!\n")
# Save model
print("💾 Saving model and tokenizer...")
model.save_pretrained("./model")
tokenizer.save_pretrained("./model")
print("✅ Model saved to ./model/\n")
print("🎉 Training complete! You can now run: uvicorn main:app --reload")
if __name__ == "__main__":
main()
|