Update app.py

app.py

@@ -5,6 +5,7 @@ from fastapi import FastAPI, Request
 from pydantic import BaseModel
 import pickle
 import logging
+import os
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -20,9 +21,10 @@ except FileNotFoundError as e:
     logger.error(f"Failed to load label encoders: {e}")
     raise
 
-# Load tokenizer
+# Load tokenizer from local directory
 try:
-    tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+    tokenizer = BertTokenizer.from_pretrained("./tokenizer")
+    logger.info("Tokenizer loaded successfully")
 except Exception as e:
     logger.error(f"Failed to load tokenizer: {e}")
     raise
@@ -31,7 +33,7 @@ except Exception as e:
 class BERTFNN(nn.Module):
     def __init__(self, num_main_classes, num_sub_classes):
         super(BERTFNN, self).__init__()
-        self.bert = BertModel.from_pretrained("bert-base-uncased")
+        self.bert = BertModel.from_pretrained("bert-base-uncased", cache_dir="./cache")
         self.fc_main = nn.Linear(self.bert.config.hidden_size, num_main_classes)
         self.fc_sub = nn.Linear(self.bert.config.hidden_size + num_main_classes, num_sub_classes)
 
@@ -52,6 +54,7 @@ try:
     model = BERTFNN(num_main_classes, num_sub_classes).to(device)
     model.load_state_dict(torch.load("expense_categorization_5k.pth", map_location=device))
     model.eval()
+    logger.info("Model loaded successfully")
 except FileNotFoundError as e:
     logger.error(f"Failed to load model weights: {e}")
     raise