Update app.py

app.py

@@ -1,73 +1,73 @@
-import torch
-import torch.nn as nn
-from transformers import BertTokenizer, BertModel
-from fastapi import FastAPI
-from pydantic import BaseModel
-import pickle
-
-# Load label encoders
-with open("main_category_encoder_5k.pkl", "rb") as f:
-    main_category_encoder = pickle.load(f)
-
-with open("sub_category_encoder_5k.pkl", "rb") as f:
-    sub_category_encoder = pickle.load(f)
-
-# Load tokenizer
-tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
-
-# Define the model
-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.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)
-
-    def forward(self, input_ids, attention_mask):
-        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
-        cls_embedding = outputs.last_hidden_state[:, 0, :]
-        main_logits = self.fc_main(cls_embedding)
-        main_pred = torch.softmax(main_logits, dim=1)
-        combined_input = torch.cat((cls_embedding, main_pred), dim=1)
-        sub_logits = self.fc_sub(combined_input)
-        return main_logits, sub_logits
-
-# Load trained model
-num_main_classes = len(main_category_encoder.classes_)
-num_sub_classes = len(sub_category_encoder.classes_)
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-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()
-
-# Initialize FastAPI
-app = FastAPI()
-
-# Define request body
-class TransactionInput(BaseModel):
-    description: str
-
-# Define predict function
-@app.post("/predict")
-def predict_category(transaction: TransactionInput):
-    tokens = tokenizer(transaction.description, return_tensors="pt", truncation=True, padding="max_length", max_length=64)
-    input_ids = tokens["input_ids"].to(device)
-    attention_mask = tokens["attention_mask"].to(device)
-
-    with torch.no_grad():
-        main_logits, sub_logits = model(input_ids, attention_mask)
-
-    main_category = torch.argmax(main_logits, dim=1).cpu().item()
-    sub_category = torch.argmax(sub_logits, dim=1).cpu().item()
-
-    return {
-        "description": transaction.description,
-        "main_category": main_category_encoder.inverse_transform([main_category])[0],
-        "sub_category": sub_category_encoder.inverse_transform([sub_category])[0]
-    }
-
-# Run the API (for local testing)
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+import torch
+import torch.nn as nn
+from transformers import BertTokenizer, BertModel
+from fastapi import FastAPI
+from pydantic import BaseModel
+import pickle
+
+# Load label encoders
+with open("main_category_encoder_5k.pkl", "rb") as f:
+    main_category_encoder = pickle.load(f)
+
+with open("sub_category_encoder_5k.pkl", "rb") as f:
+    sub_category_encoder = pickle.load(f)
+
+# Load tokenizer
+tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+
+# Define the model
+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.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)
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+        cls_embedding = outputs.last_hidden_state[:, 0, :]
+        main_logits = self.fc_main(cls_embedding)
+        main_pred = torch.softmax(main_logits, dim=1)
+        combined_input = torch.cat((cls_embedding, main_pred), dim=1)
+        sub_logits = self.fc_sub(combined_input)
+        return main_logits, sub_logits
+
+# Load trained model
+num_main_classes = len(main_category_encoder.classes_)
+num_sub_classes = len(sub_category_encoder.classes_)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+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()
+
+# Initialize FastAPI
+app = FastAPI()
+
+# Define request body
+class TransactionInput(BaseModel):
+    description: str
+
+# Define predict function
+@app.post("/predict")
+def predict_category(transaction: TransactionInput):
+    tokens = tokenizer(transaction.description, return_tensors="pt", truncation=True, padding="max_length", max_length=64)
+    input_ids = tokens["input_ids"].to(device)
+    attention_mask = tokens["attention_mask"].to(device)
+
+    with torch.no_grad():
+        main_logits, sub_logits = model(input_ids, attention_mask)
+
+    main_category = torch.argmax(main_logits, dim=1).cpu().item()
+    sub_category = torch.argmax(sub_logits, dim=1).cpu().item()
+
+    return {
+        "description": transaction.description,
+        "main_category": main_category_encoder.inverse_transform([main_category])[0],
+        "sub_category": sub_category_encoder.inverse_transform([sub_category])[0]
+    }
+
+# Run the API (for local testing)
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)