Update app.py

app.py

@@ -1,21 +1,8 @@
 import streamlit as st
 import pandas as pd
-import string
-import re
-import nltk
-from nltk.corpus import stopwords
-from nltk.tokenize import word_tokenize
-from nltk.stem import WordNetLemmatizer
-from nltk.corpus import wordnet
 from transformers import pipeline
 from PIL import Image
 
-# Download necessary NLTK data
-nltk.download("stopwords")
-nltk.download("punkt")
-nltk.download("wordnet")
-nltk.download("averaged_perceptron_tagger")
-
 # Load Models
 news_classifier = pipeline("text-classification", model="Oneli/News_Classification")
 qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
@@ -32,63 +19,79 @@ label_mapping = {
 # Store classified article for QA
 context_storage = {"context": "", "bulk_context": "", "num_articles": 0}
 
-# Preprocessing functions
-def remove_punctuation(text):
-    return text.translate(str.maketrans('', '', string.punctuation))
-
-def remove_special_characters(text):
-    return re.sub(r'[^A-Za-z\s]', '', text)
-
-def remove_stopwords(text):
-    stop_words = set(stopwords.words('english'))
-    return " ".join([word for word in text.split() if word not in stop_words])
-
-def tokenize_text(text):
-    return word_tokenize(text)
-
-def lemmatize_tokens(tokens):
-    lemmatizer = WordNetLemmatizer()
-    wordnet_map = {"N": wordnet.NOUN, 'V': wordnet.VERB, 'J': wordnet.ADJ, 'R': wordnet.ADV}
-    return [lemmatizer.lemmatize(token, wordnet_map.get(nltk.pos_tag([token])[0][1][0].upper(), wordnet.NOUN)) for token in tokens]
-
-def preprocess_text(text):
-    text = text.lower()
-    text = remove_punctuation(text)
-    text = remove_special_characters(text)
-    text = remove_stopwords(text)
-    tokens = tokenize_text(text)
-    tokens = lemmatize_tokens(tokens)
-    return " ".join(tokens)
-
-# Classification functions
+# Define the functions
 def classify_text(text):
-    cleaned_text = preprocess_text(text)
-    result = news_classifier(cleaned_text)[0]
+    result = news_classifier(text)[0]
     category = label_mapping.get(result['label'], "Unknown")
     confidence = round(result['score'] * 100, 2)
-    context_storage["context"] = cleaned_text
+
+    # Store context for QA
+    context_storage["context"] = text
+
     return category, f"Confidence: {confidence}%"
 
-def classify_csv(file):
+def classify_csv(file_path):
     try:
-        df = pd.read_csv(file, encoding="utf-8")
-        text_column = df.columns[0]
-        df["Cleaned_Text"] = df[text_column].astype(str).apply(preprocess_text)
-        df["Encoded Prediction"] = df["Cleaned_Text"].apply(lambda x: news_classifier(x)[0]['label'])
+        df = pd.read_csv(file_path, encoding="utf-8")
+
+        # Automatically detect the column containing text
+        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["Decoded Prediction"] = df["Encoded Prediction"].map(label_mapping)
-        df["Confidence"] = df["Cleaned_Text"].apply(lambda x: round(news_classifier(x)[0]['score'] * 100, 2))
-        context_storage["bulk_context"] = " ".join(df["Cleaned_Text"].dropna().tolist())
+        df["Confidence"] = df[text_column].apply(lambda x: round(news_classifier(str(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())
         context_storage["num_articles"] = len(df)
+
         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)}"
 
-# Streamlit App
+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?",
+        "how are you": "🤖 I'm just a bot, but I'm here to help!",
+        "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!")
+    history.append([user_input, response])
+    return history, ""
+
+# Streamlit App Layout
 st.set_page_config(page_title="News Classifier", page_icon="📰")
-st.image("cover.png", caption="News Classifier 📢", use_column_width=True)
 
+# Load Cover Image
+cover_image = Image.open("cover.png")  # Ensure this image exists
+st.image(cover_image, caption="News Classifier 📢", use_column_width=True)
+
+# Section for Single Article Classification
 st.subheader("📰 Single Article Classification")
 text_input = st.text_area("Enter News Text", placeholder="Type or paste news content here...")
 if st.button("🔍 Classify"):
@@ -99,6 +102,7 @@ if st.button("🔍 Classify"):
     else:
         st.warning("Please enter some text to classify.")
 
+# Section for Bulk CSV Classification
 st.subheader("📂 Bulk Classification (CSV)")
 file_input = st.file_uploader("Upload CSV File", type="csv")
 if file_input:
@@ -113,3 +117,15 @@ if file_input:
         )
     else:
         st.error(f"Error processing file: {output_file}")
+
+# Section for Chatbot Interaction
+st.subheader("💬 AI Chat Assistant")
+history = []
+user_input = st.text_input("Ask about news classification or topics", placeholder="Type a message...")
+source_toggle = st.radio("Select Context Source", ["Single Article", "Bulk Classification"])
+if st.button("✉ Send"):
+    history, bot_response = chatbot_response(history, user_input, source_toggle)
+    st.write("*Chatbot Response:*")
+    for q, a in history:
+        st.write(f"*Q:* {q}")
+        st.write(f"*A:* {a}")