bert-sentiment-classifier / sentiment_classificator.py
santit96's picture
Initialize project and create a first sentiment classifier prototype with streamlit and a pretrained bert model
93109de
raw
history blame
897 Bytes
"""
Module to classify text into positive or negative sentiments
"""
import sys
import tensorflow as tf
from models.models import load_sentiments_model
sentiments_model = load_sentiments_model()
MAX_NEG = 0.4
MIN_POS = 0.6
def classify_sentiment(input_text: str) -> str:
"""
Receives a string and classifies it in positive, negative or none
"""
result = tf.sigmoid(sentiments_model(tf.constant([input_text])))
if result < MAX_NEG:
return "negative"
elif result > MIN_POS:
return "positive"
else:
return "-"
if __name__ == "__main__":
if len(sys.argv) < 2:
print(
f"Usage: python {sys.argv[0]} <text to classify>")
sys.exit(1)
# Get the input string from command line argument
input_text = sys.argv[1]
sentiment = classify_sentiment(input_text)
print("Sentiment of the sentence: ", sentiment)