Update README.md

README.md

@@ -51,19 +51,66 @@ Use the code below to get started with EcommerceClassifier:
 
 ```python
 import torch
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoTokenizer, AutoModel
 import json
 import requests
 from PIL import Image
 from torchvision import transforms
 import urllib.request
+import torch.nn as nn
+
+# --- Define the Model ---
+class FineGrainedClassifier(nn.Module):
+    def __init__(self, num_classes=434):  # Updated to 434 classes
+        super(FineGrainedClassifier, self).__init__()
+        self.image_encoder = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
+        self.image_encoder.fc = nn.Identity()
+        self.text_encoder = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en')
+        self.classifier = nn.Sequential(
+            nn.Linear(2048 + 768, 1024),
+            nn.BatchNorm1d(1024),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(1024, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, num_classes)  # Updated to 434 classes
+        )
+    
+    def forward(self, image, input_ids, attention_mask):
+        image_features = self.image_encoder(image)
+        text_output = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask)
+        text_features = text_output.last_hidden_state[:, 0, :]
+        combined_features = torch.cat((image_features, text_features), dim=1)
+        output = self.classifier(combined_features)
+        return output
 
 # Load the label-to-class mapping from Hugging Face
 label_map_url = "https://huggingface.co/Maverick98/EcommerceClassifier/resolve/main/label_to_class.json"
 label_to_class = requests.get(label_map_url).json()
 
-# Load the model and tokenizer
-model = AutoModel.from_pretrained("Maverick98/EcommerceClassifier")
+# Load the custom model
+model = FineGrainedClassifier(num_classes=len(label_to_class))
+checkpoint_url = f"https://huggingface.co/Maverick98/EcommerceClassifier/resolve/main/model_checkpoint.pth"
+checkpoint = torch.hub.load_state_dict_from_url(checkpoint_url, map_location=torch.device('cpu'))
+
+# Clean up the state dictionary
+state_dict = checkpoint.get('model_state_dict', checkpoint)
+new_state_dict = {}
+for k, v in state_dict.items():
+    if k.startswith("module."):
+        new_key = k[7:]  # Remove "module." prefix
+    else:
+        new_key = k
+
+    # Check if the new_key exists in the model's state_dict, only add if it does
+    if new_key in model.state_dict():
+        new_state_dict[new_key] = v
+
+model.load_state_dict(new_state_dict)
+
+# Load the tokenizer from Jina
 tokenizer = AutoTokenizer.from_pretrained("jinaai/jina-embeddings-v2-base-en")
 
 # Define image preprocessing
@@ -124,13 +171,6 @@ print("Prediction Results:")
 for class_name, prob in results.items():
     print(f"Class: {class_name}, Probability: {prob}")
 
-# Map the top 3 indices to class names
-top3_classes = [label_to_class[str(idx.item())] for idx in top3_indices[0]]
-
-# Output the class names and their probabilities
-for i in range(3):
-    print(f"Class: {top3_classes[i]}, Probability: {top3_probabilities[0][i].item()}")
-
 ```
 
 # Training Details