Viewing File: /home/ubuntu/codegamaai-test/broker_bot/help_apis/intent_class.py

import torch
import torch.nn.functional as F
import pandas as pd
from transformers import BertTokenizer, BertForSequenceClassification
from sklearn.model_selection import train_test_split
from torch.utils.data import Dataset
import os

output_dir = os.environ['INTENT_MODEL_DIR']
checkpoint_dir = os.path.join(os.environ['INTENT_MODEL_DIR'], "checkpoint-200")

# Load the model and tokenizer from the checkpoint
model = BertForSequenceClassification.from_pretrained(checkpoint_dir)
tokenizer = BertTokenizer.from_pretrained(output_dir)

# Prediction function with probabilities
def predict_with_probabilities(text, model, tokenizer):
    model.eval()  # Set model to evaluation mode
    predictions = []
    probabilities = []

    with torch.no_grad():  # Disable gradient calculation
        
        inputs = tokenizer(text, padding=True, truncation=True, max_length=128, return_tensors="pt")
        outputs = model(**inputs)
        logits = outputs.logits
        probs = F.softmax(logits, dim=1)
        _, predicted_class = torch.max(probs, dim=1)
        predicted_prob = torch.max(probs, dim=1).values

        # create dictionary with predicted class original names instead of numbers
        di = {0: 'buy_stocks', 1: 'historical_stock_data', 2: 'sell_stocks', 3: 'stock_info', 4: 'view_stocks'}
        predicted_class = di[predicted_class.item()]
        response_dict = {
            "predicted_class": predicted_class,
            "predicted_prob": predicted_prob.item()
        }
        

    return response_dict

# Intent classification function
def classify_user_query(text):
    response = predict_with_probabilities(text, model, tokenizer)
    return response