# apaiglow/ag-news-classification architecture

> Architecture diagram of the GitHub repository apaiglow/ag-news-classification, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-28). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/apaiglow/ag-news-classification
- Repository: https://github.com/apaiglow/ag-news-classification

## Overview

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: [`src/train.py`](https://github.com/apaiglow/ag-news-classification/blob/main/src/train.py) (Train), [`src/inference.py`](https://github.com/apaiglow/ag-news-classification/blob/main/src/inference.py) (Infer), [`src/evaluate.py`](https://github.com/apaiglow/ag-news-classification/blob/main/src/evaluate.py) (Evaluate), [`src/eda.py`](https://github.com/apaiglow/ag-news-classification/blob/main/src/eda.py) (EDA Script)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_runtime["Runtime"]
  node_main_py["Main<br/>entrypoint<br/>[main.py]"]
  node_streamlit_app_py["Streamlit<br/>ui<br/>[streamlit_app.py]"]
end

subgraph group_core["Core ML"]
  node_train_py["Train<br/>orchestrator<br/>[train.py]"]
  node_evaluate_py["Evaluate<br/>metrics<br/>[evaluate.py]"]
  node_inference_py["Infer<br/>runtime<br/>[inference.py]"]
  node_preprocessing_py["Preprocess<br/>text pipeline<br/>[preprocessing.py]"]
  node_vocabulary_py["Vocabulary<br/>token map<br/>[vocabulary.py]"]
  node_dataset_py["Dataset<br/>data layer<br/>[dataset.py]"]
  node_model_py["LSTM<br/>classifier<br/>[model.py]"]
end

subgraph group_artifacts["Artifacts"]
  node_best_model_pth[("Best Model<br/>checkpoint<br/>[best_model.pth]")]
  node_best_model_weights_pth[("Model Weights<br/>checkpoint")]
  node_vocab_pkl[("Vocab File<br/>artifact<br/>[vocab.pkl]")]
  node_classification_report_txt["Report<br/>evaluation output"]
  node_eda_summary_txt["EDA Summary<br/>analysis output<br/>[eda_summary.txt]"]
  node_vocabulary_coverage_csv["Coverage<br/>analysis output"]
  node_words_frequency_table_csv["Frequency<br/>analysis output"]
end

subgraph group_analysis["Analysis"]
  node_eda_py["EDA Script<br/>analysis<br/>[eda.py]"]
  node_eda_notebook["EDA Notebook<br/>analysis notebook<br/>[EDA.ipynb]"]
end

node_main_py -->|"start training"| node_train_py
node_main_py -->|"run eval"| node_evaluate_py
node_train_py -->|"clean text"| node_preprocessing_py
node_train_py -->|"build vocab"| node_vocabulary_py
node_train_py -->|"batch data"| node_dataset_py
node_train_py -->|"fit model"| node_model_py
node_train_py -->|"save best"| node_best_model_pth
node_train_py -->|"save weights"| node_best_model_weights_pth
node_train_py -->|"persist vocab"| node_vocab_pkl
node_evaluate_py -->|"load model"| node_best_model_pth
node_evaluate_py -->|"write report"| node_classification_report_txt
node_evaluate_py -->|"write summary"| node_eda_summary_txt
node_inference_py -->|"normalize input"| node_preprocessing_py
node_inference_py -->|"lookup tokens"| node_vocab_pkl
node_inference_py -->|"load weights"| node_best_model_weights_pth
node_inference_py -->|"predict"| node_model_py
node_streamlit_app_py -->|"serve prediction"| node_inference_py
node_eda_py -->|"analyze text"| node_preprocessing_py
node_eda_py -->|"export coverage"| node_vocabulary_coverage_csv
node_eda_py -->|"export frequencies"| node_words_frequency_table_csv
node_eda_notebook -.->|"prototype"| node_eda_py

click node_main_py "https://github.com/apaiglow/ag-news-classification/blob/main/main.py"
click node_streamlit_app_py "https://github.com/apaiglow/ag-news-classification/blob/main/streamlit_app.py"
click node_train_py "https://github.com/apaiglow/ag-news-classification/blob/main/src/train.py"
click node_evaluate_py "https://github.com/apaiglow/ag-news-classification/blob/main/src/evaluate.py"
click node_inference_py "https://github.com/apaiglow/ag-news-classification/blob/main/src/inference.py"
click node_preprocessing_py "https://github.com/apaiglow/ag-news-classification/blob/main/src/preprocessing.py"
click node_vocabulary_py "https://github.com/apaiglow/ag-news-classification/blob/main/src/vocabulary.py"
click node_dataset_py "https://github.com/apaiglow/ag-news-classification/blob/main/src/dataset.py"
click node_model_py "https://github.com/apaiglow/ag-news-classification/blob/main/src/model.py"
click node_eda_py "https://github.com/apaiglow/ag-news-classification/blob/main/src/eda.py"
click node_eda_notebook "https://github.com/apaiglow/ag-news-classification/blob/main/notebooks/EDA.ipynb"
click node_best_model_pth "https://github.com/apaiglow/ag-news-classification/blob/main/models/best_model.pth"
click node_best_model_weights_pth "https://github.com/apaiglow/ag-news-classification/blob/main/models/best_model_weights.pth"
click node_vocab_pkl "https://github.com/apaiglow/ag-news-classification/blob/main/models/vocab.pkl"
click node_classification_report_txt "https://github.com/apaiglow/ag-news-classification/blob/main/results/classification_report.txt"
click node_eda_summary_txt "https://github.com/apaiglow/ag-news-classification/blob/main/results/eda_summary.txt"
click node_vocabulary_coverage_csv "https://github.com/apaiglow/ag-news-classification/blob/main/results/EDAGraphs/vocabulary_coverage.csv"
click node_words_frequency_table_csv "https://github.com/apaiglow/ag-news-classification/blob/main/results/EDAGraphs/words_frequency_table.csv"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_main_py,node_streamlit_app_py toneBlue
class node_train_py,node_evaluate_py,node_inference_py,node_preprocessing_py,node_vocabulary_py,node_dataset_py,node_model_py toneAmber
class node_best_model_pth,node_best_model_weights_pth,node_vocab_pkl,node_classification_report_txt,node_eda_summary_txt,node_vocabulary_coverage_csv,node_words_frequency_table_csv toneMint
class node_eda_py,node_eda_notebook toneRose
```

## Components

### Runtime

Entry points for training, evaluation, and interactive prediction.

- [Main](https://github.com/apaiglow/ag-news-classification/blob/main/main.py): entrypoint · CLI workflow for training and evaluation. · `main.py`
- [Streamlit](https://github.com/apaiglow/ag-news-classification/blob/main/streamlit_app.py): ui · Interactive prediction app. · `streamlit_app.py`

### Core ML

Reusable preprocessing, data, vocabulary, and model components.

- [Train](https://github.com/apaiglow/ag-news-classification/blob/main/src/train.py): orchestrator · Coordinates training, validation, and checkpointing. · `src/train.py`
- [Evaluate](https://github.com/apaiglow/ag-news-classification/blob/main/src/evaluate.py): metrics · Generates reports and diagnostics from trained models. · `src/evaluate.py`
- [Infer](https://github.com/apaiglow/ag-news-classification/blob/main/src/inference.py): runtime · Loads saved artifacts and predicts article classes. · `src/inference.py`
- [Preprocess](https://github.com/apaiglow/ag-news-classification/blob/main/src/preprocessing.py): text pipeline · Normalizes raw title/description text into model-ready sequences. · `src/preprocessing.py`
- [Vocabulary](https://github.com/apaiglow/ag-news-classification/blob/main/src/vocabulary.py): token map · Builds and applies token-to-index mappings shared across training and inference. · `src/vocabulary.py`
- [Dataset](https://github.com/apaiglow/ag-news-classification/blob/main/src/dataset.py): data layer · Turns processed text and labels into batched tensors. · `src/dataset.py`
- [LSTM](https://github.com/apaiglow/ag-news-classification/blob/main/src/model.py): classifier · Embedding-LSTM-fully-connected news classifier. · `src/model.py`

### Artifacts

Persisted model and vocabulary files plus generated evaluation outputs.

- [Best Model](https://github.com/apaiglow/ag-news-classification/blob/main/models/best_model.pth): checkpoint · Persisted trained model artifact. · `models/best_model.pth`
- [Model Weights](https://github.com/apaiglow/ag-news-classification/blob/main/models/best_model_weights.pth): checkpoint · Saved model weights for inference and reuse. · `models/best_model_weights.pth`
- [Vocab File](https://github.com/apaiglow/ag-news-classification/blob/main/models/vocab.pkl): artifact · Serialized vocabulary shared by training and inference. · `models/vocab.pkl`
- [Report](https://github.com/apaiglow/ag-news-classification/blob/main/results/classification_report.txt): evaluation output · Classification metrics summary. · `results/classification_report.txt`
- [EDA Summary](https://github.com/apaiglow/ag-news-classification/blob/main/results/eda_summary.txt): analysis output · Written summary of exploratory findings. · `results/eda_summary.txt`
- [Coverage](https://github.com/apaiglow/ag-news-classification/blob/main/results/EDAGraphs/vocabulary_coverage.csv): analysis output · Vocabulary coverage results from EDA. · `results/EDAGraphs/vocabulary_coverage.csv`
- [Frequency](https://github.com/apaiglow/ag-news-classification/blob/main/results/EDAGraphs/words_frequency_table.csv): analysis output · Token frequency table from EDA. · `results/EDAGraphs/words_frequency_table.csv`

### Analysis

Exploratory data analysis code and notebook.

- [EDA Script](https://github.com/apaiglow/ag-news-classification/blob/main/src/eda.py): analysis · Reproducible exploratory analysis workflow. · `src/eda.py`
- [EDA Notebook](https://github.com/apaiglow/ag-news-classification/blob/main/notebooks/EDA.ipynb): analysis notebook · Interactive exploratory analysis and experimentation. · `notebooks/EDA.ipynb`

## Connections

- 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.