Update app.py

app.py

@@ -2,30 +2,39 @@
 import streamlit as st
 from transformers import pipeline
 
-# Caching the model pipeline
+# Caching the text classification models
 @st.cache_resource
-def load_pipeline():
-    return pipeline("text2text-generation", model="facebook/blenderbot_small-90M")
+def load_pipeline(model_name):
+    return pipeline("text-classification", model=model_name)
 
-# Initialize the model once using cache
-pipe = load_pipeline()
-
-# Initialize session state for conversation history and bot response
+# Initialize session state for conversation history, bot response, and selected model
 if 'conversation_history' not in st.session_state:
     st.session_state.conversation_history = ""
 if 'bot_response' not in st.session_state:
     st.session_state.bot_response = ""
+if 'selected_model' not in st.session_state:
+    st.session_state.selected_model = "distilbert/distilbert-base-uncased-finetuned-sst-2-english"
 
-def converse(user_message):
+def classify_text(user_message):
     # Update the conversation history
     st.session_state.conversation_history += f"User: {user_message}\n"
-    result = pipe(st.session_state.conversation_history)[0]['generated_text']
-    st.session_state.conversation_history += f"Bot: {result}\n"
+    pipe = load_pipeline(st.session_state.selected_model)
+    result = pipe(user_message)[0]  # pipe returns a list of results
+    st.session_state.conversation_history += f"Bot: {result['label']} (Score: {result['score']:.2f})\n"
     st.session_state.bot_response = result
     return result
 
 # Sidebar options
 st.sidebar.title("App Settings")
+
+# Model selection
+model_options = {
+    "DistilBERT Sentiment Analysis": "distilbert/distilbert-base-uncased-finetuned-sst-2-english",
+    "BERT Multilingual Sentiment Analysis": "nlptown/bert-base-multilingual-uncased-sentiment"
+}
+selected_model = st.sidebar.selectbox("Select model:", list(model_options.keys()))
+st.session_state.selected_model = model_options[selected_model]
+
 show_history = st.sidebar.checkbox("Show conversation history", value=True)
 character_limit = st.sidebar.slider("Set character limit for input:", min_value=50, max_value=500, value=200)
 
@@ -36,21 +45,21 @@ if st.sidebar.button("Reset Conversation"):
     st.sidebar.success("Conversation history cleared.")
 
 # Streamlit app layout
-st.title("🤖 AI Chatbot")
-st.subheader("Chat with an AI-powered bot!")
+st.title("🧠 Text Classification Bot")
+st.subheader("Classify your text with a sentiment analysis model!")
 
 # Input field with character limit
 user_message = st.text_input(f"Enter your message (max {character_limit} characters):", max_chars=character_limit)
 
-# Send button to generate bot response
-if st.button("Send"):
+# Send button to generate classification
+if st.button("Classify"):
     if user_message:
-        # Get response from the chatbot
-        bot_response = converse(user_message)
+        # Get classification from the selected model
+        classification_result = classify_text(user_message)
         
         # Display bot's response in a dedicated area
-        st.markdown("### Bot's Response")
-        st.success(bot_response)
+        st.markdown("### Classification Result")
+        st.success(f"**Label:** {classification_result['label']}\n**Score:** {classification_result['score']:.2f}")
         
         if show_history:
             # Display conversation history in a text area for better scrolling
@@ -58,12 +67,12 @@ if st.button("Send"):
             st.text_area("Conversation", value=st.session_state.conversation_history, height=250, max_chars=None)
     else:
         # Show a warning if no message is provided
-        st.warning("Please enter a message before sending.")
+        st.warning("Please enter a message before classifying.")
 
 # About section
 st.markdown("---")
 st.markdown("### About this App")
-st.info("This chatbot is powered by a pre-trained model from the Hugging Face Transformers library. You can chat with the bot, and the conversation history will be maintained during the session.")
+st.info("This app uses pre-trained models for sentiment analysis. You can select a model and enter text to see its classification and sentiment score.")
 
 st.sidebar.markdown("---")
 st.sidebar.write("Created by [Your Name](https://github.com/yourprofile)")