|
--- |
|
language: |
|
- en |
|
tags: |
|
- text-classification |
|
- e-commerce |
|
- product-classification |
|
- distilbert |
|
license: apache-2.0 |
|
datasets: |
|
- Adnan-AI-Labs/products_categories_data |
|
model-index: |
|
- name: DistilBERT-ProductClassifier |
|
results: |
|
- task: |
|
type: text-classification |
|
name: Product Category Classification |
|
dataset: |
|
name: Product Classification and Categorization |
|
type: lakritidis/product-classification-and-categorization |
|
metrics: |
|
- type: accuracy |
|
value: 96.17 |
|
name: Accuracy |
|
base_model: |
|
- lxs1/DistilBertForSequenceClassification_6h_768dim |
|
- distilbert/distilbert-base-uncased |
|
pipeline_tag: text-classification |
|
--- |
|
|
|
# Model Card for DistilBERT-ProductClassifier |
|
|
|
This is a fine-tuned version of the DistilBERT model, specifically trained for product classification tasks within the e-commerce domain. The model distinguishes between various categories like "CPUs," "Digital Cameras," "Dishwashers," and more, making it useful for organizing and categorizing products in online retail platforms. |
|
|
|
## Model Details |
|
|
|
### Model Description |
|
|
|
The DistilBERT-ProductClassifier model is trained on an e-commerce dataset to classify products into specific categories. It is optimized for efficient text classification tasks and is designed to work well with limited computational resources. This model leverages the compact DistilBERT architecture, making it suitable for real-time applications, including mobile and web environments. |
|
|
|
- **Developed by:** Adnan AI Labs |
|
- **Model type:** DistilBERT fine-tuned for text classification |
|
- **Language:** English |
|
- **License:** Apache 2.0 |
|
- **Finetuned from model:** [DistilBERT](https://huggingface.co/distilbert-base-uncased) |
|
|
|
|
|
If you find this project useful, consider buying me a coffee to support further development! ☕️ |
|
|
|
<a href="https://buymeacoffee.com/adnanailabs"> |
|
<img src="https://cdn.buymeacoffee.com/buttons/v2/default-yellow.png" alt="Buy Me a Coffee"> |
|
</a> |
|
|
|
|
|
## Model Sources |
|
|
|
- **Repository:** [Adnan-AI-Labs/DistilBERT-ProductClassifier](https://huggingface.co/Adnan-AI-Labs/DistilBERT-ProductClassifier) |
|
|
|
## Uses |
|
|
|
### Direct Use |
|
|
|
The model is intended for classifying products in text-based product listings for e-commerce platforms. Users can use this model to categorize products automatically, reducing the need for manual tagging and improving product discoverability. |
|
|
|
### Out-of-Scope Use |
|
|
|
This model should not be used for tasks unrelated to product classification or outside the e-commerce context. It is not designed for sentiment analysis, general text generation, or other unrelated NLP tasks. |
|
|
|
## Bias, Risks, and Limitations |
|
|
|
The model was trained on e-commerce data and may not perform well on products or categories outside the provided training scope. Additionally, some categories may have less representation in the training data, potentially affecting classification accuracy for those classes. |
|
|
|
### Recommendations |
|
|
|
For use cases involving other languages or highly specialized product categories not included in the training data, further fine-tuning may be necessary. Users should validate the model's output before using it for high-stakes decisions. |
|
|
|
## How to Get Started with the Model |
|
|
|
Use the code below to get started with the model for product classification: |
|
|
|
```python |
|
import torch |
|
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification |
|
|
|
# Load the model and tokenizer from the Hugging Face Hub |
|
def load_model_and_tokenizer(model_name, num_labels): |
|
tokenizer = DistilBertTokenizer.from_pretrained(model_name) |
|
model = DistilBertForSequenceClassification.from_pretrained(model_name, num_labels=num_labels) |
|
model.eval() # Set the model to evaluation mode |
|
return model, tokenizer |
|
|
|
# Predict categories for the provided prompts |
|
def predict(model, tokenizer, prompts, category_mapping, device): |
|
model.to(device) |
|
inputs = tokenizer(prompts, padding=True, truncation=True, return_tensors='pt', max_length=128) |
|
|
|
with torch.no_grad(): |
|
input_ids = inputs['input_ids'].to(device) |
|
attention_mask = inputs['attention_mask'].to(device) |
|
|
|
outputs = model(input_ids, attention_mask=attention_mask) |
|
logits = outputs.logits |
|
predictions = torch.argmax(logits, dim=1).cpu().numpy() |
|
|
|
predicted_categories = [category_mapping[pred] for pred in predictions] |
|
return predicted_categories |
|
|
|
# Main execution block |
|
if __name__ == "__main__": |
|
# Define some example prompts for prediction |
|
prompts = [ |
|
"Intel Core i7 CPU", |
|
"Nikon D3500 Digital Camera", |
|
"Bosch Series 6 Dishwasher", |
|
"Samsung 32 inch Smart TV", |
|
"Apple iPhone 13" |
|
] |
|
|
|
# Create the category mapping based on provided comments |
|
category_mapping = { |
|
0: 'cpus', |
|
1: 'digital cameras', |
|
2: 'dishwashers', |
|
3: 'fridge freezers', |
|
4: 'microwaves', |
|
5: 'mobile phones', |
|
6: 'tvs', |
|
7: 'washing machines' |
|
} |
|
|
|
model_name = 'Adnan-AI-Labs/DistilBERT-ProductClassifier' |
|
|
|
# Load the model and tokenizer |
|
print(f"Loading model and tokenizer from Hugging Face Hub: {model_name}") |
|
model, tokenizer = load_model_and_tokenizer(model_name, len(category_mapping)) |
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
|
# Make predictions |
|
predicted_categories = predict(model, tokenizer, prompts, category_mapping, device) |
|
|
|
# Display the predictions |
|
for prompt, category in zip(prompts, predicted_categories): |
|
print(f"Prompt: '{prompt}' | Predicted Category: '{category}'") |
|
|
|
``` |
|
## Output |
|
Loading model and tokenizer from Hugging Face Hub: Adnan-AI-Labs/DistilBERT-ProductClassifier |
|
Prompt: 'Intel Core i7 CPU' | Predicted Category: 'cpus' |
|
Prompt: 'Nikon D3500 Digital Camera' | Predicted Category: 'digital cameras' |
|
Prompt: 'Bosch Series 6 Dishwasher' | Predicted Category: 'dishwashers' |
|
Prompt: 'Samsung 32 inch Smart TV' | Predicted Category: 'tvs' |
|
Prompt: 'Apple iPhone 13' | Predicted Category: 'mobile phones' |
|
|
|
# Training Details |
|
## Training Data |
|
The model was trained on an e-commerce dataset that includes various product categories such as CPUs, Digital Cameras, Dishwashers, Fridge Freezers, Microwaves, Mobile Phones, TVs, and Washing Machines. The data was preprocessed by removing duplicates, lowercasing, and tokenizing text. |
|
|
|
## Training Procedure |
|
1. Preprocessing: Text data was cleaned, lowercased, and tokenized. Product descriptions were truncated to 128 tokens for uniformity. |
|
2. Hyperparameters: Fine-tuning was conducted with a learning rate of 2e-5 and batch size of 16 for 3 epochs. |
|
3. Training Hardware: The model was trained on a single NVIDIA Tesla V100 GPU for approximately 3 hours. |
|
|
|
## Evaluation |
|
Testing Data, Factors & Metrics |
|
The model was evaluated on a separate test set of product descriptions, using precision, recall, and F1-score as the evaluation metrics. |
|
|
|
## Summary |
|
The model achieved an overall accuracy of 96.16%, with strong performance across multiple product categories. The F1-scores indicate that the model performs particularly well in the "CPUs" and "Digital Cameras" categories. |
|
|
|
## Technical Specifications |
|
Model Architecture and Objective |
|
The DistilBERT-ProductClassifier model utilizes the DistilBERT architecture, fine-tuned with a text classification head for e-commerce product categorization tasks. |
|
|
|
## Compute Infrastructure |
|
The model is optimized to run efficiently on CPUs and small GPUs, making it suitable for real-time applications. |
|
|
|
## Hardware |
|
This model requires a minimum of 4GB of RAM for efficient inference, and a modern CPU or basic GPU is recommended. |
|
|
|
## Software |
|
Transformers library: v4.26.0 |
|
Python version: 3.8 or higher |
|
CUDA [optional]: 10.2 or higher (if running on GPU) |
|
|
|
## Citation |
|
If you use this model, please cite as follows: |
|
|
|
|
|
@misc{distilbert_product_classifier, |
|
author = {Adnan AI Labs}, |
|
title = {DistilBERT-ProductClassifier for E-commerce}, |
|
year = {2024}, |
|
url = {https://huggingface.co/Adnan-AI-Labs/DistilBERT-ProductClassifier} |
|
} |