Distilbert ProductClassifier

🚀 DistilBERT-ProductClassifier

A fine - tuned DistilBERT model for product classification in the e - commerce domain, efficiently categorizing products and enhancing discoverability.

🚀 Quick Start

Use the code below to get started with the model for product classification:

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}'")

✨ Features

• Efficient Classification: The model is optimized for efficient text classification tasks and can work well with limited computational resources.
• Real - time Application: Leveraging the compact DistilBERT architecture, it is suitable for real - time applications in mobile and web environments.
• High Accuracy: Achieved an overall accuracy of 96.16% across multiple product categories.

📦 Installation

The code example above uses the transformers library. You can install it using the following command:

pip install transformers

💻 Usage Examples

Basic Usage

The code in the "Quick Start" section shows the basic usage of the model for product classification. It loads the model and tokenizer, and then predicts the categories for a set of product descriptions.

Advanced Usage

If you want to fine - tune the model further for your specific product categories, you can modify the hyperparameters in the training procedure. For example, you can change the learning rate, batch size, or the number of epochs.

# Example code for further fine - tuning (pseudo - code)
from transformers import DistilBertForSequenceClassification, DistilBertTokenizer, AdamW
import torch

model_name = 'Adnan-AI-Labs/DistilBERT-ProductClassifier'
tokenizer = DistilBertTokenizer.from_pretrained(model_name)
model = DistilBertForSequenceClassification.from_pretrained(model_name)

# Prepare your custom dataset
# ...

# Define optimizer and hyperparameters
optimizer = AdamW(model.parameters(), lr=1e-5)
epochs = 5
batch_size = 8

# Training loop
for epoch in range(epochs):
    # Training steps
    # ...
    pass

📚 Documentation

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

Model Sources

• Repository: 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.

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

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

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)

🔧 Technical Details

The model uses the DistilBERT architecture, which is a distilled version of BERT. It has a smaller number of parameters compared to BERT, making it more lightweight and faster to train and run. The fine - tuning process involves adding a text classification head on top of the pre - trained DistilBERT model and training it on the e - commerce product classification dataset.

📄 License

This project is licensed under the Apache 2.0 license.

📚 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}
}

📋 Model Information

⚠️ Important Note

The model was trained on e - commerce data and may not perform well on products or categories outside the provided training scope.

💡 Usage Tip

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.