NanG01/bert-txn-classifier

bert-base-multilingual-cased
  → [CLS] token (768-d)
  → Dropout(0.3)
  → Linear(768 → 7)

import torch
import torch.nn as nn
from transformers import BertTokenizer, BertModel
from huggingface_hub import hf_hub_download

# Model class (must match training definition)
class BERTTxnClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.bert       = BertModel.from_pretrained("bert-base-multilingual-cased")
        self.drop       = nn.Dropout(0.3)
        self.classifier = nn.Linear(768, 7)

    def forward(self, input_ids, attention_mask):
        cls = self.bert(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0, :]
        return self.classifier(self.drop(cls))

CATEGORIES = ["Food", "Transport", "EMIs", "Entertainment", "Utilities", "Investments", "Other"]

# Download model + tokenizer
model_path = hf_hub_download(repo_id="NanG01/bert-txn-classifier", filename="bert_classifier.pt")

device    = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = BertTokenizer.from_pretrained("NanG01/bert-txn-classifier")
model     = BERTTxnClassifier()
model.load_state_dict(torch.load(model_path, map_location=device))
model.to(device).eval()

# Inference
def predict(text: str) -> dict:
    enc = tokenizer(text, max_length=64, padding="max_length",
                    truncation=True, return_tensors="pt")
    with torch.no_grad():
        probs = torch.softmax(
            model(enc["input_ids"].to(device), enc["attention_mask"].to(device)), dim=-1
        ).squeeze(0)
    pred = probs.argmax().item()
    return {"category": CATEGORIES[pred], "confidence": round(probs[pred].item(), 4)}

predict("SWIGGY ORDER PAYMENT")
# → {"category": "Food", "confidence": 0.9821}

predict("HDFC BANK PERSONAL LOAN EMI")
# → {"category": "EMIs", "confidence": 0.9743}

predict("OLA RIDE PAYMENT")
# → {"category": "Transport", "confidence": 0.9512}

predict("ZERODHA MUTUAL FUND")
# → {"category": "Investments", "confidence": 0.9301}

predict("बिजली बिल भुगतान")
# → {"category": "Utilities", "confidence": 0.9104}

predict("ਖਾਣੇ ਦਾ ਭੁਗਤਾਨ")
# → {"category": "Food", "confidence": 0.8932}