Update app.py

app.py

@@ -1,219 +1,115 @@
 import streamlit as st
 import pandas as pd
-import torch
+import string
 import re
+import nltk
 from nltk.corpus import stopwords
 from nltk.tokenize import word_tokenize
 from nltk.stem import WordNetLemmatizer
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
-import nltk
+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")
 
-# Set page configuration
-st.set_page_config(page_title="News Analysis App", layout="wide")
+# Load Models
+news_classifier = pipeline("text-classification", model="Oneli/News_Classification")
+qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
 
-# Download required NLTK resources
-@st.cache_resource
-def download_nltk_resources():
-    nltk.download('punkt')
-    nltk.download('stopwords')
-    nltk.download('wordnet')
+# Label Mapping
+label_mapping = {
+    "LABEL_0": "Business",
+    "LABEL_1": "Opinion",
+    "LABEL_2": "Political Gossip",
+    "LABEL_3": "Sports",
+    "LABEL_4": "World News"
+}
 
-download_nltk_resources()
+# Store classified article for QA
+context_storage = {"context": "", "bulk_context": "", "num_articles": 0}
 
-# Initialize preprocessor components
-stop_words = set(stopwords.words('english'))
-lemmatizer = WordNetLemmatizer()
+# Preprocessing functions
+def remove_punctuation(text):
+    return text.translate(str.maketrans('', '', string.punctuation))
 
-# Load the fine-tuned model for classification
-@st.cache_resource
-def load_classification_model():
-    model_name = "Oneli/News_Classification"  # Replace with your actual model path
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSequenceClassification.from_pretrained(model_name)
-    return model, tokenizer
+def remove_special_characters(text):
+    return re.sub(r'[^A-Za-z\s]', '', text)
 
-# Load Q&A pipeline
-@st.cache_resource
-def load_qa_pipeline():
-    qa_pipeline = pipeline("question-answering", model="distilbert-base-cased-distilled-squad")
-    return qa_pipeline
+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]
 
-# Text preprocessing function
 def preprocess_text(text):
-    if pd.isna(text):
-        return ""
-    
-    # Convert to lowercase
     text = text.lower()
-    
-    # Remove URLs
-    text = re.sub(r'http\S+|www\S+|https\S+', '', text)
-    
-    # Remove HTML tags
-    text = re.sub(r'<.*?>', '', text)
-    
-    # Remove special characters and numbers
-    text = re.sub(r'[^a-zA-Z\s]', '', text)
-    
-    # Tokenize
-    tokens = word_tokenize(text)
-    
-    # Remove stopwords and lemmatize
-    cleaned_tokens = [lemmatizer.lemmatize(token) for token in tokens if token not in stop_words]
-    
-    # Join tokens back into text
-    cleaned_text = ' '.join(cleaned_tokens)
-    
-    return cleaned_text
-
-# Function to classify news articles (bulk processing)
-def classify_news(df, model, tokenizer):
-    # Preprocess the text
-    df['cleaned_content'] = df['content'].apply(preprocess_text)
-    
-    # Prepare for classification
-    texts = df['cleaned_content'].tolist()
-    
-    # Get predictions
-    predictions = []
-    batch_size = 16
-    
-    for i in range(0, len(texts), batch_size):
-        batch_texts = texts[i:i+batch_size]
-        inputs = tokenizer(batch_texts, padding=True, truncation=True, max_length=512, return_tensors="pt")
-        
-        with torch.no_grad():
-            outputs = model(**inputs)
-            logits = outputs.logits
-            batch_predictions = torch.argmax(logits, dim=1).tolist()
-            predictions.extend(batch_predictions)
-    
-    # Map numeric predictions back to class labels
-    id2label = model.config.id2label
-    df['class'] = [id2label[pred] for pred in predictions]
-    
-    return df
+    text = remove_punctuation(text)
+    text = remove_special_characters(text)
+    text = remove_stopwords(text)
+    tokens = tokenize_text(text)
+    tokens = lemmatize_tokens(tokens)
+    return " ".join(tokens)
 
-# Function for single article classification
-def classify_single_article(text, model, tokenizer):
-    # Preprocess the text
+# Classification functions
+def classify_text(text):
     cleaned_text = preprocess_text(text)
-    
-    # Prepare for classification
-    inputs = tokenizer(cleaned_text, padding=True, truncation=True, max_length=512, return_tensors="pt")
-    
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
-        prediction = torch.argmax(logits, dim=1).item()
-    
-    # Map numeric prediction back to class label
-    id2label = model.config.id2label
-    category = id2label[prediction]
-    confidence = torch.nn.functional.softmax(logits, dim=1).max().item() * 100
-    
-    return category, round(confidence, 2)
+    result = news_classifier(cleaned_text)[0]
+    category = label_mapping.get(result['label'], "Unknown")
+    confidence = round(result['score'] * 100, 2)
+    context_storage["context"] = cleaned_text
+    return category, f"Confidence: {confidence}%"
 
-# Main app
-def main():
-    st.title("News Classifier 📢")
-    
-    # Sidebar for navigation
-    st.sidebar.title("Navigation")
-    app_mode = st.sidebar.radio("Choose the app mode", ["News Classification", "Question Answering"])
-    
-    # Section for Single Article Classification
-    if app_mode == "News Classification":
-        st.header("📰 Single Article Classification")
-        st.write("Enter a news article or upload a CSV file to classify the content.")
-        
-        # Text input for single article classification
-        text_input = st.text_area("Enter News Text", placeholder="Type or paste news content here...")
-        if st.button("🔍 Classify"):
-            if text_input:
-                # Load classification model
-                with st.spinner("Loading classification model..."):
-                    model, tokenizer = load_classification_model()
+def classify_csv(file):
+    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["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())
+        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)}"
 
-                # Classify the text
-                with st.spinner("Classifying the article..."):
-                    category, confidence = classify_single_article(text_input, model, tokenizer)
-                    st.write(f"*Predicted Category:* {category}")
-                    st.write(f"*Confidence Level:* {confidence}%")
-            else:
-                st.warning("Please enter some text to classify.")
-        
-        # File upload for bulk classification
-        st.subheader("📂 Bulk Classification (CSV)")
-        file_input = st.file_uploader("Upload CSV File", type="csv")
-        if file_input:
-            df = pd.read_csv(file_input)
-            
-            # Display sample of the data
-            st.subheader("Sample of uploaded data")
-            st.dataframe(df.head())
-            
-            # Check if the required column exists
-            if 'content' not in df.columns:
-                st.error("The CSV file must contain a 'content' column with the news articles text.")
-            else:
-                # Load model and tokenizer
-                with st.spinner("Loading classification model..."):
-                    model, tokenizer = load_classification_model()
-                
-                # Classify button
-                if st.button("Classify Articles"):
-                    with st.spinner("Classifying news articles..."):
-                        # Perform classification
-                        result_df = classify_news(df, model, tokenizer)
-                        
-                        # Display results
-                        st.subheader("Classification Results")
-                        st.dataframe(result_df[['content', 'class']])
-                        
-                        # Save to CSV
-                        csv = result_df.to_csv(index=False)
-                        st.download_button(
-                            label="Download output.csv",
-                            data=csv,
-                            file_name="output.csv",
-                            mime="text/csv"
-                        )
-                        
-                        # Show distribution of classes
-                        st.subheader("Class Distribution")
-                        class_counts = result_df['class'].value_counts()
-                        st.bar_chart(class_counts)
-    
-    # Section for Question Answering
-    elif app_mode == "Question Answering":
-        st.header("💬 AI Chat Assistant")
-        st.write("Ask questions about news content and get answers using a Q&A model.")
-        
-        # Text area for news content
-        news_content = st.text_area("Paste news article content here:", height=200)
-        
-        # Question input
-        question = st.text_input("Enter your question about the article:")
-        
-        if news_content and question:
-            # Load QA pipeline
-            with st.spinner("Loading Q&A model..."):
-                qa_pipeline = load_qa_pipeline()
-            
-            # Get answer
-            if st.button("Get Answer"):
-                with st.spinner("Finding answer..."):
-                    result = qa_pipeline(question=question, context=news_content)
-                    
-                    # Display results
-                    st.subheader("Answer")
-                    st.write(result["answer"])
-                    
-                    st.subheader("Confidence")
-                    st.progress(float(result["score"]))
-                    st.write(f"Confidence Score: {result['score']:.4f}")
+# Streamlit App
+st.set_page_config(page_title="News Classifier", page_icon="📰")
+st.image("cover.png", caption="News Classifier 📢", use_column_width=True)
 
-if __name__ == "__main__":
-    main()
+st.subheader("📰 Single Article Classification")
+text_input = st.text_area("Enter News Text", placeholder="Type or paste news content here...")
+if st.button("🔍 Classify"):
+    if text_input:
+        category, confidence = classify_text(text_input)
+        st.write(f"*Predicted Category:* {category}")
+        st.write(f"*Confidence Level:* {confidence}")
+    else:
+        st.warning("Please enter some text to classify.")
 
+st.subheader("📂 Bulk Classification (CSV)")
+file_input = st.file_uploader("Upload CSV File", type="csv")
+if file_input:
+    df, output_file = classify_csv(file_input)
+    if df is not None:
+        st.dataframe(df)
+        st.download_button(
+            label="Download Processed CSV",
+            data=open(output_file, 'rb').read(),
+            file_name=output_file,
+            mime="text/csv"
+        )
+    else:
+        st.error(f"Error processing file: {output_file}")