predict.py

#!/usr/bin/env python3
"""
Prediction Script for Fake News Detection
Classifies news articles as real or fake.

Usage:
    python predict.py "Your news article text here"
    python predict.py --file article.txt
    python predict.py --interactive

Author: Anik Tahabilder
"""

import argparse
import os
import sys
import re

# Add src to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from utils import load_model, print_banner


# Paths to saved models
MODEL_DIR = 'models'
VECTORIZER_PATH = os.path.join(MODEL_DIR, 'tfidf_vectorizer.pkl')
MODEL_PATH = os.path.join(MODEL_DIR, 'best_model.pkl')


def preprocess_text(text: str) -> str:
    """Clean and preprocess input text."""
    if not text:
        return ""

    # Convert to lowercase
    text = text.lower()

    # Remove URLs
    text = re.sub(r'http\S+|www\S+|https\S+', '', text)

    # Remove email addresses
    text = re.sub(r'\S+@\S+', '', text)

    # Remove special characters and digits
    text = re.sub(r'[^a-zA-Z\s]', '', text)

    # Remove extra whitespace
    text = re.sub(r'\s+', ' ', text).strip()

    return text


class FakeNewsPredictor:
    """
    Fake News Prediction Interface.

    Loads trained models and provides prediction functionality.
    """

    def __init__(self, model_path: str = MODEL_PATH,
                 vectorizer_path: str = VECTORIZER_PATH):
        """
        Initialize the predictor.

        Args:
            model_path: Path to the trained model
            vectorizer_path: Path to the TF-IDF vectorizer
        """
        self.model = None
        self.vectorizer = None

        self._load_models(model_path, vectorizer_path)

    def _load_models(self, model_path: str, vectorizer_path: str):
        """Load the model and vectorizer from disk."""
        if not os.path.exists(model_path):
            raise FileNotFoundError(
                f"Model not found at {model_path}. "
                "Please run 'python train.py' first."
            )

        if not os.path.exists(vectorizer_path):
            raise FileNotFoundError(
                f"Vectorizer not found at {vectorizer_path}. "
                "Please run 'python train.py' first."
            )

        self.model = load_model(model_path)
        self.vectorizer = load_model(vectorizer_path)

    def predict(self, text: str) -> dict:
        """
        Predict if a news article is fake or real.

        Args:
            text: News article text

        Returns:
            Dictionary with prediction results
        """
        # Preprocess
        clean_text = preprocess_text(text)

        if len(clean_text) < 20:
            return {
                'prediction': 'UNKNOWN',
                'confidence': 0.0,
                'message': 'Text too short for reliable prediction'
            }

        # Vectorize
        text_tfidf = self.vectorizer.transform([clean_text])

        # Get prediction
        if hasattr(self.model, 'predict'):
            # It's a FakeNewsClassifier wrapper
            prediction = self.model.predict(text_tfidf)[0]
            try:
                proba = self.model.predict_proba(text_tfidf)[0]
                confidence = max(proba) * 100
            except Exception:
                confidence = None
        else:
            # Direct sklearn model
            prediction = self.model.predict(text_tfidf)[0]
            if hasattr(self.model, 'predict_proba'):
                proba = self.model.predict_proba(text_tfidf)[0]
                confidence = max(proba) * 100
            else:
                confidence = None

        result = 'FAKE' if prediction == 1 else 'REAL'

        return {
            'prediction': result,
            'confidence': confidence,
            'label': int(prediction)
        }

    def predict_batch(self, texts: list) -> list:
        """
        Predict for multiple texts.

        Args:
            texts: List of news article texts

        Returns:
            List of prediction dictionaries
        """
        return [self.predict(text) for text in texts]


def interactive_mode(predictor: FakeNewsPredictor):
    """Run interactive prediction mode."""
    print_banner("INTERACTIVE MODE")
    print("\nEnter news articles to classify (type 'quit' to exit):")
    print("-" * 50)

    while True:
        print("\n")
        try:
            text = input("Enter article text:\n> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n\nExiting...")
            break

        if text.lower() in ['quit', 'exit', 'q']:
            print("Goodbye!")
            break

        if not text:
            print("Please enter some text.")
            continue

        result = predictor.predict(text)

        print(f"\n{'='*40}")
        print(f"Prediction: {result['prediction']}")
        if result['confidence']:
            print(f"Confidence: {result['confidence']:.1f}%")
        print(f"{'='*40}")


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description='Fake News Detection Predictor'
    )
    parser.add_argument(
        'text',
        nargs='?',
        help='News article text to classify'
    )
    parser.add_argument(
        '--file', '-f',
        type=str,
        help='Path to file containing news article'
    )
    parser.add_argument(
        '--interactive', '-i',
        action='store_true',
        help='Run in interactive mode'
    )
    parser.add_argument(
        '--model', '-m',
        type=str,
        default=MODEL_PATH,
        help='Path to model file'
    )

    args = parser.parse_args()

    # Initialize predictor
    try:
        predictor = FakeNewsPredictor(model_path=args.model)
    except FileNotFoundError as e:
        print(f"Error: {e}")
        print("\nPlease train the model first by running:")
        print("  python train.py")
        return 1

    # Interactive mode
    if args.interactive:
        interactive_mode(predictor)
        return 0

    # File mode
    if args.file:
        if not os.path.exists(args.file):
            print(f"Error: File not found: {args.file}")
            return 1

        with open(args.file, 'r', encoding='utf-8') as f:
            text = f.read()

        result = predictor.predict(text)

        print(f"\nFile: {args.file}")
        print(f"Prediction: {result['prediction']}")
        if result['confidence']:
            print(f"Confidence: {result['confidence']:.1f}%")
        return 0

    # Direct text mode
    if args.text:
        result = predictor.predict(args.text)

        print(f"\nPrediction: {result['prediction']}")
        if result['confidence']:
            print(f"Confidence: {result['confidence']:.1f}%")
        return 0

    # No input provided
    parser.print_help()
    return 1


if __name__ == '__main__':
    sys.exit(main())