Mental_Health_Analysis / SentBERT.py
AadityaMeshram's picture
Upload 10 files
6b2e2a5 verified
Raw
History Blame Contribute Delete
4.42 kB
# from transformers import TFBertForSequenceClassification, BertTokenizer
# import tensorflow as tf
# import numpy as np
# from sklearn.preprocessing import LabelEncoder
# # --- Configuration (ensure these match your training config) ---
# MODEL_NAME = 'bert-base-uncased'
# MAX_LEN = 128
# # Load saved model and tokenizer
# # !!! UPDATE THESE PATHS TO YOUR SAVED MODEL AND TOKENIZER LOCATION !!!
# model_path = r'C:\Users\adity\OneDrive\Desktop\VS Code Files\Project\ChatUI\bert_sentiment_model'
# tokenizer_path = r'C:\Users\adity\OneDrive\Desktop\VS Code Files\Project\ChatUI\bert_sentiment_tokenizer'
# model = TFBertForSequenceClassification.from_pretrained(model_path)
# tokenizer = BertTokenizer.from_pretrained(tokenizer_path)
# print('Model and Tokenizer loaded successfully.')
# if 'le' not in locals():
# le = LabelEncoder()
# # You would need to have access to the original 'labels' list or the 'classes_' array
# # to fit it correctly or reconstruct it. For this example, I'll use the known classes:
# original_classes = ['Anxiety', 'Bipolar', 'Depression', 'Normal', 'Personality disorder', 'Stress', 'Suicidal']
# le.fit(original_classes)
# print('LabelEncoder ready. Classes:', le.classes_)
# # --- Inference helper function ---
# # def predict_text(text):
# # enc = tokenizer(text, padding=True, truncation=True, max_length=MAX_LEN, return_tensors='tf')
# # logits = model(enc)[0]
# # pred = tf.argmax(logits, axis=1).numpy()[0]
# # proba = tf.nn.softmax(logits, axis=1).numpy()[0]
# # return {
# # 'label': le.inverse_transform([pred])[0],
# # 'label_id': int(pred),
# # 'probs': float(f"{proba[pred]:.2f}") # Probability of the top predicted label, formatted to 2 decimal places
# # }
# def predict_text(text):
# enc = tokenizer(text, padding=True, truncation=True, max_length=MAX_LEN, return_tensors='tf')
# logits = model(enc)[0]
# print("DEBUG LOGITS SHAPE:", logits.shape) # <-- ADD THIS
# pred = tf.argmax(logits, axis=1).numpy()[0]
# prob = tf.nn.softmax(logits, axis=1).numpy()[0][pred]
# # user_input_text = input("Please enter the text you want to analyze: ")
# # if user_input_text:
# # prediction = predict_text(user_input_text)
# # print(f"\nText: '{user_input_text}'")
# # print(f"Prediction: {prediction['label']}({prediction['probs']})")
# # #print(f"Probabilities: {prediction['probs']}")
# # else:
# # print("No text entered for analysis.")
# # --- Example usage ---
# # print(predict_text("I feel so sad and hopeless"))
import numpy as np
from transformers import TFBertForSequenceClassification, BertTokenizer
import tensorflow as tf
from sklearn.preprocessing import LabelEncoder
# Paths
model_path = r'C:\Users\adity\OneDrive\Desktop\VS Code Files\Project\ChatUI\bert_sentiment_model'
tokenizer_path = r'C:\Users\adity\OneDrive\Desktop\VS Code Files\Project\ChatUI\bert_sentiment_tokenizer'
# Load model & tokenizer
model = TFBertForSequenceClassification.from_pretrained(model_path)
tokenizer = BertTokenizer.from_pretrained(tokenizer_path)
# Label order (Very Important)
classes = ['Anxiety', 'Bipolar', 'Depression', 'Normal',
'Personality disorder', 'Stress', 'Suicidal']
# Label encoder (consistent with model output)
label_encoder = LabelEncoder()
label_encoder.fit(classes)
print("All model checkpoint layers were used when initializing TFBertForSequenceClassification.")
print(f"All the layers of TFBertForSequenceClassification were initialized from: {model_path}")
print("Model and Tokenizer loaded successfully.")
print("LabelEncoder ready. Classes:", classes)
# --------------------------
# PREDICT FUNCTION
# --------------------------
def predict_text(text: str):
"""Predicts label and confidence using the trained BERT model."""
# Convert text to input tensors
inputs = tokenizer(text, return_tensors='tf', padding=True, truncation=True)
# Run model
outputs = model(inputs)
logits = outputs.logits
# Softmax → probabilities
probs = tf.nn.softmax(logits, axis=1).numpy()[0]
# Get predicted class
pred_index = np.argmax(probs)
pred_label = classes[pred_index]
confidence = float(probs[pred_index])
return pred_label, confidence, probs