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

If you find this project useful, consider buying me a coffee to support further development! ☕️

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.

How to Get Started with the Model

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

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