Update app.py

app.py

@@ -2,15 +2,15 @@ import streamlit as st
 from transformers import DistilBertForSequenceClassification, DistilBertTokenizerFast
 import torch
 
-# Set page configuration as the very first Streamlit command
+# Set page configuration
 st.set_page_config(page_title="Spam Detection", page_icon="📧")
 
-# Load fine-tuned model and tokenizer
-model = DistilBertForSequenceClassification.from_pretrained("Anuj02003/Spam-classification-using-LLM")
-tokenizer = DistilBertTokenizerFast.from_pretrained("Anuj02003/Spam-classification-using-LLM")
+# Load the local fine-tuned model and tokenizer
+model_path = "./fine_tuned_model"
+model = DistilBertForSequenceClassification.from_pretrained(model_path)
+tokenizer = DistilBertTokenizerFast.from_pretrained(model_path)
 
-
-# Function to predict whether a message is spam or not
+# Function to predict spam
 def predict_spam(text):
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
     with torch.no_grad():
@@ -19,17 +19,18 @@ def predict_spam(text):
         prediction = torch.argmax(logits, dim=-1).item()
     return "Spam" if prediction == 1 else "Not Spam"
 
+# Streamlit app
 def main():
     st.title("Spam Detection")
     st.write("This is a Spam Detection App using a fine-tuned DistilBERT model.")
 
-    # Input text box for the user
+    # Input text box
     message = st.text_area("Enter message to classify as spam or not:")
 
     if st.button("Predict"):
         if message:
             prediction = predict_spam(message)
-            st.write(f"The message is: {prediction}")
+            st.write(f"The message is: **{prediction}**")
         else:
             st.write("Please enter a message to classify.")