File size: 1,427 Bytes
854f567 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
from typing import List, Dict
import numpy as np
import torch
from transformers import BertForSequenceClassification, BertTokenizer
from sklearn.preprocessing import LabelEncoder
from huggingface_hub import hf_hub_download
class CustomBertClassifier:
def __init__(self):
# Load model and tokenizer
self.model = BertForSequenceClassification.from_pretrained(".")
self.tokenizer = BertTokenizer.from_pretrained(".")
self.model.eval()
# Load label classes
label_path = hf_hub_download(repo_id="JaySenpai/bert-model", filename="label_classes.npy")
self.le = LabelEncoder()
self.le.classes_ = np.load(label_path, allow_pickle=True)
def __call__(self, inputs: str) -> List[Dict]:
# Tokenize input
inputs = self.tokenizer(inputs, return_tensors="pt", truncation=True, padding=True)
with torch.no_grad():
outputs = self.model(**inputs)
probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
probs = probs[0].tolist()
# Map to labels
results = []
for i, prob in enumerate(probs):
results.append({
"label": self.le.classes_[i],
"score": round(prob, 4)
})
# Sort by score descending
results = sorted(results, key=lambda x: x["score"], reverse=True)
return results
|