Update app.py

app.py

@@ -1,3 +1,123 @@
+import gradio as gr
+import pandas as pd
+import torch
+from transformers import pipeline
+import nltk
+from nltk.corpus import stopwords
+from nltk.stem import WordNetLemmatizer
+import re
+import string
+import matplotlib.pyplot as plt
+
+# ====================== 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()
+    punct_to_remove = string.punctuation.replace("'","").replace('"',"").replace("$","").replace("%","").replace("?","")
+    text = re.sub(f"[{punct_to_remove}]", " ", text)
+    tokens = nltk.word_tokenize(text)
+    tokens = [lemmatizer.lemmatize(word) for word in tokens]
+    return ' '.join(tokens)
+
+classifier_model = "Ginidu2003/Distilbert-Base-News-classifier"
+
+# ====================== BEAUTIFUL COLORED BAR CHART ======================
+def create_colored_bar_chart(category_counts):
+    if category_counts is None or len(category_counts) == 0:
+        fig, ax = plt.subplots()
+        ax.text(0.5, 0.5, "No data available", ha='center', va='center')
+        return fig
+
+    categories = category_counts["Category"]
+    counts = category_counts["Count"]
+
+    # Nice modern color palette
+    colors = ['#3498DB', '#E67E22', '#9B59B6', '#2ECC71', '#E74C3C']
+
+    fig, ax = plt.subplots(figsize=(11, 6))
+    bars = ax.bar(categories, counts, color=colors, edgecolor='white', linewidth=0.8)
+
+    # Add value on top of bars
+    for bar in bars:
+        height = bar.get_height()
+        ax.text(bar.get_x() + bar.get_width()/2, height + 0.8,
+                str(int(height)), ha='center', va='bottom', fontsize=13, fontweight='bold')
+
+    ax.set_title("Category Distribution Across 5 Classes", fontsize=16, fontweight='bold', pad=20)
+    ax.set_xlabel("Category", fontsize=12)
+    ax.set_ylabel("Count", fontsize=12)
+    plt.xticks(rotation=15)
+    plt.tight_layout()
+    return fig
+
+# ====================== CLASSIFICATION FUNCTION ======================
+@torch.no_grad()
+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
+        
+        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():
+                predictions.append("Unknown")
+            else:
+                result = classifier(text)[0]
+                predictions.append(result['label'])
+        
+        df['class'] = predictions
+        df = df.drop(columns=['clean_content'], errors='ignore')
+        
+        output_file = "output.csv"
+        df.to_csv(output_file, index=False)
+        
+        category_counts = df['class'].value_counts().reset_index()
+        category_counts.columns = ["Category", "Count"]
+        
+        fig = create_colored_bar_chart(category_counts)
+        
+        return f"✅ Success! Classified {len(df)} rows", output_file, fig
+    except Exception as e:
+        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:
+        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
+        
+        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: {str(e)}"
+
 # ====================== BEAUTIFUL UI ======================
 with gr.Blocks(
     title="English News Classifier",