TruthCheck / src /app.py
adnaan05's picture
Update src/app.py
81dd6cb verified
raw
history blame
7.65 kB
import streamlit as st
import torch
import pandas as pd
import numpy as np
from pathlib import Path
import sys
import plotly.express as px
import plotly.graph_objects as go
from transformers import BertTokenizer
import nltk
# Download required NLTK data
try:
nltk.data.find('tokenizers/punkt')
except LookupError:
nltk.download('punkt')
try:
nltk.data.find('corpora/stopwords')
except LookupError:
nltk.download('stopwords')
try:
nltk.data.find('tokenizers/punkt_tab')
except LookupError:
nltk.download('punkt_tab')
try:
nltk.data.find('corpora/wordnet')
except LookupError:
nltk.download('wordnet')
# Add project root to Python path
project_root = Path(__file__).parent.parent
sys.path.append(str(project_root))
from src.models.hybrid_model import HybridFakeNewsDetector
from src.config.config import *
from src.data.preprocessor import TextPreprocessor
# Page config is set in main app.py
@st.cache_resource
def load_model_and_tokenizer():
"""Load the model and tokenizer (cached)."""
# Initialize model
model = HybridFakeNewsDetector(
bert_model_name=BERT_MODEL_NAME,
lstm_hidden_size=LSTM_HIDDEN_SIZE,
lstm_num_layers=LSTM_NUM_LAYERS,
dropout_rate=DROPOUT_RATE
)
# Load trained weights
state_dict = torch.load(SAVED_MODELS_DIR / "final_model.pt", map_location=torch.device('cpu'))
# Filter out unexpected keys
model_state_dict = model.state_dict()
filtered_state_dict = {k: v for k, v in state_dict.items() if k in model_state_dict}
# Load the filtered state dict
model.load_state_dict(filtered_state_dict, strict=False)
model.eval()
# Initialize tokenizer
tokenizer = BertTokenizer.from_pretrained(BERT_MODEL_NAME)
return model, tokenizer
@st.cache_resource
def get_preprocessor():
"""Get the text preprocessor (cached)."""
return TextPreprocessor()
def predict_news(text):
"""Predict if the given news is fake or real."""
# Get model, tokenizer, and preprocessor from cache
model, tokenizer = load_model_and_tokenizer()
preprocessor = get_preprocessor()
# Preprocess text
processed_text = preprocessor.preprocess_text(text)
# Tokenize
encoding = tokenizer.encode_plus(
processed_text,
add_special_tokens=True,
max_length=MAX_SEQUENCE_LENGTH,
padding='max_length',
truncation=True,
return_attention_mask=True,
return_tensors='pt'
)
# Get prediction
with torch.no_grad():
outputs = model(
encoding['input_ids'],
encoding['attention_mask']
)
probabilities = torch.softmax(outputs['logits'], dim=1)
prediction = torch.argmax(outputs['logits'], dim=1)
attention_weights = outputs['attention_weights']
# Convert attention weights to numpy and get the first sequence
attention_weights_np = attention_weights[0].cpu().numpy()
return {
'prediction': prediction.item(),
'label': 'FAKE' if prediction.item() == 1 else 'REAL',
'confidence': torch.max(probabilities, dim=1)[0].item(),
'probabilities': {
'REAL': probabilities[0][0].item(),
'FAKE': probabilities[0][1].item()
},
'attention_weights': attention_weights_np
}
def plot_confidence(probabilities):
"""Plot prediction confidence."""
fig = go.Figure(data=[
go.Bar(
x=list(probabilities.keys()),
y=list(probabilities.values()),
text=[f'{p:.2%}' for p in probabilities.values()],
textposition='auto',
)
])
fig.update_layout(
title='Prediction Confidence',
xaxis_title='Class',
yaxis_title='Probability',
yaxis_range=[0, 1]
)
return fig
def plot_attention(text, attention_weights):
"""Plot attention weights."""
tokens = text.split()
attention_weights = attention_weights[:len(tokens)] # Truncate to match tokens
# Ensure attention weights are in the correct format
if isinstance(attention_weights, (list, np.ndarray)):
attention_weights = np.array(attention_weights).flatten()
# Format weights for display
formatted_weights = [f'{float(w):.2f}' for w in attention_weights]
fig = go.Figure(data=[
go.Bar(
x=tokens,
y=attention_weights,
text=formatted_weights,
textposition='auto',
)
])
fig.update_layout(
title='Attention Weights',
xaxis_title='Tokens',
yaxis_title='Attention Weight',
xaxis_tickangle=45
)
return fig
def main():
st.title("📰 Fake News Detection System")
st.write("""
This application uses a hybrid deep learning model (BERT + BiLSTM + Attention)
to detect fake news articles. Enter a news article below to analyze it.
""")
# Sidebar
st.sidebar.title("About")
st.sidebar.info("""
The model combines:
- BERT for contextual embeddings
- BiLSTM for sequence modeling
- Attention mechanism for interpretability
""")
# Main content
st.header("News Analysis")
# Text input
news_text = st.text_area(
"Enter the news article to analyze:",
height=200,
placeholder="Paste your news article here..."
)
if st.button("Analyze"):
if news_text:
with st.spinner("Analyzing the news article..."):
# Get prediction
result = predict_news(news_text)
# Display result
col1, col2 = st.columns(2)
with col1:
st.subheader("Prediction")
if result['label'] == 'FAKE':
st.error(f"🔴 This news is likely FAKE (Confidence: {result['confidence']:.2%})")
else:
st.success(f"🟢 This news is likely REAL (Confidence: {result['confidence']:.2%})")
with col2:
st.subheader("Confidence Scores")
st.plotly_chart(plot_confidence(result['probabilities']), use_container_width=True)
# Show attention visualization
st.subheader("Attention Analysis")
st.write("""
The attention weights show which parts of the text the model focused on
while making its prediction. Higher weights indicate more important tokens.
""")
st.plotly_chart(plot_attention(news_text, result['attention_weights']), use_container_width=True)
# Show model explanation
st.subheader("Model Explanation")
if result['label'] == 'FAKE':
st.write("""
The model identified this as fake news based on:
- Linguistic patterns typical of fake news
- Inconsistencies in the content
- Attention weights on suspicious phrases
""")
else:
st.write("""
The model identified this as real news based on:
- Credible language patterns
- Consistent information
- Attention weights on factual statements
""")
else:
st.warning("Please enter a news article to analyze.")
if __name__ == "__main__":
main()