TA_Project / app.py
Ashendilantha's picture
Upload 2 files
79578f9 verified
import streamlit as st
import pandas as pd
import numpy as np
import re
import nltk
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
from nltk.stem import WordNetLemmatizer
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
import requests
from io import BytesIO
# Set page configuration
st.set_page_config(page_title="News Analysis App", layout="wide")
# Download required NLTK resources
@st.cache_resource
def download_nltk_resources():
nltk.download('punkt')
nltk.download('stopwords')
nltk.download('wordnet')
download_nltk_resources()
# Initialize preprocessor components
stop_words = set(stopwords.words('english'))
lemmatizer = WordNetLemmatizer()
# 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
# 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
# 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
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
# Main app
def main():
st.title("News Analysis Application")
# Sidebar for navigation
st.sidebar.title("Navigation")
app_mode = st.sidebar.radio("Choose the app mode", ["News Classification", "Question Answering"])
if app_mode == "News Classification":
st.header("News Article Classification")
st.write("Upload a CSV file containing news articles to classify them into categories.")
# File upload
uploaded_file = st.file_uploader("Choose a CSV file", type="csv")
if uploaded_file is not None:
# Load the data
df = pd.read_csv(uploaded_file)
# 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)
elif app_mode == "Question Answering":
st.header("News Article Q&A")
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}")
if __name__ == "__main__":
main()