Update app.py

app.py

@@ -2,10 +2,13 @@ import gradio as gr
 import torch
 from transformers import DistilBertTokenizer
 from model import TransformerSentimentModel  # Import your custom class
+from huggingface_hub import hf_hub_download
+
 
 # 1. Load Tokenizer and Model
+REPO_ID = "Nefflymicn/amazon-sentiment-transformer"
 device = torch.device("cpu")  # Spaces use CPU by default for free tier
-tokenizer = DistilBertTokenizer.from_pretrained("distilbert/distilbert-base-uncased")
+tokenizer = DistilBertTokenizer.from_pretrained(REPO_ID, revision="v2.0")
 
 # Architecture must match your trained version exactly
 model = TransformerSentimentModel(
@@ -16,6 +19,7 @@ model = TransformerSentimentModel(
     num_layers=4,
     output_dim=2
 )
+weights_path = hf_hub_download(repo_id=REPO_ID, filename="pytorch_model.bin", revision="v2.0")
 model.load_state_dict(torch.load("pytorch_model.bin", map_location=device))
 model.eval()