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"

parvashah/bert-product-cat

My Custom BERT Model for Text Classification

Usage

parvashah/
bert-product-cat