Statistical Domain Language Model

📊 Project Overview

This repository contains an end-to-end machine learning pipeline designed to pretrain a domain-specific generative language model. The core of the system is a 124M-parameter GPT-2 style Transformer, built and trained from scratch using PyTorch. By ingesting and processing complex statistical textbooks, the model is engineered to understand and generate domain-specific text related to probability, variance, and statistical analysis.

🛠️ System Architecture & Hardware

• Model Architecture: GPT-2 (Decoder-only Transformer)
• Parameter Count: ~124 Million Parameters
• Deep Learning Framework: PyTorch & Hugging Face Transformers
• Hardware Acceleration: Apple Silicon (MacBook M4 Pro) utilizing Metal Performance Shaders (MPS) for GPU-accelerated training.
• Corpus Size: ~943,000 Tokens

⚙️ The Data Engineering Pipeline

Transforming raw educational textbooks into a clean, training-ready corpus required a multi-stage data processing pipeline. The raw data consisted of open-access PDF textbooks (e.g., OpenIntro Statistics, Think Stats), which presented significant formatting challenges.

1. PDF Extraction (extract_text.py): Automated parsing scripts utilizing pdfplumber to extract raw text, handling complex layouts across hundreds of pages.
2. Aggregation (combine_texts.py): Unified the raw outputs into a single, cohesive text document.
3. Text Sanitization (clean_text.py): Engineered Regex cleaning functions to strip out formatting artifacts, page numbers, repeating headers/footers, and normalize whitespace.
4. Algorithmic Deduplication (deduplicate.py): Implemented paragraph-level deduplication using MD5 hashing to identify and remove redundant text segments, maximizing the variance of the training dataset.
5. Tokenization (tokenize_corpus.py): Integrated a GPT-2 specific tokenization process, converting the sanitized corpus into optimized integer sequences for the Transformer model.

🧠 Model Training & Performance

The generative model utilizes a custom GPT2Config setup, optimized for local training on Apple Silicon.

• Configuration: 12 Layers, 12 Attention Heads, 768 Embedding Dimension.
• Training Duration: 3 Epochs
• Batch Size: 4 (Sequence Length: 512)
• Optimizer: AdamW with Linear Schedule Warmup
• Final Evaluation: The model successfully converged, reducing the loss from 6.25 to 4.82 in approximately 15 minutes of MPS-accelerated training.

📂 Repository Structure

stat-language-model-project/
├── data/
│   ├── raw_pdfs/             # Input statistical textbooks
│   ├── raw_txts/             # Per-book extracted raw text
│   └── processed/            # Cleaned, deduped, and tokenized data
├── scripts/
│   ├── extract_text.py       # PDF to text conversion
│   ├── combine_texts.py      # Corpus aggregation
│   ├── clean_text.py         # Regex sanitization
│   ├── deduplicate.py        # MD5 hashing deduplication
│   └── tokenize_corpus.py    # GPT-2 tokenization and tensor saving
├── model/
│   ├── model_config.py       # PyTorch model definition
│   ├── train.py              # Main training loop with MPS support
│   └── generate.py           # Inference script for text generation
└── experiments/
    └── checkpoints/          # Saved model weights (.pt files)

🚀 Usage & Generation

To utilize the text generation capabilities of this model, ensure you have the required dependencies installed and run the inference script.

1. Install Dependencies

pip install torch transformers pdfplumber tqdm

2. Run Inference Using the generate.py script, the model generates synthetic statistical text based on user-provided prompts. The script automatically detects and utilizes Apple's MPS backend if available.

python model/generate.py

Sample Prompts included in the generation script:

• "The mean is"
• "The probability of"
• "The variance measures"

📈 Known Limitations & Future Scope

• Extraction Artifacts: Due to the complexities of PDF font encodings, some early generation outputs occasionally produce artifacts (e.g., (cid:2)). Future iterations will implement stricter OCR or PyMuPDF extraction to resolve missing math symbols.
• Corpus Scaling: The current model is trained on a foundational corpus of 3 textbooks. The immediate next step is scaling the data pipeline to process 10+ statistical textbooks to expand the model's vocabulary and conceptual understanding.