Update README.md

README.md

@@ -71,19 +71,46 @@ For use cases involving other languages or highly specialized product categories
 Use the code below to get started with the model for product classification:
 
 ```python
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
+import torch
+from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
 
-# Load the model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("Adnan-AI-Labs/DistilBERT-ProductClassifier")
-model = AutoModelForSequenceClassification.from_pretrained("Adnan-AI-Labs/DistilBERT-ProductClassifier")
+# Define the model repository name
+model_name = "Adnan-AI-Labs/DistilBERT-ProductClassifier"
 
-# Create a pipeline for text classification
-classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)
+# Load the tokenizer and model from the Hugging Face Hub
+try:
+    # Load tokenizer
+    tokenizer = DistilBertTokenizer.from_pretrained(model_name, use_fast=True)
+
+    # Load model, forcing the download to avoid any cached version
+    model = DistilBertForSequenceClassification.from_pretrained(model_name)
+
+    print("Model and tokenizer loaded successfully.")
+
+except Exception as e:
+    print(f"An error occurred while loading the model: {e}")
+    exit()
+
+# Test the model with some sample inputs
+sample_texts = [
+    "estar s20 single uk sim free mobile phone red",
+    "cello c40227dvbt2 40 full hd black led tv",
+]
+
+# Prepare the inputs for the model
+inputs = tokenizer(sample_texts, padding=True, truncation=True, return_tensors="pt")
+
+# Make predictions
+with torch.no_grad():
+    outputs = model(**inputs)
+
+# Get the predicted class indices
+predictions = torch.argmax(outputs.logits, dim=1)
+
+# Print out the predictions
+for text, pred in zip(sample_texts, predictions):
+    print(f"Text: {text} \nPredicted Class: {pred.item()}\n")
 
-# Example usage
-product_description = "High-resolution digital camera with 20MP sensor."
-result = classifier(product_description)
-print(result)
 ```
 
 # Training Details