Update app.py

app.py

@@ -9,20 +9,19 @@ import re
 import string
 
 # ====================== NLTK SETUP ======================
-
 nltk.download('wordnet', quiet=True)
 nltk.download('punkt', quiet=True)
 nltk.download('punkt_tab', quiet=True)
 
-
 lemmatizer = WordNetLemmatizer()
 
 def preprocess_text(text):
     if not isinstance(text, str):
         return ""
     text = text.lower()
+    # Remove punctuation except: ' " $ % ?
     punct_to_remove = string.punctuation.replace("'","").replace('"',"").replace("$","").replace("%","").replace("?","")
-    text = re.sub(f"[{punct_to_remove}]", " ", text)# Remove punctuation except: ' " $ % ?
+    text = re.sub(f"[{punct_to_remove}]", " ", text)  # remove the specified ones
     tokens = nltk.word_tokenize(text)
     tokens = [word for word in tokens]
     tokens = [lemmatizer.lemmatize(word) for word in tokens]
@@ -30,23 +29,24 @@ def preprocess_text(text):
 
 # ====================== MODELS ======================
 classifier_model = "Ginidu2003/Distilbert-Base-News-classifier"
-qa_model_name = "deepset/roberta-base-squad2"
-
-# Load QA pipeline using a supported method
-qa_pipeline = pipeline("document-question-answering", model=qa_model_name, device=-1)
+qa_model = "deepset/roberta-base-squad2"   # Best for news Q&A
 
 # ====================== CLASSIFICATION FUNCTION ======================
+import matplotlib.pyplot as plt
+import io
+
 @torch.no_grad()
-def classify_csv(file):
+def classify_csv_with_chart(file):
     try:
         df = pd.read_csv(file)
+
         if 'content' not in df.columns:
-            return "Error: CSV must have a column named 'content'", None
-        
+            return "Error: CSV must have a column named 'content'", None, None
+
         df['clean_content'] = df['content'].apply(preprocess_text)
-        
+
         classifier = pipeline("text-classification", model=classifier_model, device=-1)
-        
+
         predictions = []
         for text in df['clean_content']:
             if not text.strip():
@@ -54,28 +54,59 @@ def classify_csv(file):
             else:
                 result = classifier(text)[0]
                 predictions.append(result['label'])
-        
+
         df['class'] = predictions
         df = df.drop(columns=['clean_content'], errors='ignore')
-        
+
+        # Save output
         output_file = "output.csv"
         df.to_csv(output_file, index=False)
-        
-        return f"✅ Success! Classified {len(df)} rows", output_file
+
+        # ====== Category Distribution Chart ======
+        fig, ax = plt.subplots()
+        df['class'].value_counts().plot(kind='bar', ax=ax, color="skyblue")
+        ax.set_title("Category Distribution")
+        ax.set_ylabel("Count")
+
+        buf = io.BytesIO()
+        plt.savefig(buf, format="png")
+        buf.seek(0)
+
+        return f"✅ Success! Classified {len(df)} rows", output_file, buf
+
     except Exception as e:
-        return f"❌ Error: {str(e)}", None
+        return f"❌ Error: {str(e)}", None, None
 
 # ====================== Q&A FUNCTION ======================
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+
+qa_tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")
+qa_model = AutoModelForQuestionAnswering.from_pretrained("deepset/roberta-base-squad2")
+
 def answer_question(news_content, question):
     if not news_content.strip() or not question.strip():
         return "Please enter both news content and a question."
+    
     try:
-        result = qa_pipeline(question=question, context=news_content)
-        answer = result[0]['answer'] if isinstance(result, list) else result['answer']
-        score = result[0]['score'] if isinstance(result, list) else result['score']
-        return f"**Answer:** {answer}\n\n**Confidence:** {score:.2%}"
+        inputs = qa_tokenizer(question, news_content, return_tensors="pt", truncation=True, max_length=512)
+        
+        with torch.no_grad():
+            outputs = qa_model(**inputs)
+        
+        start_idx = torch.argmax(outputs.start_logits)
+        end_idx = torch.argmax(outputs.end_logits) + 1
+        
+        # Clean answer - remove question repetition and special tokens
+        answer = qa_tokenizer.decode(inputs.input_ids[0][start_idx:end_idx], 
+                                    skip_special_tokens=True, 
+                                    clean_up_tokenization_spaces=True)
+        
+        confidence = torch.max(torch.softmax(outputs.start_logits, dim=1)).item()
+        
+        return f"**Answer:** {answer.strip()}\n\n**Confidence:** {confidence:.2%}"
+    
     except Exception as e:
-        return f"Error processing question: {str(e)}"
+        return f"Error: {str(e)}"
 
 # ====================== GRADIO INTERFACE ======================
 with gr.Blocks(title="Daily Mirror News Classifier") as demo:
@@ -83,25 +114,27 @@ with gr.Blocks(title="Daily Mirror News Classifier") as demo:
     gr.Markdown("### Section 02 - Text Analytics Assignment")
 
     with gr.Tabs():
+        # Tab 1: Classification
         with gr.Tab("📊 Text Classification"):
-            gr.Markdown("Upload CSV with `content` column")
             file_input = gr.File(label="Upload CSV", file_types=[".csv"])
             classify_btn = gr.Button("🚀 Classify News", variant="primary")
             output_text = gr.Textbox(label="Status")
             output_file = gr.File(label="Download output.csv")
-            
+            dist_plot = gr.Image(label="Category Distribution")
+
             classify_btn.click(
-                fn=classify_csv,
+                fn=classify_csv_with_chart,
                 inputs=file_input,
-                outputs=[output_text, output_file]
+                outputs=[output_text, output_file, dist_plot]
             )
 
+        # Tab 2: Q&A Pipeline
         with gr.Tab("❓ Question Answering"):
             gr.Markdown("Ask any question about a news article")
-            news_input = gr.Textbox(lines=12, label="Paste News Content", placeholder="Paste the full news article here...")
-            question_input = gr.Textbox(label="Your Question", placeholder="e.g. What is the main topic?")
+            news_input = gr.Textbox(lines=10, label="Paste News Content", placeholder="Paste the full news article here...")
+            question_input = gr.Textbox(label="Your Question", placeholder="e.g., What is the main issue discussed?")
             qa_btn = gr.Button("🔍 Get Answer", variant="primary")
-            qa_output = gr.Textbox(label="Answer", lines=5)
+            qa_output = gr.Textbox(label="Answer", lines=4)
             
             qa_btn.click(
                 fn=answer_question,
@@ -111,4 +144,4 @@ with gr.Blocks(title="Daily Mirror News Classifier") as demo:
 
     gr.Markdown("Built for Text Analytics Assignment - Section 02")
 
-demo.launch()
+demo.launch()