Update app.py

app.py

@@ -1,8 +1,19 @@
 import streamlit as st
 import pandas as pd
+import re
+import string
+import nltk
+from nltk.corpus import stopwords
+from nltk.tokenize import word_tokenize
+from nltk.stem import WordNetLemmatizer
 from transformers import pipeline
 from PIL import Image
 
+# Download required NLTK data
+nltk.download('stopwords')
+nltk.download('punkt')
+nltk.download('wordnet')
+
 # Load Models
 news_classifier = pipeline("text-classification", model="Oneli/News_Classification")
 qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
@@ -19,27 +30,38 @@ label_mapping = {
 # Store classified article for QA
 context_storage = {"context": "", "bulk_context": "", "num_articles": 0}
 
+# Text Cleaning Functions
+def clean_text(text):
+    text = text.lower()
+    text = re.sub(f"[{string.punctuation}]", "", text)  # Remove punctuation
+    text = re.sub(r"[^a-zA-Z0-9\s]", "", text)  # Remove special characters
+    tokens = word_tokenize(text)
+    tokens = [word for word in tokens if word not in stopwords.words("english")]  # Remove stopwords
+    lemmatizer = WordNetLemmatizer()
+    tokens = [lemmatizer.lemmatize(word) for word in tokens]  # Lemmatize tokens
+    return " ".join(tokens)
+
 # Define the functions
 def classify_text(text):
-    result = news_classifier(text)[0]
+    cleaned_text = clean_text(text)
+    result = news_classifier(cleaned_text)[0]
     category = label_mapping.get(result['label'], "Unknown")
     confidence = round(result['score'] * 100, 2)
 
     # Store context for QA
-    context_storage["context"] = text
+    context_storage["context"] = cleaned_text
 
     return category, f"Confidence: {confidence}%"
 
-def classify_csv(file_path):
+def classify_csv(file):
     try:
-        df = pd.read_csv(file_path, encoding="utf-8")
-
-        # Automatically detect the column containing text
+        df = pd.read_csv(file, encoding="utf-8")
         text_column = df.columns[0]  # Assume first column is the text column
-
-        df["Encoded Prediction"] = df[text_column].apply(lambda x: news_classifier(str(x))[0]['label'])
+        
+        df[text_column] = df[text_column].astype(str).apply(clean_text)  # Clean text column
+        df["Encoded Prediction"] = df[text_column].apply(lambda x: news_classifier(x)[0]['label'])
         df["Decoded Prediction"] = df["Encoded Prediction"].map(label_mapping)
-        df["Confidence"] = df[text_column].apply(lambda x: round(news_classifier(str(x))[0]['score'] * 100, 2))
+        df["Confidence"] = df[text_column].apply(lambda x: round(news_classifier(x)[0]['score'] * 100, 2))
 
         # Store all text as a single context for QA
         context_storage["bulk_context"] = " ".join(df[text_column].dropna().astype(str).tolist())
@@ -47,30 +69,26 @@ def classify_csv(file_path):
 
         output_file = "output.csv"
         df.to_csv(output_file, index=False)
-
         return df, output_file
     except Exception as e:
         return None, f"Error: {str(e)}"
 
 def chatbot_response(history, user_input, source):
     user_input = user_input.lower()
-
-    # Select context based on source toggle
     context = context_storage["context"] if source == "Single Article" else context_storage["bulk_context"]
     num_articles = context_storage["num_articles"]
-
+    
     if "number of articles" in user_input or "how many articles" in user_input:
         answer = f"There are {num_articles} articles in the uploaded CSV."
         history.append([user_input, answer])
         return history, ""
-
+    
     if context:
         result = qa_pipeline(question=user_input, context=context)
         answer = result["answer"]
         history.append([user_input, answer])
         return history, ""
-
-    # Default responses if no context is available
+    
     responses = {
         "hello": "👋 Hello! How can I assist you with news today?",
         "hi": "😊 Hi there! What do you want to know about news?",
@@ -78,16 +96,12 @@ def chatbot_response(history, user_input, source):
         "thank you": "🙏 You're welcome! Let me know if you need anything else.",
         "news": "📰 I can classify news into Business, Sports, Politics, and more!",
     }
-
-    response = responses.get(user_input,
-                             "🤔 I'm here to help with news classification and general info. Ask me about news topics!")
+    response = responses.get(user_input, "🤔 I'm here to help with news classification and general info. Ask me about news topics!")
     history.append([user_input, response])
     return history, ""
 
 # Streamlit App Layout
 st.set_page_config(page_title="News Classifier", page_icon="📰")
-
-# Load Cover Image
 cover_image = Image.open("cover.png")  # Ensure this image exists
 st.image(cover_image, caption="News Classifier 📢", use_column_width=True)
 
@@ -128,4 +142,4 @@ if st.button("✉ Send"):
     st.write("*Chatbot Response:*")
     for q, a in history:
         st.write(f"*Q:* {q}")
-        st.write(f"*A:* {a}")
+        st.write(f"*A:* {a}")