|
""" |
|
This module provides a Streamlit web application for classifying Glassdoor |
|
reviews into sentiment categories using a fine-tuned BERT model. |
|
""" |
|
|
|
import numpy as np |
|
import streamlit as st |
|
import torch |
|
from transformers import BertTokenizer |
|
|
|
from config import (BERTIMBAU_FINETUNED_MODEL, BERTIMBAU_MODEL, |
|
SENTIMENT_MAPPING) |
|
from glassdoor_reviews_classifier import GlassdoorReviewsClassifier |
|
|
|
|
|
@st.cache_resource |
|
def load_model(): |
|
""" |
|
Loads the fine-tuned BERT model for sentiment classification. |
|
|
|
Returns: |
|
model (GlassdoorReviewsClassifier): The loaded model. |
|
""" |
|
try: |
|
model = GlassdoorReviewsClassifier().to(device) |
|
model.load_state_dict( |
|
torch.load(BERTIMBAU_FINETUNED_MODEL, map_location=device) |
|
) |
|
|
|
model.eval() |
|
|
|
return model |
|
except Exception as e: |
|
st.error(f"Error loading model: {e}") |
|
return None |
|
|
|
|
|
@st.cache_resource |
|
def load_tokenizer(): |
|
""" |
|
Loads the BERT tokenizer. |
|
|
|
Returns: |
|
tokenizer (BertTokenizer): The loaded tokenizer. |
|
""" |
|
return BertTokenizer.from_pretrained(BERTIMBAU_MODEL) |
|
|
|
|
|
def predict_sentiment(text): |
|
""" |
|
Predicts the sentiment of a given text. |
|
|
|
Args: |
|
text (str): The input text to classify. |
|
|
|
Returns: |
|
np.ndarray: The predicted probabilities for each sentiment class. |
|
""" |
|
outputs = [] |
|
encoded_text = tokenizer( |
|
text=text, |
|
max_length=512, |
|
add_special_tokens=True, |
|
return_token_type_ids=False, |
|
padding="max_length", |
|
truncation=True, |
|
return_attention_mask=True, |
|
return_tensors="pt", |
|
) |
|
|
|
input_ids = encoded_text["input_ids"].to(device) |
|
attention_mask = encoded_text["attention_mask"].to(device) |
|
|
|
with torch.no_grad(): |
|
output = model(input_ids, attention_mask) |
|
probabilities = torch.nn.functional.softmax(output, dim=1) |
|
outputs.append(probabilities.cpu().numpy()) |
|
|
|
return np.concatenate(outputs, axis=0) |
|
|
|
|
|
def get_sentiment_and_score(user_input): |
|
""" |
|
Gets the sentiment and score for a given user input. |
|
|
|
Args: |
|
user_input (str): The input text from the user. |
|
|
|
Returns: |
|
tuple: The predicted sentiment and its corresponding score. |
|
""" |
|
output_probabilities = predict_sentiment(user_input) |
|
|
|
predicted_index = np.argmax(output_probabilities) |
|
predicted_sentiment = SENTIMENT_MAPPING.get(predicted_index) |
|
|
|
sentiment_score = np.max(output_probabilities) |
|
|
|
return predicted_sentiment, sentiment_score |
|
|
|
|
|
if __name__ == "__main__": |
|
if torch.cuda.is_available(): |
|
device = torch.device("cuda") |
|
else: |
|
device = torch.device("cpu") |
|
|
|
model = load_model() |
|
tokenizer = load_tokenizer() |
|
|
|
user_input = st.text_input("Glassdoor Review Text") |
|
|
|
if user_input: |
|
predicted_sentiment, sentiment_score = get_sentiment_and_score(user_input) |
|
|
|
st.write( |
|
f"**Sentiment:** {predicted_sentiment}, **Score:** {sentiment_score:.4f} " |
|
) |
|
|
