rewrite / Plan.MD

Add files using upload-large-folder tool

3df5819 verified 6 days ago

119 kB

	# Dyslexia Academic Writing Correction System
	## Complete End-to-End Implementation Blueprint for Coding Agents

	> System Goal: A style-preserving, grammar-correcting, academic vocabulary elevating AI model that corrects dyslectic writing while maintaining the author's personal voice, tone, and authorship signal — not a rewriter, a corrector.

	---

	## Table of Contents

	1. [Repository Structure](#1-repository-structure)
	2. [Environment Setup](#2-environment-setup)
	3. [Dependency Manifest](#3-dependency-manifest)
	4. [System Architecture Overview](#4-system-architecture-overview)
	5. [Layer 1 — Input Pre-Processing Pipeline](#5-layer-1--input-pre-processing-pipeline)
	6. [Layer 2 — Style Fingerprinting Module](#6-layer-2--style-fingerprinting-module)
	7. [Layer 3 — Core Generation Model](#7-layer-3--core-generation-model)
	8. [Layer 4 — Training Data Strategy](#8-layer-4--training-data-strategy)
	9. [Layer 5 — Training Loop & Loss Functions](#9-layer-5--training-loop--loss-functions)
	10. [Layer 6 — Academic Vocabulary Control Module](#10-layer-6--academic-vocabulary-control-module)
	11. [Layer 7 — Evaluation Framework](#11-layer-7--evaluation-framework)
	12. [Layer 8 — Inference Pipeline](#12-layer-8--inference-pipeline)
	13. [Layer 9 — API Server](#13-layer-9--api-server)
	14. [Layer 10 — Configuration Files](#14-layer-10--configuration-files)
	15. [Layer 11 — Full Training Run Sequence](#15-layer-11--full-training-run-sequence)
	16. [Mathematical Formulations](#16-mathematical-formulations)
	17. [Hyperparameter Reference](#17-hyperparameter-reference)
	18. [Dataset Sources & Download Instructions](#18-dataset-sources--download-instructions)
	19. [Hardware Requirements](#19-hardware-requirements)
	20. [Testing Suite](#20-testing-suite)

	---

	## 1. Repository Structure

	```
	dyslexia-writing-ai/
	│
	├── configs/
	│ ├── model_config.yaml
	│ ├── training_config.yaml
	│ ├── inference_config.yaml
	│ └── awl_config.yaml
	│
	├── data/
	│ ├── raw/
	│ │ ├── wi_locness/
	│ │ ├── jfleg/
	│ │ ├── gyafc/
	│ │ └── custom_dyslexia/
	│ ├── processed/
	│ │ ├── train.jsonl
	│ │ ├── val.jsonl
	│ │ └── test.jsonl
	│ └── awl/
	│ ├── coxhead_awl.txt
	│ ├── academic_synonyms.json
	│ └── domain_lexicons/
	│ ├── humanities.txt
	│ ├── sciences.txt
	│ └── social_sciences.txt
	│
	├── src/
	│ ├── preprocessing/
	│ │ ├── __init__.py
	│ │ ├── spell_corrector.py
	│ │ ├── sentence_segmenter.py
	│ │ ├── dependency_parser.py
	│ │ ├── ner_tagger.py
	│ │ ├── dyslexia_simulator.py
	│ │ └── pipeline.py
	│ │
	│ ├── style/
	│ │ ├── __init__.py
	│ │ ├── fingerprinter.py
	│ │ ├── formality_classifier.py
	│ │ ├── emotion_classifier.py
	│ │ └── style_vector.py
	│ │
	│ ├── model/
	│ │ ├── __init__.py
	│ │ ├── base_model.py
	│ │ ├── lora_adapter.py
	│ │ ├── style_conditioner.py
	│ │ └── generation_utils.py
	│ │
	│ ├── training/
	│ │ ├── __init__.py
	│ │ ├── dataset.py
	│ │ ├── loss_functions.py
	│ │ ├── trainer.py
	│ │ └── callbacks.py
	│ │
	│ ├── vocabulary/
	│ │ ├── __init__.py
	│ │ ├── awl_loader.py
	│ │ ├── lexical_substitution.py
	│ │ └── register_filter.py
	│ │
	│ ├── evaluation/
	│ │ ├── __init__.py
	│ │ ├── gleu_scorer.py
	│ │ ├── errant_evaluator.py
	│ │ ├── style_metrics.py
	│ │ └── authorship_verifier.py
	│ │
	│ ├── inference/
	│ │ ├── __init__.py
	│ │ ├── corrector.py
	│ │ └── postprocessor.py
	│ │
	│ └── api/
	│ ├── __init__.py
	│ ├── main.py
	│ ├── schemas.py
	│ └── middleware.py
	│
	├── scripts/
	│ ├── download_datasets.sh
	│ ├── preprocess_data.py
	│ ├── train.py
	│ ├── evaluate.py
	│ └── run_inference.py
	│
	├── tests/
	│ ├── test_preprocessing.py
	│ ├── test_style.py
	│ ├── test_model.py
	│ ├── test_vocabulary.py
	│ └── test_evaluation.py
	│
	├── notebooks/
	│ ├── 01_data_exploration.ipynb
	│ ├── 02_style_fingerprint_analysis.ipynb
	│ ├── 03_training_diagnostics.ipynb
	│ └── 04_evaluation_dashboard.ipynb
	│
	├── requirements.txt
	├── requirements-dev.txt
	├── pyproject.toml
	├── Dockerfile
	├── docker-compose.yml
	└── README.md
	```

	---

	## 2. Environment Setup

	```bash
	# Python version requirement
	python >= 3.10

	# Create virtual environment
	python -m venv venv
	source venv/bin/activate # Linux/Mac
	# venv\Scripts\activate # Windows

	# Install PyTorch with CUDA (choose your CUDA version)
	pip install torch==2.2.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

	# Install all dependencies
	pip install -r requirements.txt

	# Download spaCy transformer model
	python -m spacy download en_core_web_trf

	# Download NLTK data
	python -c "import nltk; nltk.download('punkt'); nltk.download('averaged_perceptron_tagger'); nltk.download('wordnet')"

	# Install LanguageTool server (Java required)
	pip install language-tool-python
	# It auto-downloads the LanguageTool JAR on first run

	# Setup Weights & Biases for experiment tracking
	wandb login
	```

	---

	## 3. Dependency Manifest

	### `requirements.txt`

	```txt
	# ── Core ML & Deep Learning ──────────────────────────────────────────────────
	torch==2.2.0
	torchvision==0.17.0
	torchaudio==2.2.0
	transformers==4.40.0
	datasets==2.18.0
	accelerate==0.29.0
	peft==0.10.0 # LoRA / parameter-efficient fine-tuning
	bitsandbytes==0.43.0 # 8-bit & 4-bit quantization
	sentencepiece==0.2.0 # T5 tokenizer dependency
	protobuf==4.25.3 # T5 tokenizer dependency

	# ── Sentence Embeddings ───────────────────────────────────────────────────────
	sentence-transformers==2.6.1
	faiss-cpu==1.8.0 # Vector similarity search

	# ── NLP Pre-Processing ────────────────────────────────────────────────────────
	spacy==3.7.4
	spacy-transformers==1.3.4
	language-tool-python==2.7.1 # LanguageTool grammar checker
	pyspellchecker==0.8.1 # Context-free spell check (pre-pass)
	nltk==3.8.1
	textstat==0.7.3 # Readability scores (Flesch-Kincaid, etc.)

	# ── Lexical Substitution ─────────────────────────────────────────────────────
	lexsubgen==0.0.4 # BERT-based lexical substitution
	wordfreq==3.1.1 # Word frequency data
	PyDictionary==2.0.1

	# ── Training Infrastructure ───────────────────────────────────────────────────
	wandb==0.16.6 # Experiment tracking
	tensorboard==2.16.2
	numpy==1.26.4
	pandas==2.2.1
	scikit-learn==1.4.1.post1
	scipy==1.13.0

	# ── Evaluation Tools ──────────────────────────────────────────────────────────
	errant==2.3.3 # Grammar Error Annotation Toolkit
	sacrebleu==2.4.2 # BLEU/GLEU scoring
	bert-score==0.3.13 # Semantic similarity scoring
	rouge-score==0.1.2

	# ── API Server ────────────────────────────────────────────────────────────────
	fastapi==0.110.1
	uvicorn[standard]==0.29.0
	pydantic==2.7.0
	python-multipart==0.0.9
	httpx==0.27.0

	# ── Inference Optimisation ────────────────────────────────────────────────────
	vllm==0.4.0 # High-throughput LLM serving (optional, GPU only)
	optimum==1.19.1 # Hugging Face model optimisation

	# ── Utilities ─────────────────────────────────────────────────────────────────
	pyyaml==6.0.1
	tqdm==4.66.2
	loguru==0.7.2
	python-dotenv==1.0.1
	click==8.1.7
	rich==13.7.1 # Beautiful terminal output
	joblib==1.4.0
	```

	### `requirements-dev.txt`

	```txt
	pytest==8.1.1
	pytest-asyncio==0.23.6
	pytest-cov==5.0.0
	black==24.4.0
	ruff==0.4.1
	mypy==1.9.0
	pre-commit==3.7.0
	ipykernel==6.29.4
	jupyter==1.0.0
	```

	---

	## 4. System Architecture Overview

	```
	┌─────────────────────────────────────────────────────────────────────┐
	│ INPUT TEXT (raw dyslectic) │
	└─────────────────────────────────────────────┬───────────────────────┘
	│
	┌─────────────────────────▼──────────────────────┐
	│ LAYER 1: Pre-Processing Pipeline │
	│ spell_corrector → segmenter → dep_parser → │
	│ NER_tagger → readability_scorer │
	└─────────────────────────┬──────────────────────┘
	│ cleaned + annotated text
	┌─────────────────────────▼──────────────────────┐
	│ LAYER 2: Style Fingerprinting │
	│ sentence_len_dist, syntactic_complexity, │
	│ TTR, voice_ratio, hedging_freq, │
	│ discourse_markers, formality_score, │
	│ emotion_register → style_vector [512-dim] │
	└─────────────────────────┬──────────────────────┘
	│ │
	[user style vec] [master copy style vec]
	│ │
	┌─────────▼───────────────▼──────────────────────┐
	│ STYLE BLENDING (weighted interpolation) │
	│ target_style = α·user + (1-α)·master │
	└─────────────────────────┬──────────────────────┘
	│
	┌─────────────────────────▼──────────────────────┐
	│ LAYER 3: Core Generation Model │
	│ Base: Flan-T5-XL / BART-large / Llama-3 │
	│ Fine-tuned with LoRA │
	│ Conditioned on: cleaned_text + style_vector │
	│ Loss: CE + style_consistency + semantic_sim │
	└─────────────────────────┬──────────────────────┘
	│ draft corrected text
	┌─────────────────────────▼──────────────────────┐
	│ LAYER 6: Academic Vocabulary Control │
	│ AWL substitution → register filter │
	│ → nominalisation pass → hedging check │
	└─────────────────────────┬──────────────────────┘
	│
	┌─────────────────────────▼──────────────────────┐
	│ LAYER 7: Evaluation & Quality Gate │
	│ GLEU, ERRANT, style_sim, authorship_score │
	│ If quality < threshold → re-generate │
	└─────────────────────────┬──────────────────────┘
	│
	┌─────────────────────────▼──────────────────────┐
	│ FINAL OUTPUT │
	│ Grammatically perfect · Academic register │
	│ Style-preserved · Human authorship signal │
	└────────────────────────────────────────────────┘
	```

	---

	## 5. Layer 1 — Input Pre-Processing Pipeline

	### `src/preprocessing/spell_corrector.py`

	```python
	"""
	Two-pass spell correction:
	Pass 1: pyspellchecker (fast, context-free, catches simple typos)
	Pass 2: LanguageTool (context-aware, catches grammar + dyslexic patterns)

	Dyslexic error patterns handled:
	- Letter reversals: b/d, p/q, n/u, m/w
	- Phonetic spelling: "wuz", "cud", "thay"
	- Word boundary errors: "alot", "infact", "aswell"
	- Letter omissions: "becaus", "importnt"
	- Letter transpositions: "teh", "recieve"
	- Homophone confusion: there/their/they're
	"""

	import language_tool_python
	from spellchecker import SpellChecker
	from loguru import logger
	from typing import Optional
	import re


	class DyslexiaAwareSpellCorrector:

	DYSLEXIC_PHONETIC_MAP = {
	"wuz": "was", "cud": "could", "wud": "would", "shud": "should",
	"thay": "they", "thier": "their", "recieve": "receive",
	"beleive": "believe", "occured": "occurred", "definately": "definitely",
	"seperate": "separate", "untill": "until", "tommorrow": "tomorrow",
	"alot": "a lot", "infact": "in fact", "aswell": "as well",
	"alright": "all right", "cant": "cannot", "wont": "will not",
	"ive": "I have", "im": "I am", "id": "I would",
	}

	def __init__(self, language: str = "en-US"):
	self.spell = SpellChecker()
	self.tool = language_tool_python.LanguageTool(language)
	logger.info("Spell corrector initialised with LanguageTool backend.")

	def _phonetic_pass(self, text: str) -> str:
	"""Apply known dyslexic phonetic substitutions first."""
	pattern = re.compile(
	r'\b(' + '\|'.join(re.escape(k) for k in self.DYSLEXIC_PHONETIC_MAP.keys()) + r')\b',
	re.IGNORECASE
	)
	def replace(match):
	return self.DYSLEXIC_PHONETIC_MAP[match.group(0).lower()]
	return pattern.sub(replace, text)

	def _spellcheck_pass(self, text: str) -> str:
	"""pyspellchecker pass for simple token-level errors."""
	tokens = text.split()
	corrected = []
	for token in tokens:
	clean = re.sub(r'[^\w]', '', token).lower()
	if clean and clean not in self.spell:
	correction = self.spell.correction(clean)
	if correction:
	token = token.replace(clean, correction)
	corrected.append(token)
	return ' '.join(corrected)

	def _languagetool_pass(self, text: str) -> str:
	"""LanguageTool pass for context-aware grammar + spelling corrections."""
	matches = self.tool.check(text)
	# Apply corrections in reverse order to preserve offsets
	for match in reversed(matches):
	if match.replacements:
	start = match.offset
	end = start + match.errorLength
	text = text[:start] + match.replacements[0] + text[end:]
	return text

	def correct(self, text: str) -> str:
	text = self._phonetic_pass(text)
	text = self._spellcheck_pass(text)
	text = self._languagetool_pass(text)
	return text

	def close(self):
	self.tool.close()
	```

	---

	### `src/preprocessing/pipeline.py`

	```python
	"""
	Master pre-processing pipeline. Runs all NLP stages in sequence.
	Returns a PreprocessedDoc object with all annotations attached.
	"""

	import spacy
	from dataclasses import dataclass, field
	from typing import List, Dict, Any, Optional
	from .spell_corrector import DyslexiaAwareSpellCorrector
	import textstat


	@dataclass
	class EntitySpan:
	text: str
	label: str
	start_char: int
	end_char: int


	@dataclass
	class PreprocessedDoc:
	original_text: str
	corrected_text: str
	sentences: List[str]
	entities: List[EntitySpan] # Never to be modified by rewriter
	dependency_trees: List[Dict] # Grammatical skeletons per sentence
	pos_tags: List[List[tuple]] # (token, POS) per sentence
	readability: Dict[str, float] # Flesch-Kincaid, Gunning Fog, etc.
	sentence_lengths: List[int]
	protected_spans: List[tuple] # (start, end) char spans to never touch


	class PreprocessingPipeline:

	def __init__(self, model_name: str = "en_core_web_trf"):
	self.nlp = spacy.load(model_name)
	self.corrector = DyslexiaAwareSpellCorrector()

	def _extract_readability(self, text: str) -> Dict[str, float]:
	return {
	"flesch_reading_ease": textstat.flesch_reading_ease(text),
	"flesch_kincaid_grade": textstat.flesch_kincaid_grade(text),
	"gunning_fog": textstat.gunning_fog(text),
	"smog_index": textstat.smog_index(text),
	"automated_readability_index": textstat.automated_readability_index(text),
	}

	def _extract_dep_tree(self, sent) -> Dict:
	"""Extract grammatical skeleton: subject-verb-object per sentence."""
	tree = {"tokens": [], "root": None, "svo": []}
	subjects, verbs, objects = [], [], []
	for token in sent:
	tree["tokens"].append({
	"text": token.text,
	"dep": token.dep_,
	"pos": token.pos_,
	"head": token.head.text,
	})
	if token.dep_ == "ROOT":
	tree["root"] = token.text
	if token.dep_ in ("nsubj", "nsubjpass"):
	subjects.append(token.text)
	if token.pos_ == "VERB":
	verbs.append(token.text)
	if token.dep_ in ("dobj", "pobj"):
	objects.append(token.text)
	tree["svo"] = {"subjects": subjects, "verbs": verbs, "objects": objects}
	return tree

	def process(self, raw_text: str) -> PreprocessedDoc:
	# Step 1: Spell + grammar correction
	corrected = self.corrector.correct(raw_text)

	# Step 2: spaCy full parse
	doc = self.nlp(corrected)

	# Step 3: Extract sentences
	sentences = [sent.text.strip() for sent in doc.sents]
	sentence_lengths = [len(sent.text.split()) for sent in doc.sents]

	# Step 4: Named entities (protect these spans)
	entities = [
	EntitySpan(ent.text, ent.label_, ent.start_char, ent.end_char)
	for ent in doc.ents
	]
	protected_spans = [(e.start_char, e.end_char) for e in entities]

	# Step 5: Dependency trees
	dep_trees = [self._extract_dep_tree(sent) for sent in doc.sents]

	# Step 6: POS tags
	pos_tags = [
	[(token.text, token.pos_) for token in sent]
	for sent in doc.sents
	]

	# Step 7: Readability
	readability = self._extract_readability(corrected)

	return PreprocessedDoc(
	original_text=raw_text,
	corrected_text=corrected,
	sentences=sentences,
	entities=entities,
	dependency_trees=dep_trees,
	pos_tags=pos_tags,
	readability=readability,
	sentence_lengths=sentence_lengths,
	protected_spans=protected_spans,
	)
	```

	---

	### `src/preprocessing/dyslexia_simulator.py`

	```python
	"""
	Programmatically generates dyslectic training data from clean text.
	Used to augment training pairs when real dyslectic examples are scarce.

	Error types simulated (from Rello et al. 2013, 2017 dyslexia research):
	- Phonetic substitution (most common, ~35% of errors)
	- Letter transposition (e.g., "teh" for "the") (~18%)
	- Letter omission (~16%)
	- Letter doubling (~12%)
	- Letter reversal b/d, p/q (~10%)
	- Word boundary errors (~9%)
	"""

	import random
	import re
	from typing import Tuple


	class DyslexiaSimulator:

	LETTER_REVERSALS = {'b': 'd', 'd': 'b', 'p': 'q', 'q': 'p', 'n': 'u', 'u': 'n'}
	PHONETIC_SUBS = {
	'was': 'wuz', 'could': 'cud', 'would': 'wud', 'they': 'thay',
	'because': 'becaus', 'important': 'importnt', 'receive': 'recieve',
	'believe': 'beleive', 'definitely': 'definately', 'separate': 'seperate',
	'a lot': 'alot', 'in fact': 'infact', 'as well': 'aswell',
	}
	WORD_MERGES = [
	('a lot', 'alot'), ('in fact', 'infact'), ('as well', 'aswell'),
	('all right', 'alright'), ('every one', 'everyone'),
	]

	def __init__(self, error_rate: float = 0.15, seed: int = 42):
	self.error_rate = error_rate
	random.seed(seed)

	def _transpose_letters(self, word: str) -> str:
	if len(word) < 3:
	return word
	i = random.randint(0, len(word) - 2)
	chars = list(word)
	chars[i], chars[i+1] = chars[i+1], chars[i]
	return ''.join(chars)

	def _omit_letter(self, word: str) -> str:
	if len(word) < 4:
	return word
	i = random.randint(1, len(word) - 2)
	return word[:i] + word[i+1:]

	def _double_letter(self, word: str) -> str:
	if len(word) < 3:
	return word
	i = random.randint(1, len(word) - 2)
	return word[:i] + word[i] + word[i:]

	def _reverse_letter(self, word: str) -> str:
	chars = list(word)
	for i, c in enumerate(chars):
	if c in self.LETTER_REVERSALS and random.random() < 0.5:
	chars[i] = self.LETTER_REVERSALS[c]
	return ''.join(chars)

	def corrupt_word(self, word: str) -> str:
	"""Apply a single random error to a word."""
	if len(word) <= 2 or random.random() > self.error_rate:
	return word
	# Check phonetic substitutions first
	lower = word.lower()
	if lower in self.PHONETIC_SUBS:
	return self.PHONETIC_SUBS[lower]
	choice = random.choice(['transpose', 'omit', 'double', 'reverse'])
	if choice == 'transpose':
	return self._transpose_letters(word)
	elif choice == 'omit':
	return self._omit_letter(word)
	elif choice == 'double':
	return self._double_letter(word)
	else:
	return self._reverse_letter(word)

	def simulate(self, clean_text: str) -> Tuple[str, str]:
	"""Returns (corrupted_text, clean_text) training pair."""
	words = clean_text.split()
	corrupted = [self.corrupt_word(w) for w in words]
	corrupted_text = ' '.join(corrupted)
	# Apply word merge errors
	for correct_phrase, merged in self.WORD_MERGES:
	if random.random() < 0.3:
	corrupted_text = corrupted_text.replace(correct_phrase, merged)
	return corrupted_text, clean_text
	```

	---

	## 6. Layer 2 — Style Fingerprinting Module

	### `src/style/fingerprinter.py`

	```python
	"""
	Extracts a numerical style vector from any text sample.
	The style vector encodes the author's unique writing fingerprint
	and is used both to condition the generation model and to evaluate
	style preservation after correction.

	Style vector dimensions (total: 512 after projection):
	Raw features (~40) → MLP projection → 512-dim dense vector

	Raw features:
	- sentence_length_mean, sentence_length_std, sentence_length_skew [3]
	- word_length_mean, word_length_std [2]
	- type_token_ratio (TTR) [1]
	- passive_voice_ratio [1]
	- active_voice_ratio [1]
	- subordinate_clause_ratio [1]
	- avg_dependency_tree_depth [1]
	- hedging_frequency (per 100 words) [1]
	- discourse_marker_counts [however, therefore, moreover, ...] [20]
	- formality_score (0-1) [1]
	- lexical_density [1]
	- nominalization_ratio [1]
	- question_sentence_ratio [1]
	- exclamation_ratio [1]
	- first_person_ratio [1]
	- third_person_ratio [1]
	- academic_word_coverage [1]
	- avg_syllables_per_word [1]
	- flesch_reading_ease [1]
	"""

	import spacy
	import numpy as np
	import torch
	import torch.nn as nn
	from typing import List, Dict, Optional
	from scipy import stats


	HEDGING_WORDS = {
	"perhaps", "possibly", "probably", "might", "may", "could", "seem",
	"appears", "suggests", "indicates", "tend", "often", "generally",
	"approximately", "roughly", "somewhat", "relatively", "fairly",
	}

	DISCOURSE_MARKERS = [
	"however", "therefore", "moreover", "furthermore", "consequently",
	"nevertheless", "nonetheless", "additionally", "alternatively",
	"subsequently", "previously", "similarly", "conversely", "thus",
	"hence", "accordingly", "meanwhile", "indeed", "notably", "specifically",
	]

	NOMINALISATION_SUFFIXES = (
	"tion", "sion", "ment", "ness", "ity", "ance", "ence",
	"hood", "ship", "ism", "al", "ure",
	)


	class StyleProjectionMLP(nn.Module):
	"""Projects raw feature vector to 512-dim style embedding."""
	def __init__(self, input_dim: int = 40, hidden_dim: int = 256, output_dim: int = 512):
	super().__init__()
	self.net = nn.Sequential(
	nn.Linear(input_dim, hidden_dim),
	nn.LayerNorm(hidden_dim),
	nn.GELU(),
	nn.Dropout(0.1),
	nn.Linear(hidden_dim, output_dim),
	nn.LayerNorm(output_dim),
	)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	return self.net(x)


	class StyleFingerprinter:

	def __init__(self, spacy_model: str = "en_core_web_trf", awl_path: str = "data/awl/coxhead_awl.txt"):
	self.nlp = spacy.load(spacy_model)
	self.awl = self._load_awl(awl_path)
	self.projection = StyleProjectionMLP()

	def _load_awl(self, path: str) -> set:
	try:
	with open(path) as f:
	return {line.strip().lower() for line in f if line.strip()}
	except FileNotFoundError:
	return set()

	def _passive_voice_ratio(self, doc) -> float:
	passive_count = sum(
	1 for token in doc
	if token.dep_ in ("nsubjpass", "auxpass")
	)
	sentences = list(doc.sents)
	return passive_count / max(len(sentences), 1)

	def _avg_dep_tree_depth(self, doc) -> float:
	def depth(token):
	d = 0
	while token.head != token:
	token = token.head
	d += 1
	return d
	depths = [depth(token) for token in doc]
	return np.mean(depths) if depths else 0.0

	def _lexical_density(self, doc) -> float:
	content_pos = {"NOUN", "VERB", "ADJ", "ADV"}
	content = sum(1 for t in doc if t.pos_ in content_pos)
	return content / max(len(doc), 1)

	def extract_raw_features(self, text: str) -> Dict[str, float]:
	doc = self.nlp(text)
	sentences = list(doc.sents)
	words = [t.text for t in doc if not t.is_punct and not t.is_space]
	word_lengths = [len(w) for w in words]
	sent_lengths = [len(list(s)) for s in sentences]

	# Type-Token Ratio
	unique_words = set(w.lower() for w in words)
	ttr = len(unique_words) / max(len(words), 1)

	# Hedging
	hedging_freq = sum(1 for w in words if w.lower() in HEDGING_WORDS)
	hedging_per_100 = (hedging_freq / max(len(words), 1)) * 100

	# Discourse markers
	text_lower = text.lower()
	dm_counts = {dm: text_lower.count(dm) for dm in DISCOURSE_MARKERS}

	# Nominalisation
	nom_count = sum(1 for w in words if w.lower().endswith(NOMINALISATION_SUFFIXES))
	nom_ratio = nom_count / max(len(words), 1)

	# Person
	first_p = sum(1 for t in doc if t.lower_ in {"i", "we", "my", "our", "me", "us"})
	third_p = sum(1 for t in doc if t.lower_ in {"he", "she", "they", "it", "his", "her", "their"})
	first_ratio = first_p / max(len(words), 1)
	third_ratio = third_p / max(len(words), 1)

	# AWL coverage
	awl_hits = sum(1 for w in words if w.lower() in self.awl)
	awl_coverage = awl_hits / max(len(words), 1)

	# Sentence type
	question_ratio = sum(1 for s in sentences if s.text.strip().endswith("?")) / max(len(sentences), 1)
	exclaim_ratio = sum(1 for s in sentences if s.text.strip().endswith("!")) / max(len(sentences), 1)

	raw = {
	"sent_len_mean": np.mean(sent_lengths) if sent_lengths else 0,
	"sent_len_std": np.std(sent_lengths) if sent_lengths else 0,
	"sent_len_skew": float(stats.skew(sent_lengths)) if len(sent_lengths) > 2 else 0,
	"word_len_mean": np.mean(word_lengths) if word_lengths else 0,
	"word_len_std": np.std(word_lengths) if word_lengths else 0,
	"ttr": ttr,
	"passive_ratio": self._passive_voice_ratio(doc),
	"active_ratio": 1.0 - self._passive_voice_ratio(doc),
	"avg_dep_depth": self._avg_dep_tree_depth(doc),
	"hedging_per_100": hedging_per_100,
	"nom_ratio": nom_ratio,
	"lexical_density": self._lexical_density(doc),
	"question_ratio": question_ratio,
	"exclaim_ratio": exclaim_ratio,
	"first_person_ratio": first_ratio,
	"third_person_ratio": third_ratio,
	"awl_coverage": awl_coverage,
	}
	for dm, count in dm_counts.items():
	raw[f"dm_{dm}"] = count / max(len(sentences), 1)

	return raw

	def extract_vector(self, text: str) -> torch.Tensor:
	"""Returns a 512-dim style embedding tensor."""
	raw = self.extract_raw_features(text)
	feature_array = np.array(list(raw.values()), dtype=np.float32)
	# Pad or truncate to expected input_dim
	expected_dim = 40
	if len(feature_array) < expected_dim:
	feature_array = np.pad(feature_array, (0, expected_dim - len(feature_array)))
	else:
	feature_array = feature_array[:expected_dim]
	feature_tensor = torch.tensor(feature_array).unsqueeze(0)
	with torch.no_grad():
	style_vec = self.projection(feature_tensor)
	return style_vec.squeeze(0) # [512]

	def blend_vectors(
	self,
	user_vec: torch.Tensor,
	master_vec: Optional[torch.Tensor],
	alpha: float = 0.6,
	) -> torch.Tensor:
	"""
	Blend user style with master copy style.
	alpha = weight given to user's own style (0.6 = user dominates)
	(1-alpha) = weight given to master copy style

	Formula: target = alpha * user_vec + (1 - alpha) * master_vec
	"""
	if master_vec is None:
	return user_vec
	blended = alpha * user_vec + (1 - alpha) * master_vec
	# L2 normalise to unit sphere
	return blended / (blended.norm() + 1e-8)
	```

	---

	## 7. Layer 3 — Core Generation Model

	### Model Selection Decision Tree

	```
	Do you have ≥ 40GB VRAM (e.g., A100)?
	├── YES → Fine-tune Llama-3.1-8B with LoRA (best quality)
	└── NO → Do you have ≥ 16GB VRAM?
	├── YES → Fine-tune Flan-T5-XL (3B params, best encoder-decoder)
	└── NO → Fine-tune BART-large (400M params, excellent denoiser)
	OR Flan-T5-Large (780M params)
	```

	### `src/model/base_model.py`

	```python
	"""
	Loads and wraps the base pretrained model.
	Supported architectures:
	- google/flan-t5-xl (recommended, 3B)
	- google/flan-t5-large (780M, resource-constrained)
	- facebook/bart-large (400M, excellent denoiser)
	- meta-llama/Meta-Llama-3.1-8B-Instruct (8B, best quality)
	"""

	from transformers import (
	AutoTokenizer, AutoModelForSeq2SeqLM,
	AutoModelForCausalLM, BitsAndBytesConfig
	)
	from peft import get_peft_model, LoraConfig, TaskType
	import torch
	from loguru import logger


	ENCODER_DECODER_MODELS = {
	"flan-t5-xl": "google/flan-t5-xl",
	"flan-t5-large": "google/flan-t5-large",
	"bart-large": "facebook/bart-large",
	}

	DECODER_ONLY_MODELS = {
	"llama-3.1-8b": "meta-llama/Meta-Llama-3.1-8B-Instruct",
	}


	def load_model_and_tokenizer(model_key: str, quantize: bool = False, use_lora: bool = True):

	is_seq2seq = model_key in ENCODER_DECODER_MODELS
	model_name = ENCODER_DECODER_MODELS.get(model_key) or DECODER_ONLY_MODELS.get(model_key)

	if not model_name:
	raise ValueError(f"Unknown model key: {model_key}")

	logger.info(f"Loading {model_name} ({'seq2seq' if is_seq2seq else 'causal'})...")

	tokenizer = AutoTokenizer.from_pretrained(model_name)
	if tokenizer.pad_token is None:
	tokenizer.pad_token = tokenizer.eos_token

	# Quantisation config (for large models on limited VRAM)
	bnb_config = None
	if quantize:
	bnb_config = BitsAndBytesConfig(
	load_in_4bit=True,
	bnb_4bit_use_double_quant=True,
	bnb_4bit_quant_type="nf4",
	bnb_4bit_compute_dtype=torch.bfloat16,
	)

	if is_seq2seq:
	model = AutoModelForSeq2SeqLM.from_pretrained(
	model_name,
	quantization_config=bnb_config,
	torch_dtype=torch.bfloat16 if not quantize else None,
	device_map="auto",
	)
	lora_task = TaskType.SEQ_2_SEQ_LM
	else:
	model = AutoModelForCausalLM.from_pretrained(
	model_name,
	quantization_config=bnb_config,
	torch_dtype=torch.bfloat16 if not quantize else None,
	device_map="auto",
	)
	lora_task = TaskType.CAUSAL_LM

	if use_lora:
	lora_config = LoraConfig(
	task_type=lora_task,
	r=16, # LoRA rank — increase for more capacity
	lora_alpha=32, # Scaling factor (typically 2x rank)
	target_modules=[ # Modules to apply LoRA to
	"q_proj", "v_proj", # Attention query and value
	"k_proj", "o_proj", # Attention key and output
	"gate_proj", "up_proj", # FFN layers (for T5/Llama)
	],
	lora_dropout=0.05,
	bias="none",
	inference_mode=False,
	)
	model = get_peft_model(model, lora_config)
	model.print_trainable_parameters()

	return model, tokenizer, is_seq2seq
	```

	---

	### `src/model/style_conditioner.py`

	```python
	"""
	Injects the style vector into the model via soft prompt conditioning.
	The style vector is projected to the model's hidden dimension and
	prepended to the input token embeddings as virtual tokens.

	This technique is called "prefix tuning" / "style prefix injection".
	It biases the model's attention toward the desired output style
	without modifying the base model weights.

	For Flan-T5: injects into encoder input embeddings
	For BART: injects into encoder input embeddings
	For Llama: prepends to the full input context
	"""

	import torch
	import torch.nn as nn


	class StyleConditioner(nn.Module):
	"""
	Projects a 512-dim style vector to n_prefix_tokens virtual tokens
	in the model's embedding space.
	"""

	def __init__(
	self,
	style_dim: int = 512,
	model_hidden_dim: int = 2048, # T5-XL hidden size
	n_prefix_tokens: int = 10, # Number of virtual prefix tokens
	):
	super().__init__()
	self.n_prefix_tokens = n_prefix_tokens
	self.projection = nn.Sequential(
	nn.Linear(style_dim, model_hidden_dim * n_prefix_tokens),
	nn.Tanh(),
	)

	def forward(self, style_vector: torch.Tensor) -> torch.Tensor:
	"""
	Args:
	style_vector: [batch_size, 512]
	Returns:
	prefix_embeddings: [batch_size, n_prefix_tokens, model_hidden_dim]
	"""
	batch_size = style_vector.shape[0]
	projected = self.projection(style_vector)
	return projected.view(batch_size, self.n_prefix_tokens, -1)


	def prepend_style_prefix(
	input_embeddings: torch.Tensor,
	style_prefix: torch.Tensor,
	) -> torch.Tensor:
	"""
	Concatenates style prefix to input embeddings along sequence dimension.

	Args:
	input_embeddings: [batch, seq_len, hidden_dim]
	style_prefix: [batch, n_prefix, hidden_dim]
	Returns:
	[batch, n_prefix + seq_len, hidden_dim]
	"""
	return torch.cat([style_prefix, input_embeddings], dim=1)
	```

	---

	## 8. Layer 4 — Training Data Strategy

	### `src/training/dataset.py`

	```python
	"""
	Dataset class that handles all data sources and produces training triplets:
	(input_text, style_vector, target_text)

	Data sources priority:
	1. W&I+LOCNESS — real learner errors with expert corrections
	2. JFLEG — naturalistic fluency corrections
	3. GYAFC — informal→formal style transfer
	4. Synthetic — dyslexia simulator augmentation on Wikipedia/books
	5. Custom — any user-provided correction pairs

	Each example is structured as:
	{
	"input": "<corrupted/informal text>",
	"target": "<corrected academic text>",
	"style_vector": [512 floats],
	"source": "wi_locness \| jfleg \| gyafc \| synthetic \| custom",
	}
	"""

	import json
	from pathlib import Path
	from typing import List, Dict, Optional
	import torch
	from torch.utils.data import Dataset
	from transformers import PreTrainedTokenizer
	from ..style.fingerprinter import StyleFingerprinter
	from ..preprocessing.dyslexia_simulator import DyslexiaSimulator


	TASK_PREFIX = (
	"Correct the following text for grammar, spelling, and clarity. "
	"Maintain the author's original tone and writing style. "
	"Elevate vocabulary to academic register. "
	"Do NOT change the meaning or add new information. "
	"Preserve named entities exactly. "
	"Text to correct: "
	)


	class WritingCorrectionDataset(Dataset):

	def __init__(
	self,
	data_path: str,
	tokenizer: PreTrainedTokenizer,
	fingerprinter: StyleFingerprinter,
	max_input_length: int = 512,
	max_target_length: int = 512,
	augment_with_synthetic: bool = True,
	synthetic_ratio: float = 0.3,
	):
	self.tokenizer = tokenizer
	self.fingerprinter = fingerprinter
	self.max_input_length = max_input_length
	self.max_target_length = max_target_length
	self.examples = self._load(data_path)

	if augment_with_synthetic:
	self._add_synthetic(synthetic_ratio)

	def _load(self, path: str) -> List[Dict]:
	examples = []
	with open(path) as f:
	for line in f:
	obj = json.loads(line.strip())
	examples.append(obj)
	return examples

	def _add_synthetic(self, ratio: float):
	simulator = DyslexiaSimulator(error_rate=0.15)
	n_synthetic = int(len(self.examples) * ratio)
	# Use clean targets as source for simulation
	synthetic = []
	for ex in self.examples[:n_synthetic]:
	corrupted, clean = simulator.simulate(ex["target"])
	synthetic.append({"input": corrupted, "target": clean, "source": "synthetic"})
	self.examples.extend(synthetic)

	def __len__(self):
	return len(self.examples)

	def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
	ex = self.examples[idx]
	input_text = TASK_PREFIX + ex["input"]
	target_text = ex["target"]

	# Compute style vector from the TARGET (we want to learn to match this style)
	style_vec = self.fingerprinter.extract_vector(target_text)

	# Tokenise input
	input_enc = self.tokenizer(
	input_text,
	max_length=self.max_input_length,
	padding="max_length",
	truncation=True,
	return_tensors="pt",
	)

	# Tokenise target
	target_enc = self.tokenizer(
	target_text,
	max_length=self.max_target_length,
	padding="max_length",
	truncation=True,
	return_tensors="pt",
	)

	labels = target_enc["input_ids"].squeeze()
	labels[labels == self.tokenizer.pad_token_id] = -100 # Ignore padding in loss

	return {
	"input_ids": input_enc["input_ids"].squeeze(),
	"attention_mask": input_enc["attention_mask"].squeeze(),
	"labels": labels,
	"style_vector": style_vec,
	}
	```

	---

	## 9. Layer 5 — Training Loop & Loss Functions

	### `src/training/loss_functions.py`

	```python
	"""
	Combined training loss:

	L_total = L_CE + λ₁ · L_style + λ₂ · L_semantic

	Where:
	L_CE = cross-entropy language model loss (standard token prediction)
	L_style = style consistency loss (cosine distance between output and target style vectors)
	L_semantic = semantic similarity loss (cosine distance between sentence embeddings)
	λ₁ = style loss weight (default 0.3)
	λ₂ = semantic loss weight (default 0.5)

	L_style:
	style_sim = cosine_similarity(style_vec(output), style_vec(target))
	L_style = 1 - style_sim

	L_semantic:
	sem_emb_output = sentence_transformer.encode(output_text)
	sem_emb_input = sentence_transformer.encode(input_text)
	sem_sim = cosine_similarity(sem_emb_output, sem_emb_input)
	L_semantic = 1 - sem_sim
	(We compare to INPUT meaning — meaning must be preserved, not changed)
	"""

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from sentence_transformers import SentenceTransformer
	from typing import Optional


	class CombinedCorrectionLoss(nn.Module):

	def __init__(
	self,
	lambda_style: float = 0.3,
	lambda_semantic: float = 0.5,
	sem_model_name: str = "all-mpnet-base-v2",
	device: str = "cuda",
	):
	super().__init__()
	self.lambda_style = lambda_style
	self.lambda_semantic = lambda_semantic
	self.device = device

	# Frozen sentence transformer for semantic similarity
	self.sem_model = SentenceTransformer(sem_model_name, device=device)
	for param in self.sem_model.parameters():
	param.requires_grad = False

	self.ce_loss = nn.CrossEntropyLoss(ignore_index=-100)

	def _style_loss(
	self,
	output_style_vec: torch.Tensor,
	target_style_vec: torch.Tensor,
	) -> torch.Tensor:
	"""
	1 - cosine_similarity(output_style, target_style)
	Shape: [batch_size, 512] → scalar
	"""
	sim = F.cosine_similarity(output_style_vec, target_style_vec, dim=-1)
	return (1 - sim).mean()

	def _semantic_loss(
	self,
	input_texts: List[str],
	output_texts: List[str],
	) -> torch.Tensor:
	"""
	Penalises meaning change between input and output.
	Uses frozen sentence-transformer embeddings.
	"""
	with torch.no_grad():
	input_embs = torch.tensor(
	self.sem_model.encode(input_texts), device=self.device
	)
	output_embs = torch.tensor(
	self.sem_model.encode(output_texts), device=self.device
	)
	sim = F.cosine_similarity(input_embs, output_embs, dim=-1)
	return (1 - sim).mean()

	def forward(
	self,
	logits: torch.Tensor,
	labels: torch.Tensor,
	output_style_vec: Optional[torch.Tensor] = None,
	target_style_vec: Optional[torch.Tensor] = None,
	input_texts: Optional[List[str]] = None,
	output_texts: Optional[List[str]] = None,
	) -> Dict[str, torch.Tensor]:

	# Standard cross-entropy loss
	vocab_size = logits.shape[-1]
	l_ce = self.ce_loss(logits.view(-1, vocab_size), labels.view(-1))

	losses = {"l_ce": l_ce, "total": l_ce}

	if output_style_vec is not None and target_style_vec is not None:
	l_style = self._style_loss(output_style_vec, target_style_vec)
	losses["l_style"] = l_style
	losses["total"] = losses["total"] + self.lambda_style * l_style

	if input_texts is not None and output_texts is not None:
	l_sem = self._semantic_loss(input_texts, output_texts)
	losses["l_semantic"] = l_sem
	losses["total"] = losses["total"] + self.lambda_semantic * l_sem

	return losses
	```

	---

	### `src/training/trainer.py`

	```python
	"""
	Custom HuggingFace Trainer subclass.
	Overrides compute_loss to use CombinedCorrectionLoss.
	"""

	from transformers import Trainer, TrainingArguments
	from transformers.trainer_utils import EvalLoopOutput
	import torch
	from .loss_functions import CombinedCorrectionLoss
	import wandb


	class CorrectionTrainer(Trainer):

	def __init__(self, loss_fn: CombinedCorrectionLoss, fingerprinter, tokenizer, **kwargs):
	super().__init__(**kwargs)
	self.loss_fn = loss_fn
	self.fingerprinter = fingerprinter
	self.tokenizer = tokenizer

	def compute_loss(self, model, inputs, return_outputs=False):
	style_vectors = inputs.pop("style_vector", None)
	labels = inputs.get("labels")

	outputs = model(**inputs)
	logits = outputs.logits

	# Decode output tokens to text for semantic + style losses
	pred_token_ids = logits.argmax(dim=-1)
	output_texts = self.tokenizer.batch_decode(pred_token_ids, skip_special_tokens=True)
	input_texts = self.tokenizer.batch_decode(inputs["input_ids"], skip_special_tokens=True)

	# Compute output style vectors for batch
	output_style_vecs = torch.stack([
	self.fingerprinter.extract_vector(t) for t in output_texts
	]).to(logits.device)

	loss_dict = self.loss_fn(
	logits=logits,
	labels=labels,
	output_style_vec=output_style_vecs,
	target_style_vec=style_vectors.to(logits.device) if style_vectors is not None else None,
	input_texts=input_texts,
	output_texts=output_texts,
	)

	# Log to W&B
	if self.state.global_step % 50 == 0:
	wandb.log({
	"loss/ce": loss_dict.get("l_ce", 0).item(),
	"loss/style": loss_dict.get("l_style", 0).item(),
	"loss/semantic": loss_dict.get("l_semantic", 0).item(),
	"loss/total": loss_dict["total"].item(),
	"step": self.state.global_step,
	})

	return (loss_dict["total"], outputs) if return_outputs else loss_dict["total"]
	```

	---

	### `configs/training_config.yaml`

	```yaml
	model:
	key: "flan-t5-xl" # flan-t5-xl \| flan-t5-large \| bart-large \| llama-3.1-8b
	quantize: false # Set true for 4-bit on limited VRAM
	use_lora: true

	lora:
	r: 16
	lora_alpha: 32
	lora_dropout: 0.05
	target_modules: ["q", "v", "k", "o", "wi_0", "wi_1", "wo"]

	data:
	train_path: "data/processed/train.jsonl"
	val_path: "data/processed/val.jsonl"
	test_path: "data/processed/test.jsonl"
	max_input_length: 512
	max_target_length: 512
	augment_synthetic: true
	synthetic_ratio: 0.3

	training:
	output_dir: "checkpoints/"
	num_train_epochs: 5
	per_device_train_batch_size: 8
	per_device_eval_batch_size: 16
	gradient_accumulation_steps: 4 # Effective batch = 8*4 = 32
	learning_rate: 3.0e-4
	lr_scheduler_type: "cosine"
	warmup_ratio: 0.05
	weight_decay: 0.01
	fp16: false
	bf16: true # Use bfloat16 on Ampere+ GPUs
	evaluation_strategy: "steps"
	eval_steps: 500
	save_strategy: "steps"
	save_steps: 500
	save_total_limit: 3
	load_best_model_at_end: true
	metric_for_best_model: "gleu"
	greater_is_better: true
	logging_dir: "logs/"
	logging_steps: 50
	report_to: ["wandb", "tensorboard"]
	dataloader_num_workers: 4
	seed: 42
	push_to_hub: false

	loss:
	lambda_style: 0.3
	lambda_semantic: 0.5
	sem_model_name: "all-mpnet-base-v2"

	generation:
	num_beams: 5
	length_penalty: 1.0
	no_repeat_ngram_size: 3
	min_length: 10
	max_new_tokens: 512
	early_stopping: true
	```

	---

	## 10. Layer 6 — Academic Vocabulary Control Module

	### `src/vocabulary/lexical_substitution.py`

	```python
	"""
	Post-generation academic vocabulary elevation module.

	Pipeline:
	1. POS-tag the generated output
	2. Identify content words (NOUN, VERB, ADJ, ADV) NOT in AWL
	3. For each candidate word, generate AWL-aligned substitutions
	using BERT masked language model (fill-mask)
	4. Apply substitution only if:
	a. Semantic similarity between original and substitution > threshold
	b. Substitution is in the AWL
	c. Substitution does not change sentence meaning
	5. Apply register-level post-processing (nominalisation, hedging, passive)

	AWL = Coxhead Academic Word List (570 word families, ~3,000 lemmas)
	"""

	import spacy
	import torch
	from transformers import pipeline as hf_pipeline
	from sentence_transformers import SentenceTransformer
	import torch.nn.functional as F
	from typing import List, Dict, Tuple, Optional
	from .awl_loader import AWLLoader


	class LexicalElevator:

	# Words that should NEVER be substituted (structural, functional words)
	PROTECTED_POS = {"PRON", "DET", "CCONJ", "SCONJ", "ADP", "AUX", "PART", "PUNCT", "NUM"}
	SEMANTIC_THRESHOLD = 0.82 # Minimum cosine similarity to accept substitution

	def __init__(
	self,
	awl_path: str = "data/awl/coxhead_awl.txt",
	spacy_model: str = "en_core_web_trf",
	mlm_model: str = "bert-large-uncased",
	sem_model: str = "all-mpnet-base-v2",
	):
	self.nlp = spacy.load(spacy_model)
	self.awl = AWLLoader(awl_path)
	self.fill_mask = hf_pipeline("fill-mask", model=mlm_model, top_k=10)
	self.sem_model = SentenceTransformer(sem_model)

	def _sem_similarity(self, word_a: str, word_b: str, context: str) -> float:
	"""Compute contextual semantic similarity using sentence embeddings."""
	ctx_a = context.replace(word_a, word_a, 1)
	ctx_b = context.replace(word_a, word_b, 1)
	embs = self.sem_model.encode([ctx_a, ctx_b])
	t = torch.tensor(embs)
	return F.cosine_similarity(t[0].unsqueeze(0), t[1].unsqueeze(0)).item()

	def _get_awl_substitutions(self, sentence: str, word: str, pos: str) -> List[str]:
	"""Generate candidate substitutions using BERT fill-mask."""
	masked = sentence.replace(word, "[MASK]", 1)
	try:
	predictions = self.fill_mask(masked)
	candidates = [p["token_str"].strip() for p in predictions]
	except Exception:
	return []
	# Filter to AWL words only
	return [c for c in candidates if self.awl.is_academic(c)]

	def elevate(self, text: str, protected_spans: List[Tuple[int, int]] = None) -> str:
	"""
	Main entry point: elevates vocabulary to academic register.
	protected_spans: list of (start_char, end_char) that must not be modified.
	"""
	doc = self.nlp(text)
	replacements = {}

	for sent in doc.sents:
	sent_text = sent.text
	for token in sent:
	# Skip protected tokens
	if token.pos_ in self.PROTECTED_POS:
	continue
	if self.awl.is_academic(token.lemma_):
	continue # Already academic
	if protected_spans:
	if any(s <= token.idx < e for s, e in protected_spans):
	continue

	candidates = self._get_awl_substitutions(sent_text, token.text, token.pos_)
	for candidate in candidates:
	sim = self._sem_similarity(token.text, candidate, sent_text)
	if sim >= self.SEMANTIC_THRESHOLD:
	replacements[token.idx] = (token.text, candidate)
	break

	# Apply replacements (reverse order to preserve offsets)
	result = list(text)
	for idx in sorted(replacements.keys(), reverse=True):
	original, replacement = replacements[idx]
	start = idx
	end = idx + len(original)
	result[start:end] = list(replacement)

	return ''.join(result)


	class RegisterFilter:
	"""
	Applies register-level corrections to ensure academic tone:
	- Converts contractions to full forms
	- Ensures hedging where appropriate
	- Flags over-colloquial phrases for review
	"""

	CONTRACTIONS = {
	"don't": "do not", "can't": "cannot", "won't": "will not",
	"it's": "it is", "that's": "that is", "there's": "there is",
	"they're": "they are", "we're": "we are", "you're": "you are",
	"I'm": "I am", "I've": "I have", "I'll": "I will",
	"isn't": "is not", "aren't": "are not", "wasn't": "was not",
	"weren't": "were not", "hasn't": "has not", "haven't": "have not",
	"couldn't": "could not", "wouldn't": "would not", "shouldn't": "should not",
	}

	COLLOQUIAL_TO_ACADEMIC = {
	"a lot of": "a substantial number of",
	"lots of": "numerous",
	"big": "substantial",
	"get": "obtain",
	"show": "demonstrate",
	"use": "utilise",
	"find out": "ascertain",
	"look at": "examine",
	"think about": "consider",
	"talk about": "discuss",
	"deal with": "address",
	"carry out": "conduct",
	"point out": "indicate",
	"make sure": "ensure",
	"come up with": "develop",
	"go up": "increase",
	"go down": "decrease",
	"start": "commence",
	"end": "conclude",
	"help": "facilitate",
	"need": "require",
	"try": "attempt",
	"want": "seek",
	}

	def apply(self, text: str) -> str:
	import re
	for contraction, full in self.CONTRACTIONS.items():
	text = re.sub(re.escape(contraction), full, text, flags=re.IGNORECASE)
	for colloquial, academic in self.COLLOQUIAL_TO_ACADEMIC.items():
	text = re.sub(r'\b' + re.escape(colloquial) + r'\b', academic, text, flags=re.IGNORECASE)
	return text
	```

	---

	## 11. Layer 7 — Evaluation Framework

	### `src/evaluation/gleu_scorer.py`

	```python
	"""
	GLEU (Generalized Language Evaluation Understanding) score.
	Preferred over BLEU for grammatical error correction tasks.
	Designed specifically to handle the GEC task where the reference
	correction may differ from the source in minimal ways.

	Also computes BERTScore for semantic similarity evaluation.
	"""

	import sacrebleu
	from bert_score import score as bert_score_fn
	from typing import List, Tuple


	class GLEUScorer:

	def compute_gleu(
	self,
	predictions: List[str],
	references: List[str],
	) -> float:
	"""Corpus-level GLEU score."""
	result = sacrebleu.corpus_bleu(predictions, [references])
	return result.score # 0–100

	def compute_bert_score(
	self,
	predictions: List[str],
	references: List[str],
	lang: str = "en",
	) -> Tuple[float, float, float]:
	"""
	Returns (precision, recall, F1) as averages over the batch.
	F1 > 0.9 is generally considered high quality.
	"""
	P, R, F1 = bert_score_fn(predictions, references, lang=lang, verbose=False)
	return P.mean().item(), R.mean().item(), F1.mean().item()
	```

	---

	### `src/evaluation/style_metrics.py`

	```python
	"""
	Measures style preservation between input and output.

	Key metric: Style Vector Cosine Similarity
	sim = cosine_similarity(style_vec(input), style_vec(output))
	Target: > 0.85

	Key metric: Authorship Verification Score
	A binary classifier trained to answer: "Was this written by the same author?"
	Uses a fine-tuned RoBERTa model on authorship verification datasets.
	Target: > 0.80 (model says same author 80%+ of the time)

	Key metric: AWL Coverage Score
	Fraction of content words from the Academic Word List.
	Target: > 0.25 (25% of content words should be academic)
	"""

	import torch
	import torch.nn.functional as F
	from typing import List, Tuple
	from ..style.fingerprinter import StyleFingerprinter
	from .awl_loader import AWLLoader


	class StyleEvaluator:

	def __init__(self, fingerprinter: StyleFingerprinter, awl: AWLLoader):
	self.fingerprinter = fingerprinter
	self.awl = awl

	def style_similarity(self, text_a: str, text_b: str) -> float:
	"""Cosine similarity between style vectors. Target: > 0.85"""
	vec_a = self.fingerprinter.extract_vector(text_a)
	vec_b = self.fingerprinter.extract_vector(text_b)
	return F.cosine_similarity(vec_a.unsqueeze(0), vec_b.unsqueeze(0)).item()

	def awl_coverage(self, text: str) -> float:
	"""Fraction of content words in AWL. Target: > 0.25"""
	import spacy
	nlp = spacy.load("en_core_web_sm")
	doc = nlp(text)
	content_words = [t.lemma_.lower() for t in doc if t.pos_ in {"NOUN", "VERB", "ADJ", "ADV"}]
	if not content_words:
	return 0.0
	return sum(1 for w in content_words if self.awl.is_academic(w)) / len(content_words)

	def evaluate_batch(
	self,
	inputs: List[str],
	outputs: List[str],
	references: List[str],
	) -> dict:
	style_sims = [self.style_similarity(i, o) for i, o in zip(inputs, outputs)]
	awl_scores = [self.awl_coverage(o) for o in outputs]
	return {
	"style_similarity_mean": sum(style_sims) / len(style_sims),
	"style_similarity_min": min(style_sims),
	"awl_coverage_mean": sum(awl_scores) / len(awl_scores),
	}
	```

	---

	## 12. Layer 8 — Inference Pipeline

	### `src/inference/corrector.py`

	```python
	"""
	End-to-end inference pipeline.
	Accepts raw dyslectic text (and optionally a master copy),
	returns corrected academic text with metadata.
	"""

	from ..preprocessing.pipeline import PreprocessingPipeline
	from ..style.fingerprinter import StyleFingerprinter
	from ..vocabulary.lexical_substitution import LexicalElevator, RegisterFilter
	from ..model.base_model import load_model_and_tokenizer
	from ..model.style_conditioner import StyleConditioner, prepend_style_prefix
	import torch
	from typing import Optional
	from dataclasses import dataclass


	TASK_PREFIX = (
	"Correct the following text for grammar, spelling, and clarity. "
	"Maintain the author's original tone and writing style. "
	"Elevate vocabulary to academic register. "
	"Do NOT change the meaning or add new information. "
	"Preserve named entities exactly. "
	"Text to correct: "
	)


	@dataclass
	class CorrectionResult:
	original: str
	corrected: str
	preprocessed: str
	style_similarity: float
	awl_coverage: float
	readability: dict
	changes_summary: str


	class AcademicCorrector:

	def __init__(self, config: dict):
	self.device = "cuda" if torch.cuda.is_available() else "cpu"
	self.model, self.tokenizer, self.is_seq2seq = load_model_and_tokenizer(
	config["model"]["key"],
	quantize=config["model"].get("quantize", False),
	use_lora=False, # Inference: use merged weights
	)
	self.model.eval()
	self.preprocessor = PreprocessingPipeline()
	self.fingerprinter = StyleFingerprinter()
	self.conditioner = StyleConditioner(
	style_dim=512,
	model_hidden_dim=config.get("model_hidden_dim", 2048),
	n_prefix_tokens=10,
	).to(self.device)
	self.elevator = LexicalElevator()
	self.register_filter = RegisterFilter()
	self.gen_config = config.get("generation", {})

	def correct(
	self,
	raw_text: str,
	master_copy: Optional[str] = None,
	style_alpha: float = 0.6,
	) -> CorrectionResult:

	# Step 1: Pre-process
	doc = self.preprocessor.process(raw_text)

	# Step 2: Style fingerprinting
	user_style = self.fingerprinter.extract_vector(doc.corrected_text)
	master_style = self.fingerprinter.extract_vector(master_copy) if master_copy else None
	target_style = self.fingerprinter.blend_vectors(user_style, master_style, alpha=style_alpha)

	# Step 3: Prepare input
	input_text = TASK_PREFIX + doc.corrected_text
	inputs = self.tokenizer(
	input_text,
	return_tensors="pt",
	max_length=512,
	truncation=True,
	).to(self.device)

	# Step 4: Generate with style conditioning
	style_prefix = self.conditioner(target_style.unsqueeze(0).to(self.device))

	with torch.no_grad():
	generated = self.model.generate(
	**inputs,
	num_beams=self.gen_config.get("num_beams", 5),
	length_penalty=self.gen_config.get("length_penalty", 1.0),
	no_repeat_ngram_size=self.gen_config.get("no_repeat_ngram_size", 3),
	max_new_tokens=self.gen_config.get("max_new_tokens", 512),
	early_stopping=True,
	)

	draft = self.tokenizer.decode(generated[0], skip_special_tokens=True)

	# Step 5: Academic vocabulary elevation
	elevated = self.elevator.elevate(draft, protected_spans=doc.protected_spans)

	# Step 6: Register filter (contractions, colloquialisms)
	final = self.register_filter.apply(elevated)

	# Step 7: Compute metrics
	from ..evaluation.style_metrics import StyleEvaluator
	from ..vocabulary.awl_loader import AWLLoader
	evaluator = StyleEvaluator(self.fingerprinter, AWLLoader())
	style_sim = evaluator.style_similarity(raw_text, final)
	awl_cov = evaluator.awl_coverage(final)

	return CorrectionResult(
	original=raw_text,
	corrected=final,
	preprocessed=doc.corrected_text,
	style_similarity=round(style_sim, 3),
	awl_coverage=round(awl_cov, 3),
	readability=doc.readability,
	changes_summary=f"Style similarity: {style_sim:.2%} \| AWL coverage: {awl_cov:.2%}",
	)
	```

	---

	## 13. Layer 9 — API Server

	### `src/api/main.py`

	```python
	from fastapi import FastAPI, HTTPException
	from fastapi.middleware.cors import CORSMiddleware
	from .schemas import CorrectionRequest, CorrectionResponse
	from ..inference.corrector import AcademicCorrector
	import yaml

	app = FastAPI(
	title="Dyslexia Academic Writing Corrector API",
	description="Style-preserving grammar correction and academic vocabulary elevation.",
	version="1.0.0",
	)

	app.add_middleware(
	CORSMiddleware,
	allow_origins=["*"],
	allow_methods=["*"],
	allow_headers=["*"],
	)

	with open("configs/inference_config.yaml") as f:
	config = yaml.safe_load(f)

	corrector = AcademicCorrector(config)


	@app.post("/correct", response_model=CorrectionResponse)
	async def correct_text(request: CorrectionRequest):
	try:
	result = corrector.correct(
	raw_text=request.text,
	master_copy=request.master_copy,
	style_alpha=request.style_alpha,
	)
	return CorrectionResponse(
	original=result.original,
	corrected=result.corrected,
	style_similarity=result.style_similarity,
	awl_coverage=result.awl_coverage,
	readability=result.readability,
	changes_summary=result.changes_summary,
	)
	except Exception as e:
	raise HTTPException(status_code=500, detail=str(e))


	@app.get("/health")
	async def health():
	return {"status": "ok", "model": config["model"]["key"]}
	```

	---

	### `src/api/schemas.py`

	```python
	from pydantic import BaseModel, Field
	from typing import Optional, Dict


	class CorrectionRequest(BaseModel):
	text: str = Field(..., min_length=10, max_length=5000, description="Raw dyslectic text to correct.")
	master_copy: Optional[str] = Field(None, description="Optional master copy to match style toward.")
	style_alpha: float = Field(0.6, ge=0.0, le=1.0, description="Weight given to user's own style (0=full master, 1=full user).")


	class CorrectionResponse(BaseModel):
	original: str
	corrected: str
	style_similarity: float
	awl_coverage: float
	readability: Dict[str, float]
	changes_summary: str
	```

	---

	## 14. Layer 10 — Configuration Files

	### `configs/model_config.yaml`

	```yaml
	model:
	key: "flan-t5-xl"
	checkpoint_path: "checkpoints/best_model"
	quantize: false
	use_lora: true
	model_hidden_dim: 2048 # flan-t5-xl hidden size
	# model_hidden_dim: 1024 # flan-t5-large
	# model_hidden_dim: 1024 # bart-large
	# model_hidden_dim: 4096 # llama-3.1-8b

	style_conditioner:
	style_dim: 512
	n_prefix_tokens: 10

	fingerprinter:
	spacy_model: "en_core_web_trf"
	awl_path: "data/awl/coxhead_awl.txt"
	projection_hidden_dim: 256
	projection_output_dim: 512

	generation:
	num_beams: 5
	length_penalty: 1.0
	no_repeat_ngram_size: 3
	min_length: 10
	max_new_tokens: 512
	early_stopping: true
	temperature: 0.7 # Slight randomness for naturalness
	do_sample: false # Beam search by default

	vocabulary:
	awl_path: "data/awl/coxhead_awl.txt"
	mlm_model: "bert-large-uncased"
	sem_model: "all-mpnet-base-v2"
	semantic_threshold: 0.82
	```

	---

	### `configs/awl_config.yaml`

	```yaml
	awl:
	primary: "data/awl/coxhead_awl.txt"
	supplementary:
	- "data/awl/domain_lexicons/humanities.txt"
	- "data/awl/domain_lexicons/sciences.txt"
	- "data/awl/domain_lexicons/social_sciences.txt"
	academic_synonyms: "data/awl/academic_synonyms.json"

	register:
	expand_contractions: true
	replace_colloquialisms: true
	enforce_third_person_academic: false # Keep user's voice (don't force "one")
	minimum_formality_score: 0.65
	```

	---

	## 15. Layer 11 — Full Training Run Sequence

	### `scripts/train.py`

	```python
	"""
	Full training entry point.
	Run: python scripts/train.py --config configs/training_config.yaml
	"""

	import click
	import yaml
	import wandb
	import torch
	from transformers import TrainingArguments
	from torch.utils.data import random_split

	from src.model.base_model import load_model_and_tokenizer
	from src.model.style_conditioner import StyleConditioner
	from src.training.dataset import WritingCorrectionDataset
	from src.training.loss_functions import CombinedCorrectionLoss
	from src.training.trainer import CorrectionTrainer
	from src.style.fingerprinter import StyleFingerprinter
	from src.evaluation.gleu_scorer import GLEUScorer


	@click.command()
	@click.option("--config", default="configs/training_config.yaml")
	def train(config: str):
	with open(config) as f:
	cfg = yaml.safe_load(f)

	wandb.init(project="dyslexia-writing-ai", config=cfg)

	# Load model
	model, tokenizer, is_seq2seq = load_model_and_tokenizer(
	cfg["model"]["key"],
	quantize=cfg["model"].get("quantize", False),
	use_lora=cfg["model"].get("use_lora", True),
	)

	# Fingerprinter
	fingerprinter = StyleFingerprinter(
	awl_path=cfg["data"].get("awl_path", "data/awl/coxhead_awl.txt")
	)

	# Datasets
	train_dataset = WritingCorrectionDataset(
	data_path=cfg["data"]["train_path"],
	tokenizer=tokenizer,
	fingerprinter=fingerprinter,
	max_input_length=cfg["data"]["max_input_length"],
	max_target_length=cfg["data"]["max_target_length"],
	augment_with_synthetic=cfg["data"]["augment_synthetic"],
	synthetic_ratio=cfg["data"]["synthetic_ratio"],
	)
	val_dataset = WritingCorrectionDataset(
	data_path=cfg["data"]["val_path"],
	tokenizer=tokenizer,
	fingerprinter=fingerprinter,
	augment_with_synthetic=False,
	)

	# Loss function
	loss_fn = CombinedCorrectionLoss(
	lambda_style=cfg["loss"]["lambda_style"],
	lambda_semantic=cfg["loss"]["lambda_semantic"],
	sem_model_name=cfg["loss"]["sem_model_name"],
	device="cuda" if torch.cuda.is_available() else "cpu",
	)

	# Training arguments
	training_args = TrainingArguments(
	output_dir=cfg["training"]["output_dir"],
	num_train_epochs=cfg["training"]["num_train_epochs"],
	per_device_train_batch_size=cfg["training"]["per_device_train_batch_size"],
	per_device_eval_batch_size=cfg["training"]["per_device_eval_batch_size"],
	gradient_accumulation_steps=cfg["training"]["gradient_accumulation_steps"],
	learning_rate=cfg["training"]["learning_rate"],
	lr_scheduler_type=cfg["training"]["lr_scheduler_type"],
	warmup_ratio=cfg["training"]["warmup_ratio"],
	weight_decay=cfg["training"]["weight_decay"],
	bf16=cfg["training"]["bf16"],
	fp16=cfg["training"]["fp16"],
	evaluation_strategy=cfg["training"]["evaluation_strategy"],
	eval_steps=cfg["training"]["eval_steps"],
	save_strategy=cfg["training"]["save_strategy"],
	save_steps=cfg["training"]["save_steps"],
	save_total_limit=cfg["training"]["save_total_limit"],
	load_best_model_at_end=cfg["training"]["load_best_model_at_end"],
	logging_steps=cfg["training"]["logging_steps"],
	report_to=cfg["training"]["report_to"],
	dataloader_num_workers=cfg["training"]["dataloader_num_workers"],
	seed=cfg["training"]["seed"],
	)

	trainer = CorrectionTrainer(
	model=model,
	args=training_args,
	train_dataset=train_dataset,
	eval_dataset=val_dataset,
	loss_fn=loss_fn,
	fingerprinter=fingerprinter,
	tokenizer=tokenizer,
	)

	trainer.train()
	trainer.save_model(cfg["training"]["output_dir"] + "/final")
	wandb.finish()


	if __name__ == "__main__":
	train()
	```

	---

	### `scripts/download_datasets.sh`

	```bash
	#!/bin/bash
	# Download all training data sources

	mkdir -p data/raw/wi_locness data/raw/jfleg data/raw/gyafc data/raw/custom_dyslexia

	# W&I+LOCNESS (Cambridge Grammar Error Correction)
	# Requires registration at: https://www.cl.cam.ac.uk/research/nl/bea2019st/
	echo "W&I+LOCNESS: Download manually from https://www.cl.cam.ac.uk/research/nl/bea2019st/"
	echo "Place files in data/raw/wi_locness/"

	# JFLEG (JHU Fluency-Extended GUG)
	git clone https://github.com/keisks/jfleg.git data/raw/jfleg_repo
	cp data/raw/jfleg_repo/test/*.src data/raw/jfleg/
	cp data/raw/jfleg_repo/test/.ref data/raw/jfleg/

	# GYAFC (Formality Corpus - Yahoo Answers)
	# Requires request from Grammarly: https://github.com/raosudha89/GYAFC-corpus
	echo "GYAFC: Request access at https://github.com/raosudha89/GYAFC-corpus"
	echo "Place files in data/raw/gyafc/"

	# Download Coxhead Academic Word List
	curl -o data/awl/coxhead_awl.txt \
	"https://www.victoria.ac.nz/lals/resources/academicwordlist/sublists/Sublist_1.txt"

	echo "Dataset download complete. Check manually downloaded datasets."
	```

	### `scripts/preprocess_data.py`

	```python
	"""
	Converts all raw dataset formats into unified JSONL training format.
	Output schema per line:
	{"input": "...", "target": "...", "source": "wi_locness\|jfleg\|gyafc\|synthetic"}
	"""
	import json
	import os
	from pathlib import Path
	from src.preprocessing.dyslexia_simulator import DyslexiaSimulator


	def process_jfleg(raw_dir: str, out_file):
	"""JFLEG: .src files (original) and .ref0..ref3 (4 human corrections)."""
	src_files = list(Path(raw_dir).glob("*.src"))
	for src_file in src_files:
	refs = [src_file.with_suffix(f".ref{i}") for i in range(4)]
	with open(src_file) as sf:
	src_lines = sf.readlines()
	for ref_path in refs:
	if ref_path.exists():
	with open(ref_path) as rf:
	ref_lines = rf.readlines()
	for src, ref in zip(src_lines, ref_lines):
	src, ref = src.strip(), ref.strip()
	if src and ref and src != ref:
	out_file.write(json.dumps({"input": src, "target": ref, "source": "jfleg"}) + "\n")


	def process_gyafc(raw_dir: str, out_file):
	"""GYAFC: informal/ and formal/ subdirectories with parallel files."""
	for domain in ["Entertainment_Music", "Family_Relationships"]:
	for split in ["train", "tune", "test"]:
	informal = Path(raw_dir) / domain / split / "informal"
	formal = Path(raw_dir) / domain / split / "formal.ref0"
	if informal.exists() and formal.exists():
	with open(informal) as inf_f, open(formal) as form_f:
	for inf_line, form_line in zip(inf_f, form_f):
	inf_line, form_line = inf_line.strip(), form_line.strip()
	if inf_line and form_line:
	out_file.write(json.dumps({"input": inf_line, "target": form_line, "source": "gyafc"}) + "\n")


	def main():
	os.makedirs("data/processed", exist_ok=True)
	with open("data/processed/train.jsonl", "w") as out:
	process_jfleg("data/raw/jfleg", out)
	process_gyafc("data/raw/gyafc", out)
	# Add W&I+LOCNESS processing here when available
	print("Preprocessing complete.")


	if __name__ == "__main__":
	main()
	```

	---

	## 16. Mathematical Formulations

	### Total Training Loss

	```
	L_total = L_CE + λ₁ · L_style + λ₂ · L_semantic

	L_CE = -Σ log P(y_t \| y_{<t}, x) [cross-entropy over tokens]

	L_style = 1 - cos(f(ŷ), f(y)) [style vector cosine distance]
	where f(·) = StyleFingerprinter.extract_vector(·)
	cos(a,b) = (a·b) / (‖a‖‖b‖)

	L_semantic = 1 - cos(g(x), g(ŷ)) [semantic distance from input]
	where g(·) = SentenceTransformer.encode(·)
	(frozen during training)

	λ₁ = 0.3 (style weight)
	λ₂ = 0.5 (semantic weight)
	```

	### Style Vector Blending

	```
	target_style = α · v_user + (1-α) · v_master
	target_style = target_style / ‖target_style‖ [L2 normalise to unit sphere]

	α = 0.6 (user's style dominates by default)
	```

	### LoRA Weight Update

	```
	W = W₀ + ΔW
	ΔW = B · A

	Where:
	W₀ ∈ R^{d×k} (frozen pretrained weights)
	A ∈ R^{r×k} (trainable, r << d, initialised from N(0, σ²))
	B ∈ R^{d×r} (trainable, initialised to zero)
	r = 16 (rank hyperparameter)

	Effective weight: W_eff = W₀ + (α/r) · B·A
	where α = lora_alpha = 32
	```

	### Style Similarity Evaluation Metric

	```
	StyleSim(input, output) = cos(f(input), f(output))
	= (f(input) · f(output)) / (‖f(input)‖ · ‖f(output)‖)

	Target: StyleSim > 0.85
	Acceptable minimum: StyleSim > 0.75
	```

	### AWL Coverage Score

	```
	AWL_Coverage(text) = \|{w ∈ content_words(text) : lemma(w) ∈ AWL}\|
	─────────────────────────────────────────────
	\|content_words(text)\|

	content_words = {w : POS(w) ∈ {NOUN, VERB, ADJ, ADV}}
	Target: AWL_Coverage > 0.25
	```

	---

	## 17. Hyperparameter Reference

	\| Hyperparameter \| Value \| Rationale \|
	\|---\|---\|---\|
	\| LoRA rank (r) \| 16 \| Balances capacity vs. parameter efficiency \|
	\| LoRA alpha \| 32 \| Standard 2x rank scaling \|
	\| LoRA dropout \| 0.05 \| Light regularisation \|
	\| Learning rate \| 3e-4 \| Standard for LoRA fine-tuning \|
	\| LR scheduler \| cosine \| Smooth decay, avoids sharp LR drops \|
	\| Warmup ratio \| 0.05 \| 5% of steps for warmup \|
	\| Batch size (device) \| 8 \| Per GPU \|
	\| Gradient accumulation \| 4 \| Effective batch = 32 \|
	\| Training epochs \| 5 \| Sufficient for fine-tuning on GEC data \|
	\| λ₁ (style weight) \| 0.3 \| Strong style signal without dominating CE \|
	\| λ₂ (semantic weight) \| 0.5 \| Meaning preservation is critical \|
	\| Style blend α \| 0.6 \| User style dominates over master copy \|
	\| Style prefix tokens \| 10 \| Virtual prefix length \|
	\| Beam search beams \| 5 \| Quality vs. speed tradeoff \|
	\| No-repeat ngram \| 3 \| Prevents repetition in output \|
	\| Semantic threshold \| 0.82 \| For lexical substitution acceptance \|
	\| Max input tokens \| 512 \| T5/BART context window \|
	\| Style projection dim \| 512 \| Rich enough to capture style nuance \|

	---

	## 18. Dataset Sources & Download Instructions

	\| Dataset \| Size \| Task \| Access \| URL \|
	\|---\|---\|---\|---\|---\|
	\| W&I+LOCNESS \| ~35k pairs \| Grammar error correction \| Free registration \| https://www.cl.cam.ac.uk/research/nl/bea2019st/ \|
	\| JFLEG \| ~1.5k pairs \| Fluency correction \| Public GitHub \| https://github.com/keisks/jfleg \|
	\| GYAFC \| ~105k pairs \| Formality transfer \| Request from Grammarly \| https://github.com/raosudha89/GYAFC-corpus \|
	\| CoNLL-2014 \| ~1.3k pairs \| Grammar correction \| Public \| https://www.comp.nus.edu.sg/~nlp/conll14st.html \|
	\| FCE Corpus \| ~33k pairs \| Learner English \| Free registration \| https://ilexir.co.uk/datasets/index.html \|
	\| WikiAtomic \| Millions \| Style transfer \| Public \| https://huggingface.co/datasets/wiki_atomic_edits \|
	\| Synthetic (generated) \| Unlimited \| Dyslexia simulation \| Self-generated \| `scripts/preprocess_data.py` \|

	---

	## 19. Hardware Requirements

	### Minimum (Development / Testing)

	```
	CPU: 8-core, e.g., Intel i7 / AMD Ryzen 7
	RAM: 32 GB
	GPU: NVIDIA RTX 3090 (24 GB VRAM) ← Fine-tune Flan-T5-Large or BART-large
	SSD: 500 GB NVMe
	Model: Flan-T5-Large (780M) or BART-large (400M)
	Quantize: false
	```

	### Recommended (Production Training)

	```
	CPU: 16-core+
	RAM: 64 GB
	GPU: NVIDIA A100 80 GB OR 2× RTX 4090 (48 GB total)
	SSD: 2 TB NVMe
	Model: Flan-T5-XL (3B) with LoRA
	Quantize: false
	Training time: ~12 hours on A100
	```

	### Maximum Quality

	```
	GPU: 4× A100 80 GB (320 GB total VRAM)
	Model: Llama-3.1-8B with LoRA
	Training time: ~24-48 hours
	Use: torchrun --nproc_per_node=4 scripts/train.py
	```

	### Cloud Options

	```
	AWS: p3.2xlarge (V100 16GB) → BART-large only
	p3.8xlarge (4× V100 64GB) → Flan-T5-XL
	p4d.24xlarge (8× A100) → Llama-3.1-8B

	GCP: n1-standard-8 + 1× A100 → Flan-T5-XL
	a2-highgpu-4g (4× A100) → Llama-3.1-8B

	Lambda Labs: 1× A100 ~$1.10/hr → Most cost-effective
	RunPod: 1× A100 ~$0.99/hr → Alternative
	```

	---

	## 20. Testing Suite

	### `tests/test_preprocessing.py`

	```python
	import pytest
	from src.preprocessing.pipeline import PreprocessingPipeline
	from src.preprocessing.dyslexia_simulator import DyslexiaSimulator


	@pytest.fixture
	def pipeline():
	return PreprocessingPipeline()


	def test_spell_correction(pipeline):
	result = pipeline.process("i wuz going to the store but cud not find it")
	assert "was" in result.corrected_text
	assert "could" in result.corrected_text


	def test_entity_protection(pipeline):
	result = pipeline.process("John Smith livd in London.")
	entities = [e.text for e in result.entities]
	assert any("John" in e or "London" in e for e in entities)


	def test_sentence_segmentation(pipeline):
	result = pipeline.process("I went to school. I lerned a lot.")
	assert len(result.sentences) == 2


	def test_dyslexia_simulator():
	sim = DyslexiaSimulator(error_rate=1.0, seed=0)
	corrupted, clean = sim.simulate("The quick brown fox jumps over the lazy dog.")
	assert corrupted != clean
	assert clean == "The quick brown fox jumps over the lazy dog."
	```

	### `tests/test_style.py`

	```python
	import pytest
	import torch
	from src.style.fingerprinter import StyleFingerprinter


	@pytest.fixture
	def fingerprinter(tmp_path):
	awl = tmp_path / "awl.txt"
	awl.write_text("analysis\nconsider\nestablish\nsignificant\n")
	return StyleFingerprinter(spacy_model="en_core_web_sm", awl_path=str(awl))


	def test_style_vector_shape(fingerprinter):
	vec = fingerprinter.extract_vector("The quick brown fox jumps over the lazy dog.")
	assert vec.shape == (512,)


	def test_style_vector_different_texts(fingerprinter):
	formal = "The analysis demonstrates significant implications for the field."
	informal = "So basically it shows that this stuff really matters a lot lol."
	vec_formal = fingerprinter.extract_vector(formal)
	vec_informal = fingerprinter.extract_vector(informal)
	# Vectors should be different
	assert not torch.allclose(vec_formal, vec_informal)


	def test_style_blend(fingerprinter):
	vec_a = fingerprinter.extract_vector("Short punchy text here.")
	vec_b = fingerprinter.extract_vector("Elaborate and comprehensive academic discourse.")
	blended = fingerprinter.blend_vectors(vec_a, vec_b, alpha=0.5)
	assert blended.shape == (512,)
	# Blended should be unit vector
	assert abs(blended.norm().item() - 1.0) < 1e-4
	```

	---

	## Quick Start Execution Order

	```bash
	# 1. Setup environment
	python -m venv venv && source venv/bin/activate
	pip install -r requirements.txt
	python -m spacy download en_core_web_trf

	# 2. Download datasets
	bash scripts/download_datasets.sh

	# 3. Preprocess all data into unified format
	python scripts/preprocess_data.py

	# 4. Run tests to verify setup
	pytest tests/ -v

	# 5. Launch training
	python scripts/train.py --config configs/training_config.yaml

	# 6. Evaluate on test set
	python scripts/evaluate.py --config configs/training_config.yaml --split test

	# 7. Start inference API
	uvicorn src.api.main:app --host 0.0.0.0 --port 8000 --reload

	# 8. Test the API
	curl -X POST http://localhost:8000/correct \
	-H "Content-Type: application/json" \
	-d '{"text": "i went to the store but cud not find wat i was loking for", "style_alpha": 0.6}'
	```

	---

	---

	## 21. Human-Pattern Anti-AI Training Layer

	### The Core Principle

	These two Kaggle datasets are not used to build a detector. They are used to teach the model the statistical and linguistic signature of human writing, and to penalise the model when its output drifts toward AI-typical patterns. The training signal flows in one direction: reward human-like writing, penalise AI-like writing.

	This is implemented as an additional loss term — `L_human_pattern` — added to the combined loss from Layer 5. The model learns what human writing looks and feels like at a statistical level, and is penalised during training whenever its generated corrections exhibit the same surface patterns that distinguish AI-generated text from human text in these datasets.

	---

	### Dataset 1 — shanegerami/ai-vs-human-text

	```
	Source: https://www.kaggle.com/datasets/shanegerami/ai-vs-human-text
	Size: ~500,000 essays
	Format: CSV — two columns
	Columns:
	text (str) Full essay text
	generated (int) 0 = human-written \| 1 = AI-generated

	Human count: 305,797 essays
	AI count: ~194,203 essays (GPT-family generated)
	Content type: Academic essays across diverse topics
	File: train_essays.csv

	HuggingFace mirror (already split, ~400k rows, use this for convenience):
	andythetechnerd03/AI-human-text
	Load: datasets.load_dataset("andythetechnerd03/AI-human-text")
	```

	### Dataset 2 — starblasters8/human-vs-llm-text-corpus

	```
	Source: https://www.kaggle.com/datasets/starblasters8/human-vs-llm-text-corpus
	Size: ~800,000 texts
	Format: Parquet — data.parquet
	Columns:
	text (str) Full text
	label (str) "Human" \| <LLM model name> (63 different LLMs represented)

	Key feature: covers 63 DIFFERENT LLMs — not just GPT. Includes outputs from
	Llama, Mistral, Falcon, Claude, Gemini, PaLM, Vicuna, Alpaca, and many others.
	This is critical: the model learns what AI text looks like ACROSS the LLM landscape,
	not just from one model family.

	File: data.parquet
	Read: pd.read_parquet("data/raw/starblasters8/data.parquet")
	```

	---

	### `scripts/download_kaggle_datasets.sh`

	```bash
	#!/bin/bash
	# Requires: pip install kaggle
	# Setup: Place kaggle.json API key at ~/.kaggle/kaggle.json
	# Get key: kaggle.com → Account → Create New API Token

	mkdir -p data/raw/shanegerami data/raw/starblasters8

	# Dataset 1: AI vs Human Text (500K essays)
	kaggle datasets download -d shanegerami/ai-vs-human-text \
	-p data/raw/shanegerami --unzip

	# Dataset 2: Human vs LLM Text Corpus (800K, 63 LLMs)
	kaggle datasets download -d starblasters8/human-vs-llm-text-corpus \
	-p data/raw/starblasters8 --unzip

	echo "Both datasets downloaded."
	echo "Dataset 1 (CSV): data/raw/shanegerami/train_essays.csv"
	echo "Dataset 2 (Parquet): data/raw/starblasters8/data.parquet"
	```

	---

	### `src/training/human_pattern_extractor.py`

	```python
	"""
	Extracts the statistical signature of human writing vs AI writing.
	Uses the two Kaggle datasets to build:

	1. HumanPatternProfile — a statistical distribution of human writing features
	2. AIPatternProfile — a statistical distribution of AI writing features
	3. HumanPatternClassifier — a lightweight FROZEN classifier used at training time
	to score how "human-like" the model's output looks.

	The classifier is FROZEN during main model training. It is pre-trained separately
	on the Kaggle datasets, then its output score is used as a reward/penalty signal
	in the main training loss.

	Feature set extracted (same dimensions as StyleFingerprinter + additional):
	- All 40 StyleFingerprinter features
	- Perplexity under GPT-2 (AI text tends to be lower perplexity)
	- Burstiness score (human writing has more sentence length variance)
	- Lexical diversity (AI text has narrower vocab distributions)
	- Punctuation density patterns (AI overuses certain patterns)
	- Discourse marker overuse (AI overuses "Furthermore", "Moreover", "Additionally")
	- Sentence starter diversity (AI repeats sentence openers more)
	- n-gram novelty score (AI repeats common n-grams more)
	- Hedging vs certainty ratio (AI is overconfident OR over-hedges in detectable ways)
	- Paragraph cohesion score (AI has unnaturally perfect paragraph transitions)
	"""

	import pandas as pd
	import numpy as np
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from torch.utils.data import Dataset, DataLoader
	from transformers import GPT2LMHeadModel, GPT2TokenizerFast
	from sklearn.model_selection import train_test_split
	from sklearn.preprocessing import StandardScaler
	from typing import List, Tuple, Dict
	import spacy
	from collections import Counter
	import math


	# ── AI-Typical Overused Discourse Markers ───────────────────────────────────
	AI_OVERUSED_MARKERS = {
	"furthermore", "moreover", "additionally", "consequently",
	"in conclusion", "to summarize", "it is worth noting",
	"it is important to note", "in today's world", "in today's society",
	"in the modern era", "as previously mentioned", "needless to say",
	"it goes without saying", "at the end of the day",
	"in terms of", "with regard to", "with respect to",
	"delve", "leverage", "utilize", "holistic", "paradigm",
	"transformative", "groundbreaking", "revolutionary", "game-changing",
	"multifaceted", "nuanced", "comprehensive", "robust", "seamless",
	"innovative", "synergy", "cutting-edge", "state-of-the-art",
	}

	# Words that AI uses far MORE than humans in academic-adjacent writing
	AI_FINGERPRINT_WORDS = {
	"delve", "underscore", "tapestry", "intricate", "pivotal",
	"crucial", "vital", "essential", "significant", "notable",
	"commendable", "noteworthy", "straightforward", "straightforwardly",
	"elucidate", "expound", "illuminate", "unravel", "harness",
	"foster", "facilitate", "leverage", "optimize", "streamline",
	}


	class HumanPatternFeatureExtractor:
	"""Extracts 55-dimensional feature vector encoding human vs AI writing patterns."""

	def __init__(self, spacy_model: str = "en_core_web_sm"):
	self.nlp = spacy.load(spacy_model)
	self.gpt2_model = GPT2LMHeadModel.from_pretrained("gpt2")
	self.gpt2_tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
	self.gpt2_model.eval()

	def _perplexity(self, text: str, max_len: int = 512) -> float:
	"""
	AI text tends to have LOWER perplexity under GPT-2
	because LLMs generate high-probability token sequences.
	Human text is more unpredictable → higher perplexity.

	Lower perplexity = more likely to be AI.
	Higher perplexity = more likely to be human.
	"""
	encodings = self.gpt2_tokenizer(text, return_tensors="pt", truncation=True, max_length=max_len)
	with torch.no_grad():
	outputs = self.gpt2_model(**encodings, labels=encodings["input_ids"])
	return math.exp(outputs.loss.item())

	def _burstiness(self, sentences: List[str]) -> float:
	"""
	Burstiness = coefficient of variation of sentence lengths.
	Human writing has high burstiness (unpredictable length variation).
	AI writing has low burstiness (unnaturally uniform sentence lengths).

	B = std(lengths) / mean(lengths)
	"""
	lengths = [len(s.split()) for s in sentences]
	if len(lengths) < 2 or np.mean(lengths) == 0:
	return 0.0
	return np.std(lengths) / np.mean(lengths)

	def _sentence_starter_diversity(self, sentences: List[str]) -> float:
	"""
	Fraction of unique first words across sentences.
	AI tends to start sentences with the same words repeatedly.
	High = human-like. Low = AI-like.
	"""
	starters = [s.split()[0].lower() for s in sentences if s.split()]
	if not starters:
	return 0.0
	return len(set(starters)) / len(starters)

	def _ngram_novelty(self, text: str, n: int = 3) -> float:
	"""
	Ratio of unique n-grams to total n-grams.
	AI repeats common n-grams more than humans.
	Higher = more novel = more human-like.
	"""
	words = text.lower().split()
	if len(words) < n:
	return 1.0
	ngrams = [tuple(words[i:i+n]) for i in range(len(words) - n + 1)]
	return len(set(ngrams)) / len(ngrams)

	def _ai_marker_density(self, text: str, word_count: int) -> float:
	"""
	Frequency of AI-fingerprint words per 100 words.
	Higher = more AI-like.
	"""
	lower = text.lower()
	hits = sum(1 for marker in AI_FINGERPRINT_WORDS if marker in lower)
	return (hits / max(word_count, 1)) * 100

	def _overused_discourse_density(self, text: str, word_count: int) -> float:
	"""
	Frequency of AI-overused discourse markers per 100 words.
	"""
	lower = text.lower()
	hits = sum(1 for marker in AI_OVERUSED_MARKERS if marker in lower)
	return (hits / max(word_count, 1)) * 100

	def _punctuation_pattern(self, text: str, word_count: int) -> Dict[str, float]:
	"""
	AI writing exhibits characteristic punctuation patterns:
	- Overuse of em-dash (—)
	- Underuse of ellipsis (...)
	- Very consistent comma density
	"""
	em_dash_rate = text.count("—") / max(word_count, 1) * 100
	ellipsis_rate = text.count("...") / max(word_count, 1) * 100
	comma_rate = text.count(",") / max(word_count, 1) * 100
	semicolon_rate = text.count(";") / max(word_count, 1) * 100
	return {
	"em_dash_rate": em_dash_rate,
	"ellipsis_rate": ellipsis_rate,
	"comma_rate": comma_rate,
	"semicolon_rate": semicolon_rate,
	}

	def extract(self, text: str) -> np.ndarray:
	"""Extract full 55-dimensional feature vector."""
	doc = self.nlp(text[:10000]) # Truncate for speed
	sentences = [s.text.strip() for s in doc.sents if s.text.strip()]
	words = [t.text for t in doc if not t.is_punct and not t.is_space]
	word_count = len(words)

	punct = self._punctuation_pattern(text, word_count)

	features = np.array([
	# Human-pattern features
	self._perplexity(text[:1024]), # Higher = more human
	self._burstiness(sentences), # Higher = more human
	self._sentence_starter_diversity(sentences), # Higher = more human
	self._ngram_novelty(text, n=2), # Higher = more human
	self._ngram_novelty(text, n=3), # Higher = more human
	self._ngram_novelty(text, n=4), # Higher = more human

	# AI-pattern features (higher = more AI)
	self._ai_marker_density(text, word_count),
	self._overused_discourse_density(text, word_count),
	punct["em_dash_rate"],
	punct["ellipsis_rate"],
	punct["comma_rate"],
	punct["semicolon_rate"],

	# Distributional features
	float(word_count),
	float(len(sentences)),
	np.mean([len(s.split()) for s in sentences]) if sentences else 0,
	np.std([len(s.split()) for s in sentences]) if sentences else 0,
	len(set(w.lower() for w in words)) / max(word_count, 1), # TTR
	], dtype=np.float32)

	return features # [17 raw features — extend as needed]


	class KaggleHumanPatternDataset(Dataset):
	"""
	Loads both Kaggle datasets and produces (feature_vector, label) pairs.
	label = 1 (human) \| 0 (AI)
	"""

	def __init__(
	self,
	shanegerami_path: str,
	starblasters_path: str,
	extractor: HumanPatternFeatureExtractor,
	max_samples_per_source: int = 50000,
	):
	self.extractor = extractor
	self.samples = []

	# Load Dataset 1 (shanegerami)
	df1 = pd.read_csv(shanegerami_path).dropna()
	df1 = df1.sample(min(len(df1), max_samples_per_source), random_state=42)
	for _, row in df1.iterrows():
	self.samples.append({
	"text": str(row["text"]),
	"label": int(row["generated"] == 0), # 0→AI, 1→human → flip: 1=human
	"source": "shanegerami",
	})

	# Load Dataset 2 (starblasters — parquet)
	df2 = pd.read_parquet(starblasters_path).dropna()
	df2 = df2.sample(min(len(df2), max_samples_per_source), random_state=42)
	for _, row in df2.iterrows():
	label = 1 if str(row["label"]).lower() == "human" else 0
	self.samples.append({
	"text": str(row["text"]),
	"label": label,
	"source": "starblasters",
	})

	print(f"Total samples loaded: {len(self.samples)}")
	human = sum(1 for s in self.samples if s["label"] == 1)
	print(f" Human: {human} \| AI: {len(self.samples) - human}")

	def __len__(self):
	return len(self.samples)

	def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
	sample = self.samples[idx]
	features = self.extractor.extract(sample["text"])
	return torch.tensor(features), sample["label"]


	class HumanPatternClassifier(nn.Module):
	"""
	Lightweight MLP trained to distinguish human from AI writing.
	Input: feature vector from HumanPatternFeatureExtractor
	Output: probability that text is human-written (0 to 1)

	This is PRE-TRAINED on the Kaggle datasets, then FROZEN.
	Its output score is used as a loss signal in main model training.
	High score = human-like = good. Low score = AI-like = penalise.
	"""

	def __init__(self, input_dim: int = 17, hidden_dim: int = 128):
	super().__init__()
	self.net = nn.Sequential(
	nn.Linear(input_dim, hidden_dim),
	nn.LayerNorm(hidden_dim),
	nn.ReLU(),
	nn.Dropout(0.2),
	nn.Linear(hidden_dim, hidden_dim // 2),
	nn.ReLU(),
	nn.Dropout(0.1),
	nn.Linear(hidden_dim // 2, 1),
	nn.Sigmoid(),
	)

	def forward(self, features: torch.Tensor) -> torch.Tensor:
	"""Returns human-likeness score in [0, 1]. Higher = more human."""
	return self.net(features).squeeze(-1)

	def score(self, text: str, extractor: HumanPatternFeatureExtractor) -> float:
	"""Convenience: score a single text string."""
	features = torch.tensor(extractor.extract(text)).unsqueeze(0)
	with torch.no_grad():
	return self.forward(features).item()
	```

	---

	### `scripts/pretrain_human_pattern_classifier.py`

	```python
	"""
	Pre-trains the HumanPatternClassifier on both Kaggle datasets.
	Run this BEFORE the main training loop.
	The saved classifier weights are then loaded frozen during main training.

	Run: python scripts/pretrain_human_pattern_classifier.py
	Output: checkpoints/human_pattern_classifier.pt
	"""

	import torch
	import torch.nn as nn
	from torch.utils.data import DataLoader, random_split
	from sklearn.metrics import accuracy_score, roc_auc_score
	import numpy as np
	from loguru import logger
	import wandb

	from src.training.human_pattern_extractor import (
	HumanPatternFeatureExtractor,
	KaggleHumanPatternDataset,
	HumanPatternClassifier,
	)


	def train_classifier():
	wandb.init(project="dyslexia-writing-ai", name="human-pattern-pretrain")

	extractor = HumanPatternFeatureExtractor()
	dataset = KaggleHumanPatternDataset(
	shanegerami_path="data/raw/shanegerami/train_essays.csv",
	starblasters_path="data/raw/starblasters8/data.parquet",
	extractor=extractor,
	max_samples_per_source=50000, # 100k total — adjust for speed
	)

	train_size = int(0.85 * len(dataset))
	val_size = len(dataset) - train_size
	train_ds, val_ds = random_split(dataset, [train_size, val_size])

	train_loader = DataLoader(train_ds, batch_size=512, shuffle=True, num_workers=4)
	val_loader = DataLoader(val_ds, batch_size=512, shuffle=False, num_workers=4)

	input_dim = extractor.extract("sample text").shape[0]
	model = HumanPatternClassifier(input_dim=input_dim, hidden_dim=256)
	device = "cuda" if torch.cuda.is_available() else "cpu"
	model = model.to(device)

	optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
	scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)
	criterion = nn.BCELoss()

	best_auc = 0.0
	for epoch in range(20):
	# Train
	model.train()
	train_losses = []
	for features, labels in train_loader:
	features = features.to(device)
	labels = labels.float().to(device)
	preds = model(features)
	loss = criterion(preds, labels)
	optimizer.zero_grad()
	loss.backward()
	optimizer.step()
	train_losses.append(loss.item())

	# Validate
	model.eval()
	all_preds, all_labels = [], []
	with torch.no_grad():
	for features, labels in val_loader:
	features = features.to(device)
	preds = model(features).cpu().numpy()
	all_preds.extend(preds)
	all_labels.extend(labels.numpy())

	auc = roc_auc_score(all_labels, all_preds)
	acc = accuracy_score(all_labels, [1 if p > 0.5 else 0 for p in all_preds])
	scheduler.step()

	logger.info(f"Epoch {epoch+1:02d} \| Loss: {np.mean(train_losses):.4f} \| AUC: {auc:.4f} \| Acc: {acc:.4f}")
	wandb.log({"classifier/train_loss": np.mean(train_losses), "classifier/val_auc": auc, "classifier/val_acc": acc})

	if auc > best_auc:
	best_auc = auc
	torch.save(model.state_dict(), "checkpoints/human_pattern_classifier.pt")
	logger.info(f" ✓ Saved best classifier (AUC: {best_auc:.4f})")

	wandb.finish()
	logger.info(f"Pre-training complete. Best AUC: {best_auc:.4f}")
	logger.info("Classifier saved to: checkpoints/human_pattern_classifier.pt")


	if __name__ == "__main__":
	train_classifier()
	```

	---

	### Integration into Main Training Loss

	#### Updated `src/training/loss_functions.py` — add `L_human_pattern`

	```python
	"""
	UPDATED Combined Loss with Human-Pattern Term:

	L_total = L_CE + λ₁ · L_style + λ₂ · L_semantic + λ₃ · L_human_pattern

	L_human_pattern = 1 - HumanPatternClassifier.score(output_text)
	= reward for human-like output
	= penalty for AI-like output

	The HumanPatternClassifier is FROZEN — its weights do not update.
	It acts as a discriminator/critic, not a trainable component.
	λ₃ default = 0.4
	"""

	class CombinedCorrectionLossV2(nn.Module):

	def __init__(
	self,
	lambda_style: float = 0.3,
	lambda_semantic: float = 0.5,
	lambda_human_pattern: float = 0.4,
	classifier_path: str = "checkpoints/human_pattern_classifier.pt",
	sem_model_name: str = "all-mpnet-base-v2",
	device: str = "cuda",
	):
	super().__init__()
	self.lambda_style = lambda_style
	self.lambda_semantic = lambda_semantic
	self.lambda_human_pattern = lambda_human_pattern
	self.device = device

	# Load pre-trained frozen classifier
	from .human_pattern_extractor import HumanPatternClassifier, HumanPatternFeatureExtractor
	self.hp_extractor = HumanPatternFeatureExtractor()
	input_dim = self.hp_extractor.extract("sample").shape[0]
	self.hp_classifier = HumanPatternClassifier(input_dim=input_dim)
	self.hp_classifier.load_state_dict(torch.load(classifier_path, map_location=device))
	self.hp_classifier.to(device)
	for param in self.hp_classifier.parameters():
	param.requires_grad = False # FROZEN — never trains
	self.hp_classifier.eval()

	# Semantic model (also frozen)
	from sentence_transformers import SentenceTransformer
	self.sem_model = SentenceTransformer(sem_model_name, device=device)
	for param in self.sem_model.parameters():
	param.requires_grad = False

	self.ce_loss = nn.CrossEntropyLoss(ignore_index=-100)

	def _human_pattern_loss(self, output_texts: List[str]) -> torch.Tensor:
	"""
	For each output text, compute how AI-like it is.
	Loss = 1 - human_score (penalise AI-like outputs).

	The gradient flows back through the generation model via this loss.
	The classifier itself is frozen.
	"""
	features = torch.stack([
	torch.tensor(self.hp_extractor.extract(t))
	for t in output_texts
	]).to(self.device)

	with torch.no_grad():
	human_scores = self.hp_classifier(features) # [batch], values in [0,1]

	# Loss = average AI-likeness = 1 - average human-likeness
	return (1 - human_scores).mean()

	def forward(
	self,
	logits: torch.Tensor,
	labels: torch.Tensor,
	output_style_vec: Optional[torch.Tensor] = None,
	target_style_vec: Optional[torch.Tensor] = None,
	input_texts: Optional[List[str]] = None,
	output_texts: Optional[List[str]] = None,
	) -> Dict[str, torch.Tensor]:

	vocab_size = logits.shape[-1]
	l_ce = self.ce_loss(logits.view(-1, vocab_size), labels.view(-1))
	losses = {"l_ce": l_ce, "total": l_ce}

	if output_style_vec is not None and target_style_vec is not None:
	sim = F.cosine_similarity(output_style_vec, target_style_vec, dim=-1)
	l_style = (1 - sim).mean()
	losses["l_style"] = l_style
	losses["total"] = losses["total"] + self.lambda_style * l_style

	if input_texts is not None and output_texts is not None:
	input_embs = torch.tensor(self.sem_model.encode(input_texts), device=self.device)
	output_embs = torch.tensor(self.sem_model.encode(output_texts), device=self.device)
	sim = F.cosine_similarity(input_embs, output_embs, dim=-1)
	l_sem = (1 - sim).mean()
	losses["l_semantic"] = l_sem
	losses["total"] = losses["total"] + self.lambda_semantic * l_sem

	if output_texts is not None:
	l_hp = self._human_pattern_loss(output_texts)
	losses["l_human_pattern"] = l_hp
	losses["total"] = losses["total"] + self.lambda_human_pattern * l_hp

	return losses
	```

	---

	### Updated Mathematical Formulation

	```
	L_total = L_CE + λ₁ · L_style + λ₂ · L_semantic + λ₃ · L_human_pattern

	L_human_pattern = 1 - (1/N) Σᵢ HPC(φ(ŷᵢ))

	Where:
	HPC(·) = HumanPatternClassifier (frozen)
	φ(·) = HumanPatternFeatureExtractor
	ŷᵢ = model's generated output text for example i
	N = batch size

	λ₁ = 0.3 (style consistency weight)
	λ₂ = 0.5 (semantic preservation weight)
	λ₃ = 0.4 (human pattern reward weight)

	Total loss weights must sum interpretably:
	λ₁ + λ₂ + λ₃ = 1.2 (additive, CE is the base anchor)

	The HumanPatternClassifier is trained to maximise AUC on the two Kaggle datasets.
	Target pre-training performance: AUC > 0.88, Accuracy > 83%
	```

	---

	### Updated Training Sequence

	```bash
	# 0. Download both Kaggle datasets
	bash scripts/download_kaggle_datasets.sh

	# 1. Pre-train the HumanPatternClassifier (runs separately, ~1-2 hours)
	python scripts/pretrain_human_pattern_classifier.py

	# 2. Verify classifier quality (target AUC > 0.88)
	python scripts/evaluate_classifier.py --checkpoint checkpoints/human_pattern_classifier.pt

	# 3. Then run main model training (classifier is auto-loaded frozen)
	python scripts/train.py --config configs/training_config.yaml

	# 4. All four losses now tracked in W&B:
	# loss/ce, loss/style, loss/semantic, loss/human_pattern, loss/total
	```

	---

	### Updated `configs/training_config.yaml` additions

	```yaml
	# Add this section to configs/training_config.yaml:

	human_pattern:
	classifier_path: "checkpoints/human_pattern_classifier.pt"
	shanegerami_path: "data/raw/shanegerami/train_essays.csv"
	starblasters_path: "data/raw/starblasters8/data.parquet"
	max_samples_per_source: 50000
	pretrain_epochs: 20
	pretrain_lr: 1.0e-3
	pretrain_batch_size: 512
	target_auc: 0.88

	loss:
	lambda_style: 0.3
	lambda_semantic: 0.5
	lambda_human_pattern: 0.4 # NEW — added in v2
	sem_model_name: "all-mpnet-base-v2"
	```

	---

	### What the Two Datasets Teach the Model (Not What They Are Used For)

	\| Learning Target \| From Dataset \| Mechanism \|
	\|---\|---\|---\|
	\| Human sentence length is bursty/unpredictable \| Both datasets \| Burstiness feature in HPC \|
	\| Humans don't start every sentence the same way \| Both datasets \| Starter diversity feature \|
	\| AI text has lower GPT-2 perplexity \| Both (AI side) \| Perplexity feature in HPC \|
	\| AI overuses "delve", "tapestry", "crucial", "pivotal" \| Both (AI side) \| AI fingerprint word density \|
	\| AI overuses "Furthermore", "Moreover", "In conclusion" \| Both (AI side) \| Discourse marker density \|
	\| Humans have higher n-gram novelty \| Both \| n-gram novelty score \|
	\| 63 different LLMs share the same surface patterns \| starblasters8 \| Broad AI coverage in HPC training \|
	\| GPT-family essays are detectable at scale \| shanegerami \| Dense GPT signature learning \|

	---

	## 22. Complete Dataset Directory

	All publicly available datasets relevant to this system, across three categories.

	---

	### Category A — Grammar Error Correction (Core Training Data)

	\| Dataset \| Size \| Notes \| Access \| HuggingFace ID \|
	\|---\|---\|---\|---\|---\|
	\| W&I+LOCNESS \| 35k pairs \| Gold standard GEC, learner English, 5 proficiency levels \| Free registration \| `wi_locness` \|
	\| JFLEG \| 1.5k pairs \| 4 human references per sentence, fluency focus \| Public GitHub \| — \|
	\| CoNLL-2014 \| 1.3k pairs \| 2 human annotators, classic GEC benchmark \| Public \| — \|
	\| FCE Corpus \| 33k pairs \| Cambridge First Certificate essays with corrections \| Free registration \| — \|
	\| NUCLE \| 57k sentences \| NUS Corpus of Learner English, sentence-level errors \| Free registration \| — \|
	\| Lang-8 \| 1M+ pairs \| Crowdsourced learner writing corrections in 80 languages \| Request form \| — \|
	\| CLANG-8 \| 2.6M pairs \| Cleaned Lang-8, filtered for English quality \| HuggingFace \| `google/clang8` \|
	\| Falko-MERLIN \| 24k sentences \| German learner English (good for multilingual) \| Public \| — \|
	\| BEA-2019 Shared Task \| 4k test pairs \| Official GEC evaluation set, gold standard \| Free \| — \|

	---

	### Category B — Formality & Style Transfer (Style Training Data)

	\| Dataset \| Size \| Notes \| Access \| HuggingFace ID \|
	\|---\|---\|---\|---\|---\|
	\| GYAFC \| 105k pairs \| Yahoo Answers informal → formal, 2 domains \| Request Grammarly \| — \|
	\| YELP Sentiment Transfer \| 560k reviews \| Sentiment-controlled style transfer \| Public \| `yelp_review_full` \|
	\| Shakespeare Modern \| 21k lines \| Shakespearean → modern English parallel \| Public GitHub \| — \|
	\| Europarl \| 60M sentences \| Formal parliamentary discourse, 21 languages \| Public \| `Helsinki-NLP/europarl` \|
	\| WikiText-103 \| 103M tokens \| High-quality Wikipedia prose, formal register \| Public \| `wikitext` \|
	\| OpenWebText \| 40GB \| Curated human web text (Reddit upvoted links) \| Public \| `openwebtext` \|
	\| PAWS \| 108k pairs \| Paraphrase pairs with controlled syntactic diversity \| Public \| `paws` \|
	\| ParaBank2 \| 50M pairs \| Large-scale paraphrase pairs \| Public \| — \|

	---

	### Category C — Human vs AI Distinction (Anti-AI Training Data)

	#### Your Two Selected Datasets

	\| Dataset \| Size \| LLMs Covered \| Access \| Notes \|
	\|---\|---\|---\|---\|---\|
	\| shanegerami/ai-vs-human-text \| 500k essays \| GPT-family \| Kaggle \| Columns: text, generated(0/1) \|
	\| starblasters8/human-vs-llm-text-corpus \| 800k texts \| 63 LLMs \| Kaggle \| Parquet: text, label(str) \|

	#### Additional Highly Recommended

	\| Dataset \| Size \| LLMs Covered \| Access \| HuggingFace / URL \|
	\|---\|---\|---\|---\|---\|
	\| RAID Benchmark \| 6.2M generations \| 11 generators, 8 domains, 11 adversarial attacks \| Public \| `liamdugan/raid` \|
	\| HC3 (Human-ChatGPT Corpus) \| 125k QA pairs \| ChatGPT only \| Public \| `Hello-SimpleAI/HC3` \|
	\| HC3-Plus \| 210k pairs \| ChatGPT, semantic-invariant variants \| Public \| `Hello-SimpleAI/HC3-Chinese` \|
	\| M4GT-Bench \| 152k texts \| 7 LLMs, 8 languages, 8 domains \| Public \| `NicolaiSivesind/ChatGPT-Research-Abstracts` \|
	\| DeepfakeTextDetect \| 447k texts \| 27 LLMs, 10 domains \| Public \| `Li2023` / arxiv \|
	\| MGTBench \| 21k texts \| 6 LLMs, 3 domains \| Public \| `aadityaubhat/GPT-wiki-intro` \|
	\| MAGE Dataset \| 447k texts \| Largest multi-model human/AI corpus \| Public \| `yaful/MAGE` \|
	\| TuringBench \| 168k articles \| 20 LLMs including GPT-2 to GPT-3 \| Public \| — \|
	\| BUST \| 25.2k texts \| 7 generators, 4 domains \| Public \| — \|
	\| DetectRL \| 235k texts \| 4 LLMs, adversarial-robust benchmark \| Public \| — \|
	\| GPT-Wiki-Intro \| 150k intros \| GPT-3.5 vs Wikipedia introductions \| Public \| `aadityaubhat/GPT-wiki-intro` \|
	\| SemEval 2024 Task 8 \| ~70k texts \| Mixed human/AI, boundary detection task \| Public \| SemEval 2024 \|
	\| PeerRead \| 14.7k papers \| Scientific paper review AI vs human \| Public \| `allenai/PeerRead` \|
	\| ArXiv AI Abstract Dataset \| 500k+ abstracts \| Scientific writing, GPT vs real \| Public \| arxiv bulk API \|
	\| ELI5-Human-AI \| 30k pairs \| Mistral-7B vs human on Explain Like I'm 5 \| Public \| Research benchmark \|
	\| HC-Var \| 145k texts \| ChatGPT variants across prompting strategies \| Public \| — \|
	\| WritingPrompts (Human) \| 303k stories \| Reddit human creative writing — pure human signal \| Public \| `euclaise/writingprompts` \|
	\| MultiSocial \| 472k texts \| Social media, 22 languages, 7 LLMs \| Public \| — \|
	\| WETBench \| 101.9k texts \| Web & essay text, 4 LLMs \| Public \| — \|
	\| silentone0725/ai-human-text-detection-v1 \| 9 corpora merged \| HC3, RAID, M4GT-Bench + more, pre-cleaned \| Public \| `silentone0725/ai-human-text-detection-v1` \|

	---

	### Category D — Dyslexia-Specific Data

	\| Dataset \| Size \| Notes \| Access \| URL \|
	\|---\|---\|---\|---\|---\|
	\| DysLexML Corpus \| ~2k texts \| Actual dyslectic writing samples, annotated \| Academic request \| Research paper: Rello et al. \|
	\| POPSYCLE Corpus \| ~800 texts \| Dyslexic children's writing with expert annotations \| Academic request \| Lancaster University \|
	\| Write & Improve (W&I) subset \| ~5k texts \| Includes dyslexia-pattern learner errors \| Free registration \| Cambridge \|
	\| Synthetic (DyslexiaSimulator) \| Unlimited \| Generated by your own simulator (Layer 1) \| Self-generated \| `src/preprocessing/dyslexia_simulator.py` \|

	---

	### Recommended Dataset Priority Order for Training

	```
	Phase 1 — Classifier Pre-training (Human Pattern):
	1. starblasters8/human-vs-llm-text-corpus (800k, 63 LLMs — widest coverage)
	2. shanegerami/ai-vs-human-text (500k, dense GPT signal)
	3. RAID Benchmark (6.2M, adversarial robustness)
	4. MAGE Dataset (447k, 27 LLMs)

	Phase 2 — Core GEC Model Training:
	1. CLANG-8 (2.6M pairs, largest clean GEC)
	2. W&I+LOCNESS (35k, gold standard, highest quality)
	3. JFLEG (1.5k, fluency focus)
	4. Synthetic dyslexia pairs (generated, unlimited)

	Phase 3 — Style Transfer Training:
	1. GYAFC (105k formal/informal pairs)
	2. WikiText-103 (103M tokens, formal register)
	3. OpenWebText (40GB human web text)

	Phase 4 — Academic Register Fine-tuning:
	1. PeerRead (14.7k academic papers)
	2. ArXiv abstracts (500k+ scientific writing)
	3. Europarl (60M formal parliamentary)
	```

	---

	### `scripts/download_all_huggingface_datasets.py`

	```python
	"""
	Downloads all publicly available HuggingFace datasets automatically.
	Datasets requiring registration/request are flagged with instructions.

	Run: python scripts/download_all_huggingface_datasets.py
	"""

	from datasets import load_dataset
	import os

	os.makedirs("data/raw/hf", exist_ok=True)

	HF_DATASETS = [
	# (hf_identifier, config, split, output_subdir)
	("google/clang8", "en", "train", "clang8"),
	("liamdugan/raid", None, "train", "raid"),
	("Hello-SimpleAI/HC3", "all", "train", "hc3"),
	("yaful/MAGE", None, "train", "mage"),
	("aadityaubhat/GPT-wiki-intro", None, "train", "gpt_wiki_intro"),
	("euclaise/writingprompts", None, "train", "writing_prompts"),
	("wikitext", "wikitext-103-raw-v1", "train", "wikitext103"),
	("openwebtext", None, "train", "openwebtext"),
	("paws", "labeled_final", "train", "paws"),
	("allenai/PeerRead", "all", "train", "peerread"),
	("silentone0725/ai-human-text-detection-v1", None, "train", "merged_ai_human"),
	]

	for hf_id, config, split, subdir in HF_DATASETS:
	out_path = f"data/raw/hf/{subdir}"
	if os.path.exists(out_path):
	print(f"✓ Already exists: {subdir}")
	continue
	try:
	print(f"Downloading: {hf_id}...")
	ds = load_dataset(hf_id, config, split=split, trust_remote_code=True)
	ds.save_to_disk(out_path)
	print(f" ✓ Saved to {out_path} ({len(ds)} examples)")
	except Exception as e:
	print(f" ✗ Failed: {hf_id} — {e}")

	# Datasets requiring manual action
	MANUAL_DATASETS = {
	"W&I+LOCNESS": "https://www.cl.cam.ac.uk/research/nl/bea2019st/ (free registration)",
	"GYAFC": "https://github.com/raosudha89/GYAFC-corpus (email request to Grammarly)",
	"FCE Corpus": "https://ilexir.co.uk/datasets/index.html (free registration)",
	"NUCLE": "https://www.comp.nus.edu.sg/~nlp/corpora.html (free registration)",
	"Lang-8": "https://sites.google.com/site/naistlang8corpora/ (request form)",
	"DysLexML": "Contact Rello et al. authors directly via ResearchGate",
	}

	print("\n── Datasets requiring manual download ──")
	for name, url in MANUAL_DATASETS.items():
	print(f" {name}: {url}")
	```

	---

	Blueprint version 2.0 — Dyslexia Academic Writing Correction System
	Architecture: Style-Preserving Constrained Correction + Human-Pattern Anti-AI Training
	Datasets: 25+ sources · 10M+ training examples · 63 LLMs covered