Update app.py

app.py

@@ -1,43 +1,31 @@
 import streamlit as st
-import pandas as pd
-import plotly.express as px
-from transformers import pipeline, AutoTokenizer
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-def main():
-    st.set_page_config(page_title="Sentiment Analysis App")
-    st.sidebar.title("Sentiment Analysis App")
-    st.sidebar.write("Choose a pre-trained model for sentiment analysis")
-    model_name = st.sidebar.selectbox("Select Model", ["bert-base-uncased", "distilbert-base-uncased", "roberta-base"])
-    text_input = st.text_input("Enter text for sentiment analysis")
+tokenizer = AutoTokenizer.from_pretrained("nlptown/bert-base-multilingual-uncased-sentiment")
+model = AutoModelForSequenceClassification.from_pretrained("nlptown/bert-base-multilingual-uncased-sentiment")
 
-    if st.button("Analyze"):
-        if text_input and isinstance(text_input, str):
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            encoded_input = tokenizer(text_input, return_tensors='pt', padding=True, truncation=True)
-            sentiment_classifier = pipeline("sentiment-analysis", model=model_name)
-            result = sentiment_classifier(encoded_input)[0]
-            st.write(f"Sentiment: {result['label']}")
-            st.write(f"Score: {round(result['score'], 4)}")
-            
-            # create a dataframe to hold the sentiment distribution
-            df_sentiments = pd.DataFrame({"sentiment": ["Positive", "Negative", "Neutral"], "count": [0, 0, 0]})
-            
-            # get the sentiment distribution of the text using the selected model
-            results = sentiment_classifier(encoded_input)
-            for r in results:
-                if r["label"] == "POSITIVE":
-                    df_sentiments.loc[0, "count"] = r["score"]
-                elif r["label"] == "NEGATIVE":
-                    df_sentiments.loc[1, "count"] = r["score"]
-                elif r["label"] == "NEUTRAL":
-                    df_sentiments.loc[2, "count"] = r["score"]
-            
-            # plot the sentiment distribution as a pie chart
-            fig = px.pie(df_sentiments, values='count', names='sentiment', hole=.4)
-            st.plotly_chart(fig)
-            
-        else:
-            st.warning("Please enter some valid text for sentiment analysis")
+st.title("Sentiment Analysis App")
 
-if __name__ == "__main__":
-    main()
+text = st.text_input("Enter text to analyze:")
+if st.button("Analyze"):
+    encoding = tokenizer.encode_plus(text, return_tensors="pt", padding=True, truncation=True)
+    input_ids = encoding["input_ids"]
+    attention_mask = encoding["attention_mask"]
+
+    with torch.no_grad():
+        output = model(input_ids, attention_mask)
+    prediction = int(torch.argmax(output.logits))
+
+    if prediction == 0:
+        st.write("Negative")
+    elif prediction == 1:
+        st.write("Neutral")
+    else:
+        st.write("Positive")
+
+    values = [output.logits[0][0].item(), output.logits[0][1].item(), output.logits[0][2].item()]
+    labels = ["Negative", "Neutral", "Positive"]
+    fig, ax = plt.subplots()
+    ax.bar(labels, values)
+    st.pyplot(fig)