A newer version of the Streamlit SDK is available:
1.45.1
title: Product Category Classification BERT
emoji: 🌍
colorFrom: green
colorTo: indigo
sdk: streamlit
sdk_version: 1.39.0
app_file: app.py
pinned: false
short_description: Identify Category based on Title and Description
Retail Product Classification Streamlit App
This Streamlit app is a product classification tool built using a BERT model fine-tuned on a retail product dataset. The model can classify products into one of 21 categories based on their title and description.
Dataset
The model was trained using the Kaggle Retail Product Classification dataset. The dataset consists of various product descriptions and their corresponding categories. The training goal was to classify products into 21 distinct categories.
Categories and Index Mapping
| Category | Index |
|---|---|
| Electronics | 0 |
| Sports & Outdoors | 1 |
| Cell Phones & Accessories | 2 |
| Automotive | 3 |
| Toys & Games | 4 |
| Tools & Home Improvement | 5 |
| Health & Personal Care | 6 |
| Beauty | 7 |
| Grocery & Gourmet Food | 8 |
| Office Products | 9 |
| Arts, Crafts & Sewing | 10 |
| Pet Supplies | 11 |
| Patio, Lawn & Garden | 12 |
| Clothing, Shoes & Jewelry | 13 |
| Baby | 14 |
| Musical Instruments | 15 |
| Industrial & Scientific | 16 |
| Baby Products | 17 |
| Appliances | 18 |
| All Beauty | 19 |
| All Electronics | 20 |
Model Training
The model used for this app is a BERT base model (bert-base-uncased) fine-tuned using the Hugging Face transformers library. The model was trained to classify products into the 21 categories listed above. The fine-tuning was carried out using the following training arguments:
'''python
training_args = TrainingArguments(
output_dir='./results',
evaluation_strategy='epoch',
save_strategy='epoch',
logging_strategy="steps",
logging_steps=10,
per_device_train_batch_size=32,
per_device_eval_batch_size=16,
num_train_epochs=6,
weight_decay=0.01,
learning_rate=2e-5,
lr_scheduler_type='cosine',
warmup_steps=250,
logging_dir='./logs',
report_to="tensorboard",
load_best_model_at_end=True,
save_total_limit=3,
gradient_accumulation_steps=2,
seed=42,
eval_accumulation_steps=10,
)'''
How to run the APP
Provide any Retail product title and Description in the given text boxes and click on classify product. The app would return the appropriate category