My Custom BERT Model for Text Classification
This model is fine-tuned on a custom dataset for text classification.
Usage
import torch
from transformers import BertTokenizer, BertForSequenceClassification
import pickle
import os
# Get the current directory of the script
current_directory = os.path.dirname(os.path.abspath(__file__))
# Define the relative path to the saved model
pkl_path = os.path.join(current_directory, 'label_encoder.pkl')
# Load the model using the relative path
model = BertForSequenceClassification.from_pretrained(current_directory)
tokenizer = BertTokenizer.from_pretrained(current_directory)
# Load the label encoder
with open(pkl_path, 'rb') as le_file:
le = pickle.load(le_file)
def predict_phase(query_description):
# Tokenize the input description
inputs = tokenizer(query_description, return_tensors='pt', truncation=True, padding=True)
# Get the model's predictions
with torch.no_grad():
logits = model(**inputs).logits
# Convert logits to probabilities
probs = logits.softmax(dim=1)
# Get the predicted label and confidence
predicted_class = torch.argmax(probs, dim=1).item()
confidence = probs[0][predicted_class].item()
round_confidence = round(float(confidence),2)
label = le.inverse_transform([predicted_class])[0]
return label, round_confidence, "Bert-2023-10"