--- 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! ☕️ Buy Me a Coffee ## 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} }