Update app.py

app.py

@@ -9,19 +9,20 @@ 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 the specified ones
+    text = re.sub(f"[{punct_to_remove}]", " ", text)# Remove punctuation except: ' " $ % ?
     tokens = nltk.word_tokenize(text)
     tokens = [word for word in tokens]
     tokens = [lemmatizer.lemmatize(word) for word in tokens]
@@ -29,24 +30,23 @@ def preprocess_text(text):
 
 # ====================== MODELS ======================
 classifier_model = "Ginidu2003/Distilbert-Base-News-classifier"
-qa_model = "deepset/roberta-base-squad2"   # Best for news Q&A
+qa_model_name = "deepset/roberta-base-squad2"
 
-# ====================== CLASSIFICATION FUNCTION ======================
-import matplotlib.pyplot as plt
-import io
+# Load QA pipeline using a supported method
+qa_pipeline = pipeline("document-question-answering", model=qa_model_name, device=-1)
 
+# ====================== CLASSIFICATION FUNCTION ======================
 @torch.no_grad()
-def classify_csv_with_chart(file):
+def classify_csv(file):
     try:
         df = pd.read_csv(file)
-
         if 'content' not in df.columns:
-            return "Error: CSV must have a column named 'content'", None, None
-
+            return "Error: CSV must have a column named 'content'", 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,59 +54,28 @@ def classify_csv_with_chart(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)
-
-        # ====== 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
-
+        
+        return f"✅ Success! Classified {len(df)} rows", output_file
     except Exception as e:
-        return f"❌ Error: {str(e)}", None, None
+        return f"❌ Error: {str(e)}", 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:
-        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%}"
-    
+        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%}"
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"Error processing question: {str(e)}"
 
 # ====================== GRADIO INTERFACE ======================
 with gr.Blocks(title="Daily Mirror News Classifier") as demo:
@@ -114,27 +83,25 @@ 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_with_chart,
+                fn=classify_csv,
                 inputs=file_input,
-                outputs=[output_text, output_file, dist_plot]
+                outputs=[output_text, output_file]
             )
 
-        # Tab 2: Q&A Pipeline
         with gr.Tab("❓ Question Answering"):
             gr.Markdown("Ask any question about a news article")
-            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?")
+            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?")
             qa_btn = gr.Button("🔍 Get Answer", variant="primary")
-            qa_output = gr.Textbox(label="Answer", lines=4)
+            qa_output = gr.Textbox(label="Answer", lines=5)
             
             qa_btn.click(
                 fn=answer_question,
@@ -144,4 +111,4 @@ with gr.Blocks(title="Daily Mirror News Classifier") as demo:
 
     gr.Markdown("Built for Text Analytics Assignment - Section 02")
 
-demo.launch()
+demo.launch()