Update app.py

app.py

@@ -1,36 +1,32 @@
-import streamlit as st 
-import transformers
-import torch 
+import streamlit as st
+import torch
 from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
 
-# Load tokenizer and model
-tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
-model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased')
+# Load the model and tokenizer
+@st.cache_resource  # Cache model for efficiency
+def load_model():
+    tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
+    model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased')
+    return tokenizer, model
 
-# Define a function to preprocess user input
-def preprocess_input(text):
-  encoded_input = tokenizer(text, return_tensors='pt')
-  return encoded_input
+tokenizer, model = load_model()
 
-# Define a function to generate response based on user input
-def generate_response(user_input):
-  encoded_input = preprocess_input(user_input)
-  outputs = model(**encoded_input)
-  # Extract relevant information from model outputs (e.g., predicted class)
-  # Based on the extracted information, formulate a response using predefined responses or logic
-  response = "I'm still under development, but I understand you said: {}".format(user_input)
-  return response
+# Input/Output areas
+st.title("Simple Sentiment Chatbot")
+user_input = st.text_input("Enter your message:")
 
-# Start the chat loop
-while True:
-  # Get user input
-  uinput = st.text_input("Enter your input: ")
+# Preprocess and generate response when the user hits Enter
+if user_input:
+    if user_input.lower() == "quit":
+        st.stop()  # End the Streamlit app
 
-  if uinput.lower() == "quit": 
-    break
+    encoded_input = tokenizer(user_input, return_tensors='pt')
+    outputs = model(**encoded_input)
+    logits = outputs.logits
+    predicted_class_id = logits.argmax(-1).item()
 
-    
-  # Generate response based on user input
-      
-  bot_response = generate_response(line)  # Assuming generate_response is defined elsewhere
-  print("Bot:", bot_response)
+    # Example sentiment mapping (you might have your own)
+    sentiment_map = {0: "negative", 1: "neutral", 2: "positive"}
+    sentiment = sentiment_map[predicted_class_id]
+
+    st.write(f"Bot: I sense a {sentiment} sentiment in your message.")