Fix model loading to use model_checkpoint.pth and add proper text encoding

app.py

@@ -1,31 +1,16 @@
 import gradio as gr
 import torch
 import torch.nn as nn
-from torchvision import models, transforms
+from torchvision import transforms
 from PIL import Image
 import json
 from huggingface_hub import hf_hub_download
-import os
-
-# Define the model architecture
-class FineGrainedClassifier(nn.Module):
-    def __init__(self, num_classes, text_dim=768):
-        super(FineGrainedClassifier, self).__init__()
-        self.resnet = models.resnet50(pretrained=False)
-        self.resnet.fc = nn.Identity()
-        self.text_fc = nn.Linear(text_dim, 1024)
-        self.fusion_fc = nn.Linear(2048 + 1024, num_classes)
-    
-    def forward(self, images, text_embeddings):
-        image_features = self.resnet(images)
-        text_features = self.text_fc(text_embeddings)
-        combined = torch.cat((image_features, text_features), dim=1)
-        output = self.fusion_fc(combined)
-        return output
+from transformers import AutoTokenizer, AutoModel
+from model import FineGrainedClassifier
 
 # Download model files
 try:
-    model_path = hf_hub_download(repo_id="Maverick98/EcommerceClassifier", filename="best_model.pth")
+    model_path = hf_hub_download(repo_id="Maverick98/EcommerceClassifier", filename="model_checkpoint.pth")
     label_path = hf_hub_download(repo_id="Maverick98/EcommerceClassifier", filename="label_to_class.json")
     
     with open(label_path, 'r') as f:
@@ -35,6 +20,12 @@ try:
     model = FineGrainedClassifier(num_classes=num_classes)
     model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
     model.eval()
+    
+    # Load text tokenizer
+    tokenizer = AutoTokenizer.from_pretrained("jinaai/jina-embeddings-v2-base-en", trust_remote_code=True)
+    text_encoder = AutoModel.from_pretrained("jinaai/jina-embeddings-v2-base-en", trust_remote_code=True)
+    text_encoder.eval()
+    
     model_loaded = True
 except Exception as e:
     print(f"Error loading model: {e}")
@@ -48,7 +39,7 @@ transform = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
-def classify_product(image, text):
+def classify_product(image, text=""):
     if not model_loaded:
         return {"Error": "Model not loaded properly"}
     
@@ -60,8 +51,14 @@ def classify_product(image, text):
         img = Image.fromarray(image).convert('RGB')
         img_tensor = transform(img).unsqueeze(0)
         
-        # For now, use dummy text embeddings (in real implementation, use Jina embeddings)
-        text_embeddings = torch.zeros(1, 768)
+        # Process text
+        if text.strip():
+            inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
+            with torch.no_grad():
+                text_embeddings = text_encoder(**inputs).last_hidden_state.mean(dim=1)
+        else:
+            # Use zero embeddings if no text provided
+            text_embeddings = torch.zeros(1, 768)
         
         # Get predictions
         with torch.no_grad():
@@ -86,14 +83,12 @@ demo = gr.Interface(
     fn=classify_product,
     inputs=[
         gr.Image(label="Product Image"),
-        gr.Textbox(label="Product Description (optional)", placeholder="Enter product title or description...", lines=2)
+        gr.Textbox(label="Product Description (optional)", placeholder="Enter product title or description...", lines=2, value="")
     ],
     outputs=gr.Label(label="Classification Results", num_top_classes=10),
     title="🛍️ E-Commerce Product Classifier",
-    description="Fast and accurate e-commerce product classification. Upload a product image to classify it into the appropriate category.",
-    examples=[
-        ["https://raw.githubusercontent.com/gradio-app/gradio/main/guides/assets/demo_files/T-shirt.png", "Cotton T-Shirt"],
-    ],
+    description="Fast and accurate e-commerce product classification powered by EcommerceClassifier. Upload a product image and optionally provide a text description to classify it into the appropriate category.",
+    examples=[],
     theme="soft"
 )