Update README.md

README.md

@@ -74,44 +74,74 @@ Use the code below to get started with the model for product classification:
 import torch
 from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
 
-# Define the model repository name
-model_name = "Adnan-AI-Labs/DistilBERT-ProductClassifier"
-
-# 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")
+# 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