bert-base-multilingual-cased
→ [CLS] token (768-d)
→ Dropout(0.3)
→ Linear(768 → 7)
1import torch
2import torch.nn as nn
3from transformers import BertTokenizer, BertModel
4from huggingface_hub import hf_hub_download
5
6# Model class (must match training definition)
7class BERTTxnClassifier(nn.Module):
8 def __init__(self):
9 super().__init__()
10 self.bert = BertModel.from_pretrained("bert-base-multilingual-cased")
11 self.drop = nn.Dropout(0.3)
12 self.classifier = nn.Linear(768, 7)
13
14 def forward(self, input_ids, attention_mask):
15 cls = self.bert(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0, :]
16 return self.classifier(self.drop(cls))
17
18CATEGORIES = ["Food", "Transport", "EMIs", "Entertainment", "Utilities", "Investments", "Other"]
19
20# Download model + tokenizer
21model_path = hf_hub_download(repo_id="NanG01/bert-txn-classifier", filename="bert_classifier.pt")
22
23device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
24tokenizer = BertTokenizer.from_pretrained("NanG01/bert-txn-classifier")
25model = BERTTxnClassifier()
26model.load_state_dict(torch.load(model_path, map_location=device))
27model.to(device).eval()
28
29# Inference
30def predict(text: str) -> dict:
31 enc = tokenizer(text, max_length=64, padding="max_length",
32 truncation=True, return_tensors="pt")
33 with torch.no_grad():
34 probs = torch.softmax(
35 model(enc["input_ids"].to(device), enc["attention_mask"].to(device)), dim=-1
36 ).squeeze(0)
37 pred = probs.argmax().item()
38 return {"category": CATEGORIES[pred], "confidence": round(probs[pred].item(), 4)}
1predict("SWIGGY ORDER PAYMENT")
2# → {"category": "Food", "confidence": 0.9821}
3
4predict("HDFC BANK PERSONAL LOAN EMI")
5# → {"category": "EMIs", "confidence": 0.9743}
6
7predict("OLA RIDE PAYMENT")
8# → {"category": "Transport", "confidence": 0.9512}
9
10predict("ZERODHA MUTUAL FUND")
11# → {"category": "Investments", "confidence": 0.9301}
12
13predict("बिजली बिल भुगतान")
14# → {"category": "Utilities", "confidence": 0.9104}
15
16predict("ਖਾਣੇ ਦਾ ਭੁਗਤਾਨ")
17# → {"category": "Food", "confidence": 0.8932}
