How apaiglow/ag-news-classification works
1. What this repository is
- This is an offline training + inference NLP project for 4-class news article classification using a PyTorch LSTM model.
- It has two user-facing runtimes:
- main.py for training/evaluation workflow
- streamlit_app.py for interactive prediction
2. Primary data source
- The model is trained on the AG News CSV dataset described in the README.
- Input features are the title and description of an article.
- The target is a class index mapped to one of: World, Sports, Business, Sci/Tech.
3. Preprocessing boundary
- src/preprocessing.py is the first major processing stage.
- It likely normalizes raw text and prepares token sequences for downstream use.
- This is the boundary between raw tabular text and model-ready sequence data.
4. Vocabulary as a shared contract
- src/vocabulary.py and the saved artifact models/vocab.pkl form the contract between training and inference.
- The vocabulary is built during training, serialized, and then reused later so the same token-to-index mapping is available in the app and evaluation paths.
- This is important for architecture diagrams because it is a shared dependency across both offline and online flows.
5. Dataset and batching layer
- src/dataset.py provides the custom PyTorch dataset/data loading layer.
- Its role is to turn preprocessed text + labels into tensors suitable for batching.
- This is the bridge between preprocessing/vocabulary and the neural network.
6. Model subsystem
- src/model.py contains the actual classifier: Embedding → LSTM → Fully Connected output.
- The output layer produces logits for the 4 classes.
- Training uses CrossEntropyLoss and Adam, as described in the README.
- Saved model artifacts live in models/:
- best_model.pth
- best_model_weights.pth
7. Training pipeline
- src/train.py is the training orchestration layer.
- It likely handles splitting, dataloader creation, training loops, validation, and checkpointing.
- The main architectural behavior is that it automatically saves the best-performing model, so training and model selection are coupled.
8. Evaluation and reporting
- src/evaluate.py consumes the trained model and produces metrics and diagnostics.
- The generated outputs in results/ show the evaluation boundary:
- classification_report.txt
- EDA summary and graph outputs
- This layer is separate from training, which helps keep metric generation reproducible.
9. Inference path
- src/inference.py is the runtime prediction layer.
- It depends on the saved model weights and saved vocabulary to transform new text into a class prediction.
- In the architecture diagram, this should be shown as:
- user text → preprocessing/vocab lookup → trained model → predicted news category
10. Streamlit application
- streamlit_app.py is the presentation layer for interactive inference.
- It does not train the model; it loads the trained artifacts and exposes predictions through a UI.
- This is the user-facing boundary of the system.
11. EDA subsystem
- src/eda.py and notebooks/EDA.ipynb serve the exploratory analysis track.
- The notebook is for interactive experimentation, while src/eda.py provides a reproducible script version.
- Outputs are stored under results/EDAGraphs/, making EDA a distinct offline analysis pipeline rather than part of production inference.
12. Project artifacts and outputs
- The repo separates code from generated artifacts clearly:
- models/ = persisted model and vocabulary
- results/ = evaluation reports and EDA outputs
- notebooks/ = exploratory work
- This separation is useful in an architecture diagram because it shows which components are source, which are runtime dependencies, and which are generated outputs.
13. Overall system flow
- The end-to-end flow is:
- dataset ingestion → preprocessing → vocabulary creation → dataset/dataloader → LSTM training → best model saving → evaluation/report generation → inference in Streamlit
- The key architectural boundary is that training builds artifacts, while inference and Streamlit consume those artifacts.
14. Technology stack that materially affects architecture
- PyTorch drives the model, training, and tensor pipeline.
- Streamlit provides the interactive frontend.
- scikit-learn is likely used for reporting/metrics generation.
- The project is structured as a modular Python application rather than a monolithic notebook-only workflow.
Main entry points: train.py Train · inference.py Infer · evaluate.py Evaluate · eda.py EDA Script
Components 18
Runtime
Entry points for training, evaluation, and interactive prediction.
Core ML
Reusable preprocessing, data, vocabulary, and model components.
- Train · Coordinates training, validation, and checkpointing.
src/train.py - Evaluate · Generates reports and diagnostics from trained models.
src/evaluate.py - Infer · Loads saved artifacts and predicts article classes.
src/inference.py - Preprocess · Normalizes raw title/description text into model-ready sequences.
src/preprocessing.py - Vocabulary · Builds and applies token-to-index mappings shared across training and inference.
src/vocabulary.py - Dataset · Turns processed text and labels into batched tensors.
src/dataset.py - LSTM · Embedding-LSTM-fully-connected news classifier.
src/model.py
Artifacts
Persisted model and vocabulary files plus generated evaluation outputs.
- Best Model · Persisted trained model artifact.
models/best_model.pth - Model Weights · Saved model weights for inference and reuse.
models/best_model_weights.pth - Vocab File · Serialized vocabulary shared by training and inference.
models/vocab.pkl - Report · Classification metrics summary.
results/classification_report.txt - EDA Summary · Written summary of exploratory findings.
results/eda_summary.txt - Coverage · Vocabulary coverage results from EDA.
results/EDAGraphs/vocabulary_coverage.csv - Frequency · Token frequency table from EDA.
results/EDAGraphs/words_frequency_table.csv
Analysis
Exploratory data analysis code and notebook.
- EDA Script · Reproducible exploratory analysis workflow.
src/eda.py - EDA Notebook · Interactive exploratory analysis and experimentation.
notebooks/EDA.ipynb
Connections 21
- Main → Train: start training · Main entrypoint launches the training workflow.
- Main → Evaluate: run eval · Main entrypoint can trigger evaluation/reporting.
- Train → Preprocess: clean text · Training consumes normalized article text.
- Train → Vocabulary: build vocab · Training builds the shared token mapping.
- Train → Dataset: batch data · Training creates datasets and dataloaders.
- Train → LSTM: fit model · Training optimizes the LSTM classifier.
- Train → Best Model: save best · Best-performing checkpoint is persisted.
- Train → Model Weights: save weights · Weights artifact is written for reuse.
- Train → Vocab File: persist vocab · Training serializes the vocabulary contract.
- Evaluate → Best Model: load model · Evaluation reads the trained checkpoint.
- Evaluate → Report: write report · Evaluation produces classification metrics output.
- Evaluate → EDA Summary: write summary · Evaluation/analysis outputs a summary artifact.
- Infer → Preprocess: normalize input · Inference preprocesses user text before prediction.
- Infer → Vocab File: lookup tokens · Inference reuses the saved vocabulary mapping.
- Infer → Model Weights: load weights · Inference loads trained parameters.
- Infer → LSTM: predict · Inference runs the LSTM classifier to produce a class.
- Streamlit → Infer: serve prediction · UI delegates prediction to the inference layer.
- EDA Script → Preprocess: analyze text · EDA script reuses text preprocessing where needed.
- EDA Script → Coverage: export coverage · EDA writes vocabulary coverage analysis.
- EDA Script → Frequency: export frequencies · EDA writes frequency analysis output.
- EDA Notebook → EDA Script: prototype · Notebook and script represent the same analysis track.