Lightweight NLP pipeline to detect clickbait headlines using TF‑IDF features and a Logistic Regression classifier.
- Goal: classify news/article headlines as clickbait vs non‑clickbait
- Tech: Python, scikit‑learn, pandas, NLTK
- Current Version: v1.1
- Artifacts (v1.1):
models/clickbait_lr_tfidf_v1.1.pklmodels/clickbait_tfidf_vectorizer_v1.1.pkl
The v1.1 model achieved an Accuracy of 0.9809, Precision of 0.9846, Recall of 0.9771, and F1-Score of 0.9808. This represents significant improvements over v1.0. For a detailed breakdown of all metrics and version history, please refer to Model Details.
data/raw/— Original datasets.data/processed/— Cleaned datasets used for training/evaluation.src/— Source code:preprocessing.py— Text cleaning and normalization utilities.dataset.py— Dataset loading and dataset‑level preprocessing.features.py— TF‑IDF feature extraction.train_model.py— End‑to‑end training, evaluation and model saving.predict.py— Inference helpers (single, batch, and CSV‑based predictions).
models/— Saved model and vectorizer (*.pkl).main.py— CLI entry point to train and run the model.requirements.txt— Python dependencies.docs/— Project documentation and visualization figures.
- Create a virtual environment (Windows PowerShell):
python -m venv venv
.\venv\Scripts\Activate.ps1- Install dependencies:
pip install -r requirements.txt- Prepare data:
- Place your raw CSV in
data/raw/, e.g.data/raw/clickbait_data.csv. - Required columns:
headline— the text of the headlineclickbait— label (1 = clickbait, 0 = non‑clickbait)
This will train the Logistic Regression model and save the artifacts with versioned names:
python main.py train `--data-path data/raw/clickbait_data.csv `--model-path models/clickbait_lr_tfidf_v1.1.pkl `--vectorizer-path models/clickbait_tfidf_vectorizer_v1.1.pklIf you omit the arguments, main.py uses v1.0 defaults. For v1.1, specify the paths explicitly.
Use the trained model and vectorizer to classify one headline:
python main.py predict "You Won't Believe What Happened Next!"`--model-path models/clickbait_lr_tfidf_v1.1.pkl `--vectorizer-path models/clickbait_tfidf_vectorizer_v1.1.pklLow‑level text preprocessing utilities that operate on raw strings or token lists:
lowercase_text(text)— converts text to lowercase.remove_urls(text)— removeshttp:///https://style URLs.remove_punctuation(text)— strips punctuation characters.remove_numbers(text)— removes digits from the text.tokenize_text(text)— tokenizes text into word tokens using NLTK.remove_stopwords(tokens)— removes common English stopwords.lemmatize_tokens(tokens)— lemmatizes tokens using WordNet.preprocess_text(text)— runs all of the above in sequence and returns a cleaned string; this is the main function used by the rest of the pipeline.
Helpers for reading and preparing datasets:
load_dataset(filepath)- Reads a CSV from
filepathinto a pandas DataFrame. - Drops rows with missing values.
- Reads a CSV from
preprocess_dataset(df)- Takes a DataFrame with a
headlinecolumn. - Applies
preprocess_textto each headline. - Returns a copy of the DataFrame with an extra
clean_headlinecolumn.
- Takes a DataFrame with a
TF‑IDF feature extraction:
extract_tfidf_features(df, text_column="clean_headline")- Fits a
TfidfVectorizer(1‑, 2‑, and 3‑grams, up to 8000 features) on the specified text column. - v1.1 improvements: Extended n-gram range from (1,2) to (1,3) and increased max features from 5000 to 8000 for better context capture.
- Returns:
tfidf_df— a DataFrame of dense TF‑IDF features.vectorizer— the fittedTfidfVectorizerfor later use at inference time.
- Fits a
End‑to‑end training pipeline:
train_clickbait_model(data_path, test_size=0.2, random_state=42, max_iter=1000)- Loads and cleans the dataset (
load_dataset→preprocess_dataset). - Splits into train/test sets.
- Builds TF‑IDF features using
extract_tfidf_features. - Trains a
LogisticRegressionclassifier withC=10for improved generalization. - v1.1 improvements: Adjusted regularization parameter
Cfrom default (1.0) to 10 for better model performance. - Evaluates on the test set and returns:
model— trained classifiervectorizer— fitted TF‑IDF vectorizermetrics— accuracy, precision, recall, F1, confusion matrix, and classification report
- Loads and cleans the dataset (
save_model(model, vectorizer, model_path, vectorizer_path)- Saves both objects to disk using
joblib.
- Saves both objects to disk using
print_metrics(metrics)- Nicely prints the evaluation metrics to the console.
Inference utilities for using a trained model and vectorizer:
load_model(model_path="models/clickbait_lr_tfidf.pkl", vectorizer_path="models/clickbait_tfidf_vectorizer.pkl")- Loads a model and vectorizer from disk.
- For the v1.1 artifacts, call it like:
fromsrc.predictimportload_modelmodel, vectorizer=load_model(
model_path="models/clickbait_lr_tfidf_v1.1.pkl",
vectorizer_path="models/clickbait_tfidf_vectorizer_v1.1.pkl",
)predict_headline(headline, model, vectorizer, return_probability=False)- Preprocesses a single headline string and predicts 1 (clickbait) / 0 (non‑clickbait).
- If
return_probability=True, returns a dict with prediction, label, and probabilities.
predict_batch(headlines, model, vectorizer, return_probabilities=False)- Predicts for a list/Series of headlines.
- Optionally returns a DataFrame with predictions and probabilities.
predict_from_file(filepath, model, vectorizer, text_column="headline", return_probabilities=False)- Reads a CSV from
filepath. - Runs predictions on the
text_column. - Returns a DataFrame combining original data with predictions.
- Reads a CSV from
Example Python usage:
fromsrc.predictimportload_model, predict_headlinemodel, vectorizer=load_model(
model_path="models/clickbait_lr_tfidf_v1.1.pkl",
vectorizer_path="models/clickbait_tfidf_vectorizer_v1.1.pkl",
)
result=predict_headline(
"You Won't Believe What Happened Next!",
model,
vectorizer,
return_probability=True,
)
print(result)
