organized code and set up chainlit for demos

.gitignore

@@ -54,6 +54,9 @@ data/
 outputs/
 runs/
 artifacts/
+trained_models/
+!trained_models/README.md
+.chainlit/
 *.pt
 *.pth
 *.ckpt

CONTRIBUTING.md

@@ -0,0 +1,27 @@
+# Contributing to mini_transformer
+
+Thanks for your interest in contributing! Follow these steps to get started.
+
+1. **Create the development environment**
+   ```bash
+   make conda-setup
+   ```
+
+2. **Run quality checks** before opening a pull request:
+   ```bash
+   make fmt
+   make lint
+   make type
+   make test
+   ```
+
+3. **Submit changes** via a pull request. Please describe the motivation, relevant issues, and testing performed.
+
+4. **Code style**
+   - Keep imports sorted and formatted with `ruff` and `black`.
+   - Prefer type annotations and existing helper utilities.
+
+5. **Doc updates**
+   - Update the README and any tutorial notebooks if behaviour changes.
+
+Feel free to open an issue for questions or proposals before starting larger pieces of work.

Dockerfile.dev

@@ -1,6 +1,6 @@
 # syntax=docker/dockerfile:1
-FROM python:3.12-slim
-# FROM = choose a base OS + language runtime (Debian slim with Python 3.12)
+FROM python:3.11-slim
+# FROM = choose a base OS + language runtime (Debian slim with Python 3.11)
 
 ENV DEBIAN_FRONTEND=noninteractive \
     PIP_NO_CACHE_DIR=1 \
@@ -17,22 +17,12 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 # rm -rf ...                -> clean apt cache to shrink the layer
 
 WORKDIR /app
-# WORKDIR = default directory for the following Dockerfile steps and 'docker run' processes
 
-COPY requirements.txt pyproject.toml ./
-# COPY = bring files from host into the image
-# Only copying dependency metadata first lets Docker cache layer if code changes later
+COPY pyproject.toml README.md LICENSE ./
+COPY src ./src
 
-RUN pip install -U pip && pip install -r requirements.txt
-# pip install -U pip    -> upgrade pip in the image
-# pip install -r ...    -> install your Python runtime dependencies
+RUN pip install -U pip && pip install -e .[dev,server,viz,notebook]
 
 COPY . .
-# Now copy your actual code (done AFTER deps for better caching)
 
-RUN pip install -e .[dev]
-# -e . = editable install (your src is importable as a package without rebuilding)
-
-CMD ["pytest"]
-# CMD = default command when a container starts from this image
-# Here: run unit tests quietly (-q) in tests/units
+CMD ["bash"]

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Ala Boussoffara
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

Makefile

@@ -1,39 +1,85 @@
-.PHONY: help conda-setup precommit lint fmt type test clean docker-build docker-run
+ENV?=mini-transformer
+CONDA_RUN=conda run -n $(ENV)
+PYTHON?=python
+
+.PHONY: help create-env install precommit lint lint-check fmt fmt-check type test cov fetch-test fetch-small infer serve ui clean clean-all docker-build docker-run
 
 help:
-	@echo "Common Makefile commands:"
-	@echo "  conda-setup     - Create and activate conda env, install all dependencies, setup pre-commit"
-	@echo "  precommit       - Run all pre-commit hooks (lint, fmt, type)"
-	@echo "  lint            - Run ruff linter"
-	@echo "  fmt             - Run ruff formatter and black"
-	@echo "  type            - Run mypy type checks"
-	@echo "  test            - Run all tests with pytest and coverage"
-	@echo "  clean           - Remove build, cache, and venv artifacts"
-	@echo "  docker-build    - Build the Docker dev image"
-	@echo "  docker-run      - Run the Docker dev container (interactive)"
-
-conda-setup:
-	conda env create -n transformer -f environment.yml|| conda env update -n transformer -f environment.yml
-	conda run -n transformer pip install -e .[dev]
-	conda run -n transformer pre-commit install
+	@echo "Available targets:"
+	@echo "  create-env     - Create or update the $(ENV) conda environment"
+	@echo "  install        - Install project with dev/notebook/viz extras into the active environment"
+	@echo "  precommit      - Run all pre-commit hooks"
+	@echo "  lint           - Run ruff lint (auto-fixes with fmt beforehand)"
+	@echo "  lint-check     - Run ruff in check-only mode"
+	@echo "  fmt            - Format code with ruff + black"
+	@echo "  fmt-check      - Check formatting without modifying files"
+	@echo "  type           - Run mypy type checks"
+	@echo "  test           - Run unit tests"
+	@echo "  cov            - Run unit tests with coverage report"
+	@echo "  fetch-test     - Download the AlaBoussoffara/transformer_test model"
+	@echo "  fetch-small    - Download the AlaBoussoffara/transformer_small model"
+	@echo "  infer          - Quick inference using the demo model"
+	@echo "  serve          - Start the FastAPI server (reload mode)"
+	@echo "  ui             - Launch the Chainlit UI"
+	@echo "  clean          - Remove build/test caches"
+	@echo "  clean-all      - Clean caches plus notebooks outputs"
+	@echo "  docker-build   - Build the development Docker image"
+	@echo "  docker-run     - Drop into the development container"
+
+create-env:
+	conda env create -n $(ENV) -f environment.yml || conda env update -n $(ENV) -f environment.yml
+	$(CONDA_RUN) pip install -e .[dev,notebook,viz]
+	$(CONDA_RUN) pre-commit install
+
+install:
+	$(PYTHON) -m pip install -e .[dev,notebook,viz]
 
 precommit:
-	conda run -n transformer pre-commit run --all-files
+	$(CONDA_RUN) pre-commit run --all-files
+
+lint: fmt
+	$(CONDA_RUN) ruff check .
 
-lint:
-	conda run -n transformer ruff check .
+lint-check:
+	$(CONDA_RUN) ruff check .
 
 fmt:
-	conda run -n transformer ruff format .
-	conda run -n transformer black src tests
+	$(CONDA_RUN) ruff format .
+	$(CONDA_RUN) black src tests
+
+fmt-check:
+	$(CONDA_RUN) ruff format . --check
+	$(CONDA_RUN) black src tests --check
 
 type:
-	conda run -n transformer mypy src
+	$(CONDA_RUN) mypy src
 
 test:
-	conda run -n transformer pytest
+	$(CONDA_RUN) pytest
+
+cov:
+	$(CONDA_RUN) pytest --cov=mini_transformer --cov-report=term-missing
+
+fetch-test:
+	$(CONDA_RUN) mini-transformer-fetch AlaBoussoffara/transformer_test --force
+
+fetch-small:
+	$(CONDA_RUN) mini-transformer-fetch AlaBoussoffara/transformer_small --force
+
+infer:
+	$(CONDA_RUN) mini-transformer-infer --model small_model_v1 -t "translate this sentence."
+
+serve:
+	$(CONDA_RUN) mini-transformer-serve --model small_model_v1 --reload
+
+ui:
+	$(CONDA_RUN) mini-transformer-ui --model small_model_v1 --host 0.0.0.0 --port 8000
+
 clean:
-	rm -rf .venv .pytest_cache .mypy_cache dist build *.egg-info
+	rm -rf .pytest_cache .mypy_cache .ruff_cache dist build *.egg-info
+
+clean-all: clean
+	rm -rf notebooks/.ipynb_checkpoints outputs/*
 
 docker-build:
 	docker compose build dev --no-cache

README.md

@@ -1,8 +1,130 @@
-# transformer
+# mini-transformer
 
-Minimal encoder–decoder Transformer.
+A compact, encoder–decoder Transformer packaged so it can be used both as a Python library and as a ready-to-run demo. The repository mirrors the installable package layout so cloning the repo or `pip install mini-transformer` gives the same structure and tooling.
 
-## Quickstart
+## Highlights
+- Typed, unit-tested implementation of an encoder–decoder Transformer with reusable building blocks (`mini_transformer/modules`).
+- Hydra configuration system with both packaged defaults (`mini_transformer/conf`) and editable configs in the repo.
+- Ready-made UIs: FastAPI REST server and Chainlit chat interface.
+- CLI convenience commands for inference, serving, UI launch, and Hugging Face model downloads.
+- Makefile shortcuts, notebooks, and Docker dev environment for day-to-day work.
+
+## Installation
+```bash
+# inside the repository (editable install with optional extras)
+pip install -e .[server,viz]
+```
+Or install the published package:
+```bash
+pip install "mini-transformer[server,viz]"
+```
+
+## Repository Layout
+- `src/mini_transformer/` – installable package (model code, CLI, apps, packaged Hydra configs).
+- `configs/` – editable Hydra configs for local experiments.
+- `trained_models/` – place downloaded or exported checkpoints here (see below).
+- `notebooks/` – exploratory notebooks (`train.ipynb`, `tokenizer.ipynb`, etc.).
+- `tests/` – unit and smoke tests.
+- Supporting files: `pyproject.toml`, `Makefile`, `Dockerfile.dev`, `environment.yml`, etc.
+
+## Preparing Models
+Download the demo models hosted on Hugging Face:
+```bash
+mini-transformer-fetch --model-id AlaBoussoffara/transformer_test
+mini-transformer-fetch --model-id AlaBoussoffara/transformer_small
+```
+Use `--name` to customise the local directory and `--force` to refresh an existing download.
+The local folder name defaults to `<author>__<repo>` (e.g. `AlaBoussoffara__transformer_small`). The Hugging Face repo is downloaded as a whole, so if the model files live inside a subfolder (for example `transformer_small/small_model_v1/`), extract or move those contents up into the root of `trained_models/<local-name>/` so the structure contains `configs/`, `checkpoints/`, and `tokenizer/` directly.
+In other words, after fetching `AlaBoussoffara/transformer_small`, move the `transformer_small/small_model_v1/` directory into `trained_models/AlaBoussoffara__transformer_small/`.
+
+You can also place your own model manually under `trained_models/<model-name>/`:
+```
+trained_models/
+  my-model/
+    configs/
+      config_inference.yaml
+    checkpoints/
+      best.pt
+    tokenizer/
+      tokenizer.json
+```
+Relative paths in `config_inference.yaml` should stay inside the model folder. Set `MINI_TRANSFORMER_MODELS=/path/to/trained_models` if you store models elsewhere.
+
+## CLI Usage
+Install the optional extras (see Installation) and use the commands below.
+
+```bash
+mini-transformer-infer --model AlaBoussoffara__transformer_small -t "Once upon a time"
+mini-transformer-serve --model AlaBoussoffara__transformer_small --reload
+mini-transformer-ui --model AlaBoussoffara__transformer_small --host 0.0.0.0 --port 8000
+```
+If you see “Tokenizer file not found”, update the model config or set `MINI_TRANSFORMER_TOKENIZER_PATH` to the correct JSON file.
+
+## Chainlit Demo
+Follow these steps to spin up the bundled Chainlit chat UI for a local inference demo.
+
+1. Install the project with the server extras (or run `make create-env`):
+   ```bash
+   pip install -e .[server]
+   ```
+2. Download a demo checkpoint (skip if you already have one under `trained_models/`):
+   ```bash
+   mini-transformer-fetch --model-id AlaBoussoffara/transformer_small
+   ```
+3. Launch the Chainlit UI, pointing to the model folder you want to use (defaults to the first available model):
+   ```bash
+   mini-transformer-ui --model AlaBoussoffara__transformer_small --host 0.0.0.0 --port 8000
+   ```
+4. Open http://localhost:8000 in your browser, send a prompt, and use `/model` in chat to switch between downloaded checkpoints.
+
+### Environment Variables
+- `MINI_TRANSFORMER_MODELS` – override the models root directory.
+- `MINI_TRANSFORMER_MODEL_NAME` – preselect a model for the server/UI.
+- `MINI_TRANSFORMER_CONFIG_DIR` / `MINI_TRANSFORMER_CONFIG_NAME` – point to custom Hydra configs.
+- `MINI_TRANSFORMER_UI_HOST` / `MINI_TRANSFORMER_UI_PORT` – defaults for Chainlit binding.
+- Optional overrides: `MINI_TRANSFORMER_CHECKPOINT_BEST`, `MINI_TRANSFORMER_TOKENIZER_PATH`, `MINI_TRANSFORMER_OUTPUT_DIR`, etc.
+
+## Programmatic Inference
+```python
+from mini_transformer.model_loader import compose_model_config
+from mini_transformer.inference import run_inference
+
+cfg = compose_model_config("AlaBoussoffara__transformer_small")
+cfg.input_text = "Hello world"
+print(run_inference(cfg)[0])
+```
+
+## Development Quickstart
 ```bash
 make conda-setup
+make lint
+make type
 make test
+```
+
+## Makefile Shortcuts
+```bash
+make help          # list common tasks
+make create-env    # create/update the conda env and install extras
+make lint          # format + lint (ruff + black)
+make lint-check    # lint without auto-fixes
+make fmt           # format code
+make type          # mypy
+make test          # pytest
+make cov           # pytest with coverage
+make fetch-test    # download AlaBoussoffara/transformer_test
+make fetch-small   # download AlaBoussoffara/transformer_small
+make infer         # quick demo inference
+make serve         # run FastAPI server (reload mode)
+make ui            # launch Chainlit UI
+make docker-build  # build the dev container
+make docker-run    # open an interactive shell in the container
+```
+
+## Full Workflow At A Glance
+1. Fetch a model: `mini-transformer-fetch --model-id AlaBoussoffara/transformer_small`
+2. Run inference: `mini-transformer-infer --model AlaBoussoffara__transformer_small -t "Once upon a time"`
+3. Launch FastAPI (optional): `mini-transformer-serve --model AlaBoussoffara__transformer_small --reload`
+4. Start Chainlit UI (optional): `mini-transformer-ui --model AlaBoussoffara__transformer_small --host 0.0.0.0 --port 8000`
+
+Environment variables (`MINI_TRANSFORMER_MODELS`, `MINI_TRANSFORMER_MODEL_NAME`, etc.) let you tailor the workflow to your setup.

architecture.txt

@@ -1,28 +0,0 @@
-Transformer:
-    embedding:
-        input_embedding (same for encoder and decoder)
-        positional_embedding (learned or sinusoidal)
-        dropout
-    encoder
-        encoder_layer *N
-            multihead_self_attention
-            dropout+norm
-            feedforward:
-                linear
-                relu
-                linear
-            dropout+norm
-    decoder
-        decoder_layer*N
-            masked_multihead_self_attention
-            dropout+norm
-            multihead_cross_attention
-            dropout+norm
-            feedforward:
-                linear
-                relu
-                linear
-            dropout+norm
-    lm_head:
-        linear (same parameters as the embedding)
-        softmax

chainlit.md

@@ -0,0 +1,14 @@
+# Welcome to Chainlit! 🚀🤖
+
+Hi there, Developer! 👋 We're excited to have you on board. Chainlit is a powerful tool designed to help you prototype, debug and share applications built on top of LLMs.
+
+## Useful Links 🔗
+
+- **Documentation:** Get started with our comprehensive [Chainlit Documentation](https://docs.chainlit.io) 📚
+- **Discord Community:** Join our friendly [Chainlit Discord](https://discord.gg/k73SQ3FyUh) to ask questions, share your projects, and connect with other developers! 💬
+
+We can't wait to see what you create with Chainlit! Happy coding! 💻😊
+
+## Welcome screen
+
+To modify the welcome screen, edit the `chainlit.md` file at the root of your project. If you do not want a welcome screen, just leave this file empty.

configs/generation/generation.yaml

@@ -1,4 +1,4 @@
-max_new_tokens: 64
+max_new_tokens: 128
 temperature: 1
 top_k: 10
 top_p: 0.6

configs/model/large_model.yaml

@@ -2,9 +2,9 @@ name: large_model
 best_checkpoint_path: ./outputs/checkpoints/large_model/best.ckpt
 latest_checkpoint_path: ./outputs/checkpoints/large_model/latest.ckpt
 tokenizer: bpe_32k
-d_model: 512
-num_heads: 8
-num_layers: 6
+d_model: 2048
+num_heads: 32
+num_layers: 16
 d_ff: 2048
 dropout_rate: 0.1
 vocab_size: 32768
@@ -12,3 +12,4 @@ max_seq_len: 256
 pad_id: 0
 bos_id: 1
 eos_id: 2
+layer_norm_style: pre

configs/model/medium_model.yaml

@@ -2,13 +2,14 @@ name: medium_model
 best_checkpoint_path: ./outputs/checkpoints/medium_model/best.ckpt
 latest_checkpoint_path: ./outputs/checkpoints/medium_model/latest.ckpt
 tokenizer: bpe_16k
-d_model: 512
-num_heads: 8
-num_layers: 6
+d_model: 1024
+num_heads: 16
+num_layers: 8
 d_ff: 2048
 dropout_rate: 0.1
 vocab_size: 16384
-max_seq_len: 256
+max_seq_len: 128
 pad_id: 0
 bos_id: 1
 eos_id: 2
+layer_norm_style: pre

configs/model/small_model.yaml

@@ -2,8 +2,8 @@ name: small_model
 best_checkpoint_path: ./outputs/checkpoints/small_model/best.ckpt
 latest_checkpoint_path: ./outputs/checkpoints/small_model/latest.ckpt
 tokenizer: bpe_8k
-d_model: 256
-num_heads: 8
+d_model: 512
+num_heads: 4
 num_layers: 4
 d_ff: 1024
 dropout_rate: 0.1
@@ -12,3 +12,4 @@ max_seq_len: 128
 pad_id: 0
 bos_id: 1
 eos_id: 2
+layer_norm_style: post

configs/model/test_model.yaml

@@ -3,8 +3,8 @@ best_checkpoint_path: ./outputs/checkpoints/test_model/best.pt
 latest_checkpoint_path: ./outputs/checkpoints/test_model/latest.pt
 tokenizer: bpe_8k
 d_model: 512
-num_heads: 8
-num_layers: 8
+num_heads: 4
+num_layers: 4
 d_ff: 1024
 dropout_rate: 0.1
 vocab_size: 8192
@@ -12,3 +12,4 @@ max_seq_len: 128
 pad_id: 0
 bos_id: 1
 eos_id: 2
+layer_norm_style: post

docker-compose.yml

@@ -3,12 +3,19 @@ services:
     build:
       context: .
       dockerfile: Dockerfile.dev
-    image: transformer:dev
-    container_name: transformer_dev
+    image: mini-transformer:dev
+    container_name: mini_transformer_dev
     working_dir: /app
     volumes:
-      - ./artifacts:/app/artifacts
+     - ./outputs:/app/outputs
+     - ./trained_models:/app/trained_models
+    environment:
+      MINI_TRANSFORMER_MODELS: /app/trained_models
+      MINI_TRANSFORMER_UI_HOST: 0.0.0.0
+      MINI_TRANSFORMER_UI_PORT: 8000
     stdin_open: true
     tty: true
     gpus: all
     command: bash
+    ports:
+      - "8000:8000"

environment.yml

@@ -1,17 +1,28 @@
-name: transformer
+name: mini-transformer
 channels:
   - defaults
 dependencies:
-  - python=3.12
+  - python=3.11
   - pip
   - pip:
-      - --extra-index-url https://download.pytorch.org/whl/cu129
-      - torch==2.8.0+cu129
+      - --extra-index-url https://download.pytorch.org/whl/cu121
+      - torch==2.3.1
+      - numpy
       - transformers
       - omegaconf
       - hydra-core
-      - jupyterlab
       - tokenizers
       - datasets
+      - huggingface_hub
+      - matplotlib
+      - jupyterlab
       - fastapi
       - uvicorn[standard]
+      - chainlit
+      - pytest
+      - pytest-cov
+      - ruff
+      - black
+      - mypy
+      - pre-commit
+      - httpx

{scripts → notebooks}/load_data.ipynb

@@ -40,7 +40,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/home/ala-boussoffara/miniconda3/envs/transformer/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "/home/ala-boussoffara/miniconda3/envs/mini-transformer/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
       "  from .autonotebook import tqdm as notebook_tqdm\n",
       "Generating train split: 179435 examples [00:00, 1366541.45 examples/s]\n",
       "Generating validation split: 903 examples [00:00, 558869.19 examples/s]\n",
@@ -98,7 +98,7 @@
     "import yaml\n",
     "from datasets import DatasetDict\n",
     "\n",
-    "from transformer.utils import set_global_seed\n",
+    "from mini_transformer.utils import set_global_seed\n",
     "\n",
     "with open(\"configs/runtime/local.yaml\", encoding=\"utf-8\") as cfg_file:\n",
     "    runtime_cfg = yaml.safe_load(cfg_file) or {}\n",
@@ -202,10 +202,10 @@
     "import numpy as np\n",
     "from hydra import compose, initialize\n",
     "from hydra.core.global_hydra import GlobalHydra\n",
+    "from mini_transformers import PreTrainedTokenizerFast\n",
     "from omegaconf import OmegaConf\n",
-    "from transformers import PreTrainedTokenizerFast\n",
     "\n",
-    "from transformer.configs import TokenizerCfg\n",
+    "from mini_transformer.configs import TokenizerCfg\n",
     "\n",
     "\n",
     "def load_tokenizer(name: str) -> PreTrainedTokenizerFast:\n",

{scripts → notebooks}/tokenizer.ipynb

@@ -18,7 +18,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/home/ala-boussoffara/miniconda3/envs/transformer/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "/home/ala-boussoffara/miniconda3/envs/mini-transformer/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
       "  from .autonotebook import tqdm as notebook_tqdm\n"
      ]
     }
@@ -31,7 +31,7 @@
     "from tokenizers import Tokenizer, models, pre_tokenizers, trainers\n",
     "from tokenizers.processors import TemplateProcessing\n",
     "\n",
-    "from transformer.configs import TokenizerCfg"
+    "from mini_transformer.configs import TokenizerCfg"
    ]
   },
   {

{scripts → notebooks}/train.ipynb

@@ -30,15 +30,15 @@
     "from hydra import compose, initialize\n",
     "from hydra.core.global_hydra import GlobalHydra\n",
     "from hydra.utils import to_absolute_path\n",
+    "from mini_transformers import PreTrainedTokenizerFast\n",
     "from omegaconf import OmegaConf\n",
     "from torch.nn.utils import clip_grad_norm_\n",
     "from torch.utils.data import DataLoader\n",
     "from tqdm import tqdm\n",
-    "from transformers import PreTrainedTokenizerFast\n",
     "\n",
-    "from transformer import BasicEncoderDecoderTransformer\n",
-    "from transformer.configs import ModelCfg, TokenizerCfg, TrainAppCfg\n",
-    "from transformer.utils import (\n",
+    "from mini_transformer import BasicEncoderDecoderTransformer\n",
+    "from mini_transformer.configs import ModelCfg, TokenizerCfg, TrainAppCfg\n",
+    "from mini_transformer.utils import (\n",
     "    check_tokenizer_model_compatibility,\n",
     "    debug_transformer_forward,\n",
     "    make_worker_init_fn,\n",
@@ -277,7 +277,7 @@
     "        raise ValueError(\"No trainable parameters found when setting up the optimizer.\")\n",
     "\n",
     "    return torch.optim.AdamW(\n",
-    "        model.parameters(),\n",
+    "        param_groups,\n",
     "        lr=optim_cfg.lr,\n",
     "        betas=optim_cfg.betas,\n",
     "        eps=optim_cfg.eps,\n",

pyproject.toml

@@ -1,65 +1,91 @@
-[project]
-name = "transformer"
-version = "0.1.0"
-description = "A transformer implementation"
-readme = "README.md"
-requires-python = ">=3.10"
-dependencies = [
-  "torch>=2.6",
-  "transformers",
-  "omegaconf",
-  "hydra-core",
-  "jupyterlab",
-  "tokenizers",
-  "datasets",
-]
-
-[project.optional-dependencies]
-dev = [
-  "pytest",
-  "pytest-cov",
-  "ruff",
-  "black",
-  "mypy",
-  "pre-commit",
-]
-server = [
-  "fastapi",
-  "uvicorn[standard]",
-]
-
-[tool.setuptools]
-package-dir = {"" = "src"}
-
-[build-system]
-requires = ["setuptools>=42", "wheel"]
-build-backend = "setuptools.build_meta"
-
-[tool.setuptools.packages.find]
-where = ["src"]
-
-[tool.pytest.ini_options]
-minversion = "7.0"
-addopts = "-ra -q"
-testpaths = ["tests"]
-
-[tool.black]
-line-length = 100
-target-version = ["py310"]
-
-[tool.ruff]
-line-length = 100
-
-[tool.ruff.lint]
-select = ["E","F","I","B","UP"]
-ignore = ["E501"]
-
-[tool.mypy]
-python_version = "3.11"
-strict = false
-ignore_missing_imports = true
-
-# Make mypy understand the src/ layout
-files = ["src"]
-mypy_path = ["src"]
-explicit_package_bases = true
+[project]
+name = "mini-transformer"
+version = "0.1.0"
+description = "A mini encoder–decoder transformer implementation"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { file = "LICENSE" }
+authors = [
+  { name = "Ala Boussoffara" }
+]
+dependencies = [
+  "torch>=2.3",
+  "numpy",
+  "transformers",
+  "omegaconf",
+  "hydra-core",
+  "tokenizers",
+  "datasets",
+  "huggingface_hub",
+]
+
+[project.optional-dependencies]
+dev = [
+  "pytest",
+  "pytest-cov",
+  "ruff",
+  "black",
+  "mypy",
+  "pre-commit",
+  "httpx",
+  "fastapi",
+  "uvicorn[standard]",
+  "chainlit",
+]
+server = [
+  "fastapi",
+  "uvicorn[standard]",
+  "chainlit",
+]
+viz = [
+  "matplotlib",
+]
+notebook = [
+  "jupyterlab",
+]
+
+
+[project.scripts]
+mini-transformer-infer = "mini_transformer.cli:infer_main"
+mini-transformer-serve = "mini_transformer.cli:serve_main"
+mini-transformer-ui = "mini_transformer.cli:ui_main"
+mini-transformer-fetch = "mini_transformer.cli:fetch_main"
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.package-data]
+"mini_transformer" = ["conf/**/*.yaml"]
+
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+minversion = "7.0"
+addopts = "-ra -q"
+testpaths = ["tests"]
+
+[tool.black]
+line-length = 100
+target-version = ["py310"]
+
+[tool.ruff]
+line-length = 100
+
+[tool.ruff.lint]
+select = ["E","F","I","B","UP"]
+ignore = ["E501"]
+
+[tool.mypy]
+python_version = "3.11"
+strict = false
+ignore_missing_imports = true
+
+# Make mypy understand the src/ layout
+files = ["src"]
+mypy_path = ["src"]
+explicit_package_bases = true

scripts/inference.py

@@ -1,76 +0,0 @@
-"""CLI entry point for running transformer inference with deterministic sampling."""
-
-from __future__ import annotations
-
-import hydra
-import torch
-from hydra.utils import to_absolute_path
-from omegaconf import DictConfig, OmegaConf
-from transformers import PreTrainedTokenizerFast
-
-from transformer import BasicEncoderDecoderTransformer
-from transformer.configs import InferAppCfg, ModelCfg, TokenizerCfg
-from transformer.utils import check_tokenizer_model_compatibility, set_global_seed
-
-
-def inference(cfg: DictConfig):
-    scfg_temp = OmegaConf.merge(OmegaConf.structured(InferAppCfg), cfg)
-    scfg: InferAppCfg = OmegaConf.to_object(scfg_temp)
-    model_cfg = ModelCfg(**vars(scfg.model))
-    tokenizer_cfg = TokenizerCfg(**vars(scfg.tokenizer))
-
-    # Ensure reproducible sampling when requested
-    set_global_seed(scfg.runtime.seed)
-
-    # Compatibility checks
-    check_tokenizer_model_compatibility(model_cfg, tokenizer_cfg)
-
-    transformer = BasicEncoderDecoderTransformer(model_cfg)
-    tokenizer = PreTrainedTokenizerFast(
-        tokenizer_file=to_absolute_path(tokenizer_cfg.path),
-        bos_token=cfg.tokenizer.bos_token,
-        eos_token=cfg.tokenizer.eos_token,
-        unk_token=cfg.tokenizer.unk_token,
-        pad_token=cfg.tokenizer.pad_token,
-        model_max_length=tokenizer_cfg.max_seq_len,
-    )
-
-    if scfg.model.best_checkpoint_path:
-        ckpt = torch.load(to_absolute_path(scfg.model.best_checkpoint_path), map_location="cpu")
-        state = ckpt.get("model_state_dict", ckpt)
-        transformer.load_state_dict(state, strict=False)
-    transformer.eval()
-
-    text = scfg.input_text
-    if not text:
-        raise SystemExit('Pass text like: input_text="hello world"')
-
-    encoded = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
-    src_ids = encoded["input_ids"]
-    src_padd_mask = encoded["attention_mask"] == 0
-
-    tgt_ids = transformer.generate(
-        src_ids,
-        src_padd_mask,
-        max_new_tokens=scfg.generation.max_new_tokens,
-        temperature=scfg.generation.temperature,
-        top_k=scfg.generation.top_k,
-        top_p=scfg.generation.top_p,
-        do_sample=scfg.generation.do_sample,
-        presence_penalty=scfg.generation.presence_penalty,
-        frequency_penalty=scfg.generation.frequency_penalty,
-        no_repeat_ngram=scfg.generation.no_repeat_ngram,
-        min_steps_before_eos=scfg.generation.min_steps_before_eos,
-        seed=scfg.runtime.seed,
-    )
-    output = tokenizer.batch_decode(tgt_ids, skip_special_tokens=False)
-    return output
-
-
-@hydra.main(config_path="../configs", config_name="infer_mode", version_base=None)
-def main(cfg: DictConfig):
-    print(inference(cfg))
-
-
-if __name__ == "__main__":
-    main()

scripts/server.py

@@ -1,26 +0,0 @@
-from __future__ import annotations
-
-from fastapi import FastAPI
-from hydra import compose, initialize
-from omegaconf import DictConfig
-
-from scripts.inference import inference as run_inference
-
-app = FastAPI(title="Transformer Inference API")
-
-
-@app.get("/healthz")
-def healthz():
-    return {"status": "ok"}
-
-
-@app.get("/generate")
-def generate(text: str | None = None):
-    # Compose the default config exactly as inference.py expects
-    with initialize(config_path="../configs", version_base=None):
-        cfg: DictConfig = compose(config_name="infer_mode")
-    # If a text query parameter is provided, override the config's input_text
-    if text is not None and text.strip():
-        cfg.input_text = text
-    outputs = run_inference(cfg)  # returns a list of strings
-    return {"outputs": outputs}

src/mini_transformer/__init__.py

@@ -0,0 +1,14 @@
+"""mini_transformer public API."""
+
+from __future__ import annotations
+
+from importlib import metadata
+
+from .transformer import BasicEncoderDecoderTransformer
+
+try:  # pragma: no cover - during editable installs metadata may be missing
+    __version__ = metadata.version("mini-transformer")
+except metadata.PackageNotFoundError:  # pragma: no cover
+    __version__ = "0.1.0"
+
+__all__ = ["BasicEncoderDecoderTransformer", "__version__"]

src/mini_transformer/apps/__init__.py

@@ -0,0 +1,3 @@
+"""Interactive applications (Chainlit UI & FastAPI server)."""
+
+__all__ = ["chainlit_app", "server"]

src/mini_transformer/apps/chainlit_app.py

@@ -0,0 +1,220 @@
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import chainlit as cl
+from hydra.core.global_hydra import GlobalHydra
+from omegaconf import DictConfig
+
+from mini_transformer.inference import run_inference
+from mini_transformer.model_loader import (
+    MODEL_NAME_ENV,
+    compose_model_config,
+    ensure_models_root,
+    list_model_names,
+)
+
+JOB_NAME = "chainlit_app"
+MODEL_SESSION_KEY = "selected_model"
+HAS_ACTION_PROMPT = hasattr(cl, "AskActionMessage") and hasattr(cl, "Action")
+
+
+def _discover_action_fields() -> set[str]:
+    if not HAS_ACTION_PROMPT:
+        return set()
+    action_cls = cl.Action
+    for attr in ("model_fields", "__fields__"):
+        fields = getattr(action_cls, attr, None)
+        if isinstance(fields, dict):
+            return set(fields.keys())
+    return set()
+
+
+ACTION_FIELDS = _discover_action_fields()
+
+
+class ModelOption:
+    """Lightweight container describing a trained model configuration."""
+
+    def __init__(self, name: str) -> None:
+        self.name = name
+
+    @property
+    def job_name(self) -> str:
+        return f"{JOB_NAME}_{self.name}"
+
+
+def discover_model_options() -> list[ModelOption]:
+    ensure_models_root()
+    return [ModelOption(name) for name in list_model_names()]
+
+
+def restore_model_from_session() -> ModelOption | None:
+    data: dict[str, str] | None = cl.user_session.get(MODEL_SESSION_KEY)
+    if not data:
+        env_model = os.environ.get(MODEL_NAME_ENV)
+        if env_model and env_model in list_model_names():
+            return ModelOption(name=env_model)
+        return None
+    name = data.get("name")
+    if not name:
+        return None
+    if name not in list_model_names():
+        return None
+    return ModelOption(name=name)
+
+
+def persist_model_selection(option: ModelOption) -> None:
+    cl.user_session.set(MODEL_SESSION_KEY, {"name": option.name})
+
+
+def extract_action_value(result: Any) -> str | None:
+    if result is None:
+        return None
+    if isinstance(result, str):
+        return result
+    if isinstance(result, dict):
+        payload = result.get("payload")
+        if isinstance(payload, dict):
+            selected = payload.get("model") or payload.get("value") or payload.get("name")
+            if selected:
+                return str(selected)
+        for key in ("value", "name", "action"):
+            value = result.get(key)
+            if value:
+                return str(value)
+        return None
+    value = getattr(result, "value", None)
+    if value:
+        return str(value)
+    name = getattr(result, "name", None)
+    if name:
+        return str(name)
+    payload = getattr(result, "payload", None)
+    if isinstance(payload, dict):
+        selected = payload.get("model") or payload.get("value") or payload.get("name")
+        if selected:
+            return str(selected)
+    return None
+
+
+async def prompt_model_selection(header: str | None = None) -> ModelOption | None:
+    options = discover_model_options()
+    if not options:
+        await cl.Message(
+            content=(
+                "No trained models found in `trained_models/`. "
+                "Place a model folder with configs/config_inference.yaml there and restart the chat."
+            )
+        ).send()
+        return None
+
+    if not HAS_ACTION_PROMPT:
+        chosen = options[0]
+        await cl.Message(
+            content=f"Defaulting to `{chosen.name}`. Action prompts are unavailable."
+        ).send()
+        return chosen
+
+    fields = ACTION_FIELDS or {"name", "value", "label", "payload"}
+    actions = []
+    for option in options:
+        kwargs: dict[str, Any] = {"name": option.name}
+        if "label" in fields:
+            kwargs["label"] = option.name
+        if "value" in fields:
+            kwargs["value"] = option.name
+        if not fields or "payload" in fields:
+            kwargs["payload"] = {"model": option.name}
+        actions.append(cl.Action(**kwargs))
+    message = header or "Select a trained model to use for inference:"
+    result = await cl.AskActionMessage(content=message, actions=actions).send()
+    selected_name = extract_action_value(result)
+    if not selected_name:
+        selected_name = options[0].name
+    return next((opt for opt in options if opt.name == selected_name), options[0])
+
+
+def build_config(user_input: str, model: ModelOption) -> DictConfig:
+    if GlobalHydra.instance().is_initialized():
+        GlobalHydra.instance().clear()
+    cfg = compose_model_config(model.name, job_name=model.job_name)
+    cfg.input_text = user_input
+    return cfg
+
+
+async def ensure_model_selected(header: str | None = None) -> ModelOption | None:
+    model = restore_model_from_session()
+    if model:
+        return model
+    model = await prompt_model_selection(header=header)
+    if model:
+        persist_model_selection(model)
+    return model
+
+
+async def handle_model_command(content: str) -> None:
+    parts = content.split(maxsplit=1)
+    options = discover_model_options()
+    if not options:
+        await cl.Message(
+            content="No trained models available to select. Add a model under `trained_models/` and try again."
+        ).send()
+        return
+
+    if len(parts) == 1:
+        model = await prompt_model_selection(header="Pick a new model:")
+    else:
+        requested = parts[1].strip()
+        model = next((opt for opt in options if opt.name == requested), None)
+        if not model:
+            available = ", ".join(opt.name for opt in options)
+            await cl.Message(
+                content=f"Model `{requested}` not found. Available models: {available}."
+            ).send()
+            return
+    if model:
+        persist_model_selection(model)
+        await cl.Message(
+            content=f"Switched to `{model.name}`. Provide text to generate a continuation."
+        ).send()
+
+
+@cl.on_chat_start
+async def on_chat_start() -> None:
+    model = await ensure_model_selected()
+    if not model:
+        return
+    await cl.Message(
+        content=(
+            f"Using trained model `{model.name}`. Send text for generation, or type `/model` to switch models."
+        )
+    ).send()
+
+
+@cl.on_message
+async def on_message(message: cl.Message) -> None:
+    user_input = message.content.strip()
+
+    if user_input.lower().startswith("/model"):
+        await handle_model_command(user_input)
+        return
+
+    if not user_input:
+        await cl.Message(content="Please provide some text for the model.").send()
+        return
+
+    model = await ensure_model_selected(header="Select a model before requesting generations:")
+    if not model:
+        return
+
+    cfg = build_config(user_input, model)
+    try:
+        outputs = run_inference(cfg)
+    except Exception as exc:  # pragma: no cover
+        await cl.Message(content=f"Inference error: {exc}").send()
+        return
+
+    response = outputs[0] if outputs else "No output generated."
+    await cl.Message(content=response).send()

src/mini_transformer/apps/server.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import os
+
+from fastapi import FastAPI, HTTPException
+from omegaconf import DictConfig
+
+from mini_transformer.inference import run_inference
+from mini_transformer.model_loader import (
+    CONFIG_DIR_ENV,
+    CONFIG_NAME_ENV,
+    MODEL_NAME_ENV,
+    compose_config_from_dir,
+    compose_model_config,
+    ensure_models_root,
+    list_model_names,
+)
+
+app = FastAPI(title="Mini-Transformer Inference API")
+
+
+@app.get("/healthz")
+def healthz():
+    return {"status": "ok"}
+
+
+@app.get("/generate")
+def generate(text: str | None = None):
+    config_dir = os.environ.get(CONFIG_DIR_ENV)
+    config_name = os.environ.get(CONFIG_NAME_ENV, "config_inference")
+
+    if config_dir:
+        cfg: DictConfig = compose_config_from_dir(config_dir, config_name=config_name)
+    else:
+        ensure_models_root()
+        model_name = os.environ.get(MODEL_NAME_ENV)
+        if not model_name:
+            names = list_model_names()
+            if not names:
+                raise HTTPException(
+                    status_code=503,
+                    detail=(
+                        "No models available under `trained_models/`. "
+                        "Start the server with --model or place a model folder."
+                    ),
+                )
+            model_name = names[0]
+        cfg = compose_model_config(model_name, config_name=config_name)
+
+    if text is not None and text.strip():
+        cfg.input_text = text
+    outputs = run_inference(cfg)
+    return {"outputs": outputs}

src/mini_transformer/cli.py

@@ -0,0 +1,237 @@
+"""Command-line entry points for mini_transformer."""
+
+from __future__ import annotations
+
+import argparse
+import os
+import subprocess
+import sys
+from collections.abc import Sequence
+from pathlib import Path
+
+import uvicorn
+from omegaconf import DictConfig
+
+from .apps import server as server_module
+from .inference import run_inference
+from .model_loader import (
+    CONFIG_DIR_ENV,
+    CONFIG_NAME_ENV,
+    MODEL_NAME_ENV,
+    MODELS_ENV,
+    compose_config_from_dir,
+    compose_model_config,
+    ensure_models_root,
+    list_model_names,
+    snapshot_to_local,
+)
+
+DEFAULT_CONFIG_NAME = "config_inference"
+
+__all__ = ["infer_main", "serve_main", "ui_main", "fetch_main"]
+
+
+def _select_model(model: str | None) -> str:
+    ensure_models_root()
+    names = list_model_names()
+    if not names:
+        raise SystemExit(
+            "No models found in `trained_models/`. Place a model folder with a configs/config_inference.yaml file there first."
+        )
+    if model and model in names:
+        return model
+    if model:
+        raise SystemExit(f"Model '{model}' not found. Available models: {', '.join(names)}.")
+    return names[0]
+
+
+def infer_main(argv: Sequence[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Run mini_transformer inference")
+    parser.add_argument("-t", "--input-text", help="Prompt to feed the model")
+    parser.add_argument("-m", "--model", help="Model folder under trained_models/")
+    parser.add_argument("--model-root", help="Override path to trained_models root")
+    parser.add_argument(
+        "-c", "--config-dir", help="Use Hydra configs from this directory instead of a model folder"
+    )
+    parser.add_argument(
+        "--config-name",
+        default=os.environ.get(CONFIG_NAME_ENV, DEFAULT_CONFIG_NAME),
+        help="Hydra config name to compose (default: config_inference)",
+    )
+    parser.add_argument(
+        "-o",
+        "--override",
+        action="append",
+        default=[],
+        help="Hydra override in dotlist form (e.g. generation.temperature=0.7)",
+    )
+    args = parser.parse_args(argv)
+
+    overrides = args.override or []
+    if args.model_root:
+        os.environ[MODELS_ENV] = args.model_root
+
+    if args.config_dir or os.environ.get(CONFIG_DIR_ENV):
+        config_dir = args.config_dir or os.environ.get(CONFIG_DIR_ENV)
+        if config_dir is None:
+            raise SystemExit("Config directory not provided")
+        cfg: DictConfig = compose_config_from_dir(
+            config_dir,
+            config_name=args.config_name,
+            overrides=overrides,
+        )
+    else:
+        model_name = _select_model(args.model)
+        cfg = compose_model_config(
+            model_name,
+            config_name=args.config_name,
+            overrides=overrides,
+        )
+    if args.input_text:
+        cfg.input_text = args.input_text
+
+    outputs = run_inference(cfg)
+    print("\n".join(outputs))
+    return 0
+
+
+def serve_main(argv: Sequence[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Launch the FastAPI inference server")
+    parser.add_argument("--host", default="0.0.0.0", help="Bind address (default: 0.0.0.0)")
+    parser.add_argument("--port", type=int, default=8000, help="Port for uvicorn (default: 8000)")
+    parser.add_argument("--reload", action="store_true", help="Enable auto-reload")
+    parser.add_argument("-m", "--model", help="Model folder under trained_models/")
+    parser.add_argument("--model-root", help="Override path to trained_models root")
+    parser.add_argument(
+        "-c", "--config-dir", help="Use configs from this directory instead of a model folder"
+    )
+    parser.add_argument(
+        "--config-name",
+        default=os.environ.get(CONFIG_NAME_ENV, DEFAULT_CONFIG_NAME),
+        help="Hydra config name to compose",
+    )
+    args = parser.parse_args(argv)
+
+    if args.model_root:
+        os.environ[MODELS_ENV] = args.model_root
+
+    if args.config_dir:
+        os.environ[CONFIG_DIR_ENV] = args.config_dir
+        os.environ.pop(MODEL_NAME_ENV, None)
+    else:
+        model_name = _select_model(args.model)
+        os.environ[MODEL_NAME_ENV] = model_name
+        os.environ.pop(CONFIG_DIR_ENV, None)
+
+    os.environ[CONFIG_NAME_ENV] = args.config_name
+
+    uvicorn.run(
+        "mini_transformer.apps.server:app",
+        host=args.host,
+        port=args.port,
+        reload=args.reload,
+    )
+    return 0
+
+
+def ui_main(argv: Sequence[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Launch the Chainlit chat UI")
+    parser.add_argument("--model-root", help="Override path to trained_models root")
+    parser.add_argument("-m", "--model", help="Pre-select a model when the UI starts")
+    parser.add_argument(
+        "-c", "--config-dir", help="Use configs from this directory instead of trained_models"
+    )
+    parser.add_argument(
+        "--host",
+        default=os.environ.get("MINI_TRANSFORMER_UI_HOST", "127.0.0.1"),
+        help="Host to bind Chainlit to (default: 127.0.0.1)",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=int(os.environ.get("MINI_TRANSFORMER_UI_PORT", 8000)),
+        help="Port for Chainlit (default: 8000)",
+    )
+    parser.add_argument(
+        "--config-name",
+        default=os.environ.get(CONFIG_NAME_ENV, DEFAULT_CONFIG_NAME),
+        help="Hydra config name to compose",
+    )
+    args = parser.parse_args(argv)
+
+    env = os.environ.copy()
+    if args.model_root:
+        env[MODELS_ENV] = args.model_root
+    if args.config_dir:
+        env[CONFIG_DIR_ENV] = args.config_dir
+    if args.model:
+        env[MODEL_NAME_ENV] = args.model
+    env[CONFIG_NAME_ENV] = args.config_name
+    env.setdefault("MINI_TRANSFORMER_UI_HOST", args.host)
+    env.setdefault("MINI_TRANSFORMER_UI_PORT", str(args.port))
+
+    app_path = Path(server_module.__file__).with_name("chainlit_app.py")
+    cmd = [
+        "chainlit",
+        "run",
+        str(app_path),
+        "--host",
+        args.host,
+        "--port",
+        str(args.port),
+    ]
+    try:
+        proc = subprocess.run(cmd, env=env)
+    except FileNotFoundError as exc:  # pragma: no cover
+        raise SystemExit(
+            "The 'chainlit' executable is missing. Install mini-transformer[server] to run the UI."
+        ) from exc
+    return proc.returncode
+
+
+def fetch_main(argv: Sequence[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Download a model from the Hugging Face Hub")
+    parser.add_argument(
+        "model_id", help="Repo ID on Hugging Face Hub (e.g. AlaBoussoffara/transformer_small)"
+    )
+    parser.add_argument(
+        "-n",
+        "--name",
+        help="Local directory name under trained_models/ (default derived from model id)",
+    )
+    parser.add_argument("--revision", help="Optional revision or branch to download")
+    parser.add_argument("--repo-type", default="model", help="Hub repo type (default: model)")
+    parser.add_argument("--token", help="Authentication token for private repos")
+    parser.add_argument("--cache-dir", help="Custom Hugging Face cache directory")
+    parser.add_argument("--model-root", help="Override path to trained_models root")
+    parser.add_argument("--force", action="store_true", help="Overwrite existing local directory")
+    args = parser.parse_args(argv)
+
+    if args.model_root:
+        os.environ[MODELS_ENV] = args.model_root
+
+    try:
+        destination = snapshot_to_local(
+            args.model_id,
+            local_name=args.name,
+            revision=args.revision,
+            repo_type=args.repo_type,
+            token=args.token,
+            cache_dir=args.cache_dir,
+            force=args.force,
+        )
+    except FileExistsError as exc:  # pragma: no cover
+        raise SystemExit(str(exc)) from exc
+
+    print(f"Downloaded {args.model_id} → {destination}")
+    config_path = destination / "configs" / f"{DEFAULT_CONFIG_NAME}.yaml"
+    if not config_path.exists():
+        print(
+            f"[warning] Expected config file {config_path} not found. The inference CLI requires this file.",
+            file=sys.stderr,
+        )
+    return 0
+
+
+if __name__ == "__main__":  # pragma: no cover
+    sys.exit(infer_main())

src/mini_transformer/conf/__init__.py

@@ -0,0 +1 @@
+"""Hydra configuration package included with mini_transformer."""

src/mini_transformer/conf/generation/generation.yaml

@@ -0,0 +1,9 @@
+max_new_tokens: 128
+temperature: 1.0
+top_k: null
+top_p: null
+do_sample: false
+presence_penalty: 0.0
+frequency_penalty: 0.0
+no_repeat_ngram: null
+min_steps_before_eos: 0

src/mini_transformer/conf/infer_mode.yaml

@@ -0,0 +1,11 @@
+defaults:
+  - model: small
+  - runtime: local
+  - tokenizer: bpe_8k
+  - generation: generation
+  - _self_
+input_text: ""
+hydra:
+  run:
+    dir: .
+  output_subdir: null

src/mini_transformer/conf/model/small.yaml

@@ -0,0 +1,15 @@
+name: small
+best_checkpoint_path: ${oc.env:MINI_TRANSFORMER_CHECKPOINT_BEST, ""}
+latest_checkpoint_path: ${oc.env:MINI_TRANSFORMER_CHECKPOINT_LATEST, ""}
+tokenizer: bpe_8k
+d_model: 512
+num_layers: 4
+num_heads: 8
+d_ff: 1024
+dropout_rate: 0.1
+vocab_size: 8192
+max_seq_len: 128
+pad_id: 0
+bos_id: 1
+eos_id: 2
+layer_norm_style: post

src/mini_transformer/conf/runtime/local.yaml

@@ -0,0 +1,7 @@
+seed: ${oc.env:MINI_TRANSFORMER_SEED, 42}
+device: ${oc.env:MINI_TRANSFORMER_DEVICE, "cpu"}
+output_dir: ${oc.env:MINI_TRANSFORMER_OUTPUT_DIR, "./outputs"}
+data_dir: ${oc.env:MINI_TRANSFORMER_DATA_DIR, "./data"}
+cache_dir: ${oc.env:MINI_TRANSFORMER_CACHE_DIR, "./data/cache"}
+tokenizer_dir: ${oc.env:MINI_TRANSFORMER_TOKENIZER_DIR, "./tokenizer"}
+checkpoint_path: ${oc.env:MINI_TRANSFORMER_CHECKPOINT_DIR, "./outputs/checkpoints"}

src/mini_transformer/conf/tokenizer/bpe_8k.yaml

@@ -0,0 +1,15 @@
+name: bpe_8k
+path: ${oc.env:MINI_TRANSFORMER_TOKENIZER_PATH, ""}
+corpus: ${oc.env:MINI_TRANSFORMER_CORPUS_PATH, ""}
+dataset: ${oc.env:MINI_TRANSFORMER_DATASET_NAME, ""}
+vocab_size: 8192
+max_seq_len: 128
+pad_token: <pad>
+bos_token: <bos>
+eos_token: <eos>
+unk_token: <unk>
+special_tokens: ["<pad>", "<bos>", "<eos>", "<unk>"]
+pad_id: 0
+bos_id: 1
+eos_id: 2
+unk_id: 3

src/{transformer → mini_transformer}/configs.py

@@ -34,6 +34,7 @@ class ModelCfg:
     pad_id: int
     bos_id: int
     eos_id: int
+    layer_norm_style: str | None = None
 
 
 @dataclass

src/mini_transformer/inference.py

@@ -0,0 +1,176 @@
+"""Programmatic inference helper for mini_transformer."""
+
+from __future__ import annotations
+
+import warnings
+from pathlib import Path
+
+import torch
+from hydra.utils import to_absolute_path
+from omegaconf import DictConfig, OmegaConf
+from transformers import PreTrainedTokenizerFast
+
+from .configs import InferAppCfg, ModelCfg, TokenizerCfg
+from .transformer import BasicEncoderDecoderTransformer
+from .utils import check_tokenizer_model_compatibility, set_global_seed
+
+
+def _resolve_device(device_str: str | None) -> torch.device:
+    """Pick a device honoring configs while falling back to CPU if unavailable."""
+
+    if not device_str:
+        return torch.device("cpu")
+    try:
+        device = torch.device(device_str)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"Invalid runtime.device value: {device_str!r}") from exc
+
+    if device.type == "cuda":
+        if not torch.cuda.is_available():
+            warnings.warn(
+                "CUDA device requested but torch.cuda.is_available() is False; falling back to CPU.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+            return torch.device("cpu")
+        if device.index is not None and device.index >= torch.cuda.device_count():
+            warnings.warn(
+                f"CUDA device index {device.index} out of range; using cuda:0 instead.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+            return torch.device("cuda:0")
+    return device
+
+
+_SESSION_CACHE: dict[
+    tuple[str, ...], tuple[BasicEncoderDecoderTransformer, PreTrainedTokenizerFast]
+] = {}
+
+
+def _session_cache_key(
+    model_cfg: ModelCfg,
+    tokenizer_cfg: TokenizerCfg,
+    tokenizer_path: Path,
+    checkpoint_path: Path | None,
+    device: torch.device,
+) -> tuple[str, ...]:
+    special_tokens = ",".join(tokenizer_cfg.special_tokens) if tokenizer_cfg.special_tokens else ""
+    return (
+        model_cfg.name,
+        str(tokenizer_path),
+        str(checkpoint_path) if checkpoint_path else "",
+        str(device),
+        str(model_cfg.vocab_size),
+        str(model_cfg.d_model),
+        str(model_cfg.num_heads),
+        str(model_cfg.num_layers),
+        str(model_cfg.d_ff),
+        f"{model_cfg.dropout_rate}",
+        str(model_cfg.max_seq_len),
+        str(model_cfg.pad_id),
+        str(model_cfg.bos_id),
+        str(model_cfg.eos_id),
+        tokenizer_cfg.name,
+        str(tokenizer_cfg.vocab_size),
+        str(tokenizer_cfg.max_seq_len),
+        tokenizer_cfg.pad_token or "",
+        tokenizer_cfg.bos_token or "",
+        tokenizer_cfg.eos_token or "",
+        tokenizer_cfg.unk_token or "",
+        special_tokens,
+    )
+
+
+def _get_or_create_session(
+    *,
+    model_cfg: ModelCfg,
+    tokenizer_cfg: TokenizerCfg,
+    tokenizer_path: Path,
+    checkpoint_path: Path | None,
+    device: torch.device,
+) -> tuple[BasicEncoderDecoderTransformer, PreTrainedTokenizerFast]:
+    key = _session_cache_key(model_cfg, tokenizer_cfg, tokenizer_path, checkpoint_path, device)
+    if key in _SESSION_CACHE:
+        return _SESSION_CACHE[key]
+
+    transformer = BasicEncoderDecoderTransformer(model_cfg)
+    if checkpoint_path:
+        ckpt = torch.load(checkpoint_path, map_location=device)
+        state = ckpt.get("model_state_dict", ckpt)
+        transformer.load_state_dict(state, strict=True)
+    transformer.to(device)
+    transformer.eval()
+
+    tokenizer = PreTrainedTokenizerFast(
+        tokenizer_file=str(tokenizer_path),
+        bos_token=tokenizer_cfg.bos_token,
+        eos_token=tokenizer_cfg.eos_token,
+        unk_token=tokenizer_cfg.unk_token,
+        pad_token=tokenizer_cfg.pad_token,
+        model_max_length=tokenizer_cfg.max_seq_len,
+    )
+
+    _SESSION_CACHE[key] = (transformer, tokenizer)
+    return transformer, tokenizer
+
+
+def run_inference(cfg: DictConfig) -> list[str]:
+    """Run inference using a composed Hydra configuration."""
+
+    scfg_temp = OmegaConf.merge(OmegaConf.structured(InferAppCfg), cfg)
+    scfg: InferAppCfg = OmegaConf.to_object(scfg_temp)
+    model_cfg = ModelCfg(**vars(scfg.model))
+    tokenizer_cfg = TokenizerCfg(**vars(scfg.tokenizer))
+
+    set_global_seed(scfg.runtime.seed)
+    check_tokenizer_model_compatibility(model_cfg, tokenizer_cfg)
+
+    device = _resolve_device(getattr(scfg.runtime, "device", None))
+
+    if not tokenizer_cfg.path:
+        raise FileNotFoundError(
+            "Tokenizer path is empty. Update your config or set MINI_TRANSFORMER_TOKENIZER_PATH."
+        )
+
+    tokenizer_path = Path(to_absolute_path(tokenizer_cfg.path))
+    if not tokenizer_path.is_file():
+        raise FileNotFoundError(
+            f"Tokenizer file not found: {tokenizer_path}. Check your model's config or environment variables."
+        )
+
+    checkpoint_path: Path | None = None
+    if scfg.model.best_checkpoint_path:
+        checkpoint_path = Path(to_absolute_path(scfg.model.best_checkpoint_path))
+
+    transformer, tokenizer = _get_or_create_session(
+        model_cfg=model_cfg,
+        tokenizer_cfg=tokenizer_cfg,
+        tokenizer_path=tokenizer_path,
+        checkpoint_path=checkpoint_path,
+        device=device,
+    )
+
+    text_input = scfg.input_text
+    if not text_input:
+        raise SystemExit('Pass text like: input_text="hello world"')
+
+    encoded = tokenizer(text_input, padding=True, truncation=True, return_tensors="pt")
+    src_ids = encoded["input_ids"].to(device)
+    src_padd_mask = (encoded["attention_mask"] == 0).to(device)
+
+    tgt_ids = transformer.generate(
+        src_ids,
+        src_padd_mask,
+        max_new_tokens=scfg.generation.max_new_tokens,
+        temperature=scfg.generation.temperature,
+        top_k=scfg.generation.top_k,
+        top_p=scfg.generation.top_p,
+        do_sample=scfg.generation.do_sample,
+        presence_penalty=scfg.generation.presence_penalty,
+        frequency_penalty=scfg.generation.frequency_penalty,
+        no_repeat_ngram=scfg.generation.no_repeat_ngram,
+        min_steps_before_eos=scfg.generation.min_steps_before_eos,
+        seed=scfg.runtime.seed,
+    )
+    return tokenizer.batch_decode(tgt_ids.cpu(), skip_special_tokens=True)

src/mini_transformer/model_loader.py

@@ -0,0 +1,172 @@
+"""Helpers for loading model configurations from local trained_models folders."""
+
+from __future__ import annotations
+
+import os
+import shutil
+from collections.abc import Sequence
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+from hydra import compose, initialize_config_dir
+from hydra.core.global_hydra import GlobalHydra
+from omegaconf import DictConfig
+
+MODELS_ENV = "MINI_TRANSFORMER_MODELS"
+MODEL_NAME_ENV = "MINI_TRANSFORMER_MODEL_NAME"
+CONFIG_DIR_ENV = "MINI_TRANSFORMER_CONFIG_DIR"
+CONFIG_NAME_ENV = "MINI_TRANSFORMER_CONFIG_NAME"
+DEFAULT_CONFIG_NAME = "config_inference"
+CONFIG_FILENAME = f"{DEFAULT_CONFIG_NAME}.yaml"
+
+
+def get_models_root() -> Path:
+    """Return the directory where local models are stored."""
+
+    env_root = os.environ.get(MODELS_ENV)
+    if env_root:
+        return Path(env_root).expanduser().resolve()
+    return (Path.cwd() / "trained_models").resolve()
+
+
+def ensure_models_root() -> Path:
+    root = get_models_root()
+    root.mkdir(parents=True, exist_ok=True)
+    return root
+
+
+def list_model_names() -> list[str]:
+    root = ensure_models_root()
+    names: list[str] = []
+    for entry in sorted(root.iterdir()):
+        if entry.is_dir() and (entry / "configs" / CONFIG_FILENAME).exists():
+            names.append(entry.name)
+    return names
+
+
+def _resolve_relative_paths(cfg: DictConfig, base_dir: Path) -> None:
+    def _resolve(value: str | None) -> str | None:
+        if not value:
+            return value
+        path = Path(value)
+        if path.is_absolute():
+            return str(path)
+        parts = path.parts
+        if len(parts) >= 2 and parts[0] == "trained_models" and parts[1] == base_dir.name:
+            path = Path(*parts[2:])
+        elif parts and parts[0] == base_dir.name:
+            path = Path(*parts[1:])
+        return str((base_dir / path).resolve())
+
+    if "model" in cfg:
+        for key in ("best_checkpoint_path", "latest_checkpoint_path"):
+            if key in cfg.model:
+                cfg.model[key] = _resolve(cfg.model[key])
+    if "tokenizer" in cfg and "path" in cfg.tokenizer:
+        cfg.tokenizer["path"] = _resolve(cfg.tokenizer["path"])
+    if "runtime" in cfg:
+        for key in ("output_dir", "data_dir", "cache_dir", "tokenizer_dir", "checkpoint_path"):
+            if key in cfg.runtime:
+                cfg.runtime[key] = _resolve(cfg.runtime[key])
+
+
+def compose_model_config(
+    model_name: str,
+    *,
+    config_name: str = DEFAULT_CONFIG_NAME,
+    overrides: Sequence[str] | None = None,
+    job_name: str | None = None,
+) -> DictConfig:
+    """Compose the Hydra config for a given model folder."""
+
+    models_root = ensure_models_root()
+    model_dir = models_root / model_name
+    config_dir = model_dir / "configs"
+    if not config_dir.is_dir():
+        raise FileNotFoundError(
+            f"Config directory not found for model '{model_name}' at {config_dir}."
+        )
+
+    overrides = list(overrides or [])
+    job = job_name or f"mini_transformer_{model_name}"
+    if GlobalHydra.instance().is_initialized():
+        GlobalHydra.instance().clear()
+    with initialize_config_dir(config_dir=str(config_dir), job_name=job, version_base=None):
+        cfg = compose(config_name=config_name, overrides=overrides)
+    _resolve_relative_paths(cfg, model_dir)
+    return cfg
+
+
+def compose_config_from_dir(
+    config_dir: str,
+    *,
+    config_name: str = DEFAULT_CONFIG_NAME,
+    overrides: Sequence[str] | None = None,
+    job_name: str = "mini_transformer_cli",
+) -> DictConfig:
+    overrides = list(overrides or [])
+    config_path = Path(config_dir).resolve()
+    base_dir = config_path.parent if config_path.name == "configs" else config_path
+    if GlobalHydra.instance().is_initialized():
+        GlobalHydra.instance().clear()
+    with initialize_config_dir(config_dir=str(config_path), job_name=job_name, version_base=None):
+        cfg = compose(config_name=config_name, overrides=overrides)
+    _resolve_relative_paths(cfg, base_dir)
+    return cfg
+
+
+__all__ = [
+    "CONFIG_FILENAME",
+    "DEFAULT_CONFIG_NAME",
+    "MODELS_ENV",
+    "MODEL_NAME_ENV",
+    "CONFIG_DIR_ENV",
+    "CONFIG_NAME_ENV",
+    "compose_model_config",
+    "compose_config_from_dir",
+    "ensure_models_root",
+    "get_models_root",
+    "list_model_names",
+    "snapshot_to_local",
+]
+
+
+def snapshot_to_local(
+    model_id: str,
+    *,
+    local_name: str | None = None,
+    revision: str | None = None,
+    repo_type: str = "model",
+    token: str | None = None,
+    cache_dir: str | None = None,
+    force: bool = False,
+) -> Path:
+    """Download a repository from the Hugging Face Hub into `trained_models/`."""
+
+    root = ensure_models_root()
+    name = local_name or model_id.replace("/", "__")
+    destination = root / name
+
+    if destination.exists() and any(destination.iterdir()):
+        if not force:
+            raise FileExistsError(
+                f"Destination {destination} already exists. Use --force to overwrite."
+            )
+        for child in destination.iterdir():
+            if child.is_dir():
+                shutil.rmtree(child)
+            else:
+                child.unlink()
+
+    destination.mkdir(parents=True, exist_ok=True)
+
+    snapshot_download(
+        repo_id=model_id,
+        repo_type=repo_type,
+        local_dir=str(destination),
+        local_dir_use_symlinks=False,
+        revision=revision,
+        token=token,
+        cache_dir=cache_dir,
+    )
+    return destination

src/mini_transformer/modules/__init__.py

@@ -0,0 +1,6 @@
+from .decoder import TransformerDecoder
+from .embedding import InputEmbedding
+from .encoder import TransformerEncoder
+from .lm_head import LMHead
+
+__all__ = ["TransformerEncoder", "TransformerDecoder", "InputEmbedding", "LMHead"]

src/{transformer → mini_transformer}/modules/attention.py

@@ -8,7 +8,7 @@ import torch
 import torch.nn as nn
 from torch import Tensor
 
-from transformer.utils import (
+from ..utils import (
     calculate_attention,
     combine_masks,
     create_qk_padding_mask,

src/{transformer → mini_transformer}/modules/decoder.py

@@ -7,8 +7,8 @@ import torch.nn as nn
 import torch.utils.checkpoint as ckpt
 from torch import Tensor
 
-from transformer.modules.attention import MultiHeadAttention
-from transformer.modules.feedforward import FeedForwardLayer
+from .attention import MultiHeadAttention
+from .feedforward import FeedForwardLayer
 
 __all__ = ["DecoderLayer", "TransformerDecoder"]
 
@@ -16,7 +16,15 @@ __all__ = ["DecoderLayer", "TransformerDecoder"]
 class DecoderLayer(nn.Module):
     """Single decoder block with self-attention, cross-attention, and feed-forward."""
 
-    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout_rate: float) -> None:
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        d_ff: int,
+        dropout_rate: float,
+        *,
+        layer_norm_style: str = "post",
+    ) -> None:
         super().__init__()
 
         if not isinstance(d_model, int):
@@ -27,6 +35,8 @@ class DecoderLayer(nn.Module):
             raise TypeError(f"d_ff must be an int, got {type(d_ff)}")
         if not isinstance(dropout_rate, float):
             raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
+        if not isinstance(layer_norm_style, str):
+            raise TypeError(f"layer_norm_style must be a string, got {type(layer_norm_style)}")
 
         if d_model <= 0:
             raise ValueError("d_model must be strictly greater than 0")
@@ -37,6 +47,12 @@ class DecoderLayer(nn.Module):
         if not 0.0 <= dropout_rate < 1.0:
             raise ValueError("dropout_rate must be in [0, 1)")
 
+        style = layer_norm_style.lower()
+        if style not in {"pre", "post"}:
+            raise ValueError("layer_norm_style must be either 'pre' or 'post' (case-insensitive)")
+        self.layer_norm_style = style
+        self.pre_norm = style == "pre"
+
         self.self_attention_layer = MultiHeadAttention(d_model, num_heads, dropout_rate)
         self.cross_attention_layer = MultiHeadAttention(d_model, num_heads, dropout_rate)
         self.feed_forward = FeedForwardLayer(d_model, d_ff, dropout_rate)
@@ -56,7 +72,7 @@ class DecoderLayer(nn.Module):
         tgt_padding_mask: Tensor,
         tgt_causal_mask: Tensor | None,
     ) -> Tensor:
-        """Run one decoder layer using post-layer-normalisation formulation."""
+        """Run one decoder layer using the configured layer-normalisation style."""
 
         if not isinstance(x, torch.Tensor):
             raise TypeError("x must be a torch.Tensor")
@@ -85,7 +101,33 @@ class DecoderLayer(nn.Module):
                     f" got dtype {mask.dtype} and shape {tuple(mask.shape)}"
                 )
 
-        # Self-attention
+        if self.pre_norm:
+            normed = self.norm1(y)
+            self_attn_out = self.self_attention_layer(
+                normed,
+                normed,
+                normed,
+                tgt_padding_mask,
+                tgt_padding_mask,
+                tgt_causal_mask,
+            )
+            y = y + self.dropout1(self_attn_out)
+
+            normed_y = self.norm2(y)
+            cross_attn_out = self.cross_attention_layer(
+                normed_y,
+                x,
+                x,
+                tgt_padding_mask,
+                src_padding_mask,
+            )
+            y = y + self.dropout2(cross_attn_out)
+
+            ff_out = self.feed_forward(self.norm3(y))
+            y = y + self.dropout3(ff_out)
+            return y
+
+        # Self-attention (post-LN)
         self_attn_out = self.self_attention_layer(
             y,
             y,
@@ -121,6 +163,8 @@ class TransformerDecoder(nn.Module):
         d_ff: int,
         num_layers: int,
         dropout_rate: float,
+        *,
+        layer_norm_style: str = "post",
     ) -> None:
         super().__init__()
 
@@ -129,8 +173,24 @@ class TransformerDecoder(nn.Module):
         if num_layers <= 0:
             raise ValueError("num_layers must be strictly greater than 0")
 
+        if not isinstance(layer_norm_style, str):
+            raise TypeError(f"layer_norm_style must be a string, got {type(layer_norm_style)}")
+        style = layer_norm_style.lower()
+        if style not in {"pre", "post"}:
+            raise ValueError("layer_norm_style must be either 'pre' or 'post' (case-insensitive)")
+
+        self.layer_norm_style = style
         self.layers = nn.ModuleList(
-            [DecoderLayer(d_model, num_heads, d_ff, dropout_rate) for _ in range(num_layers)]
+            [
+                DecoderLayer(
+                    d_model,
+                    num_heads,
+                    d_ff,
+                    dropout_rate,
+                    layer_norm_style=style,
+                )
+                for _ in range(num_layers)
+            ]
         )
         self.use_ckpt = False
 

src/{transformer → mini_transformer}/modules/embedding.py

@@ -8,7 +8,7 @@ import torch
 import torch.nn as nn
 from torch import Tensor
 
-from transformer.utils import sinusoidal_positional_encoding
+from ..utils import sinusoidal_positional_encoding
 
 __all__ = ["InputEmbedding", "PositionalEmbedding", "LearnedPositionalEmbedding"]
 

src/{transformer → mini_transformer}/modules/encoder.py

@@ -7,16 +7,24 @@ import torch.nn as nn
 import torch.utils.checkpoint as ckpt
 from torch import Tensor
 
-from transformer.modules.attention import MultiHeadAttention
-from transformer.modules.feedforward import FeedForwardLayer
+from .attention import MultiHeadAttention
+from .feedforward import FeedForwardLayer
 
 __all__ = ["EncoderLayer", "TransformerEncoder"]
 
 
 class EncoderLayer(nn.Module):
-    """Self-attention + feed-forward block with post-layer-normalisation."""
+    """Self-attention + feed-forward block supporting pre/post layer normalisation."""
 
-    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout_rate: float) -> None:
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        d_ff: int,
+        dropout_rate: float,
+        *,
+        layer_norm_style: str = "post",
+    ) -> None:
         super().__init__()
 
         if not isinstance(d_model, int):
@@ -27,6 +35,8 @@ class EncoderLayer(nn.Module):
             raise TypeError(f"d_ff must be an int, got {type(d_ff)}")
         if not isinstance(dropout_rate, float):
             raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
+        if not isinstance(layer_norm_style, str):
+            raise TypeError(f"layer_norm_style must be a string, got {type(layer_norm_style)}")
 
         if d_model <= 0:
             raise ValueError("d_model must be strictly greater than 0")
@@ -37,6 +47,12 @@ class EncoderLayer(nn.Module):
         if not 0.0 <= dropout_rate < 1.0:
             raise ValueError("dropout_rate must be in [0, 1)")
 
+        style = layer_norm_style.lower()
+        if style not in {"pre", "post"}:
+            raise ValueError("layer_norm_style must be either 'pre' or 'post' (case-insensitive)")
+        self.layer_norm_style = style
+        self.pre_norm = style == "pre"
+
         self.attention_layer = MultiHeadAttention(d_model, num_heads, dropout_rate)
         self.feed_forward = FeedForwardLayer(d_model, d_ff, dropout_rate)
         self.norm1 = nn.LayerNorm(d_model)
@@ -59,6 +75,16 @@ class EncoderLayer(nn.Module):
                 f" got dtype {src_padding_mask.dtype} and shape {tuple(src_padding_mask.shape)}"
             )
 
+        if self.pre_norm:
+            normed = self.norm1(x)
+            attn_out = self.attention_layer(
+                normed, normed, normed, src_padding_mask, src_padding_mask
+            )
+            x = x + self.dropout1(attn_out)
+            ff_out = self.feed_forward(self.norm2(x))
+            x = x + self.dropout2(ff_out)
+            return x
+
         attn_out = self.attention_layer(x, x, x, src_padding_mask, src_padding_mask)
         x = self.norm1(x + self.dropout1(attn_out))
         ff_out = self.feed_forward(x)
@@ -75,6 +101,8 @@ class TransformerEncoder(nn.Module):
         d_ff: int,
         num_layers: int,
         dropout_rate: float,
+        *,
+        layer_norm_style: str = "post",
     ) -> None:
         super().__init__()
 
@@ -83,8 +111,24 @@ class TransformerEncoder(nn.Module):
         if num_layers <= 0:
             raise ValueError("num_layers must be strictly greater than 0")
 
+        if not isinstance(layer_norm_style, str):
+            raise TypeError(f"layer_norm_style must be a string, got {type(layer_norm_style)}")
+        style = layer_norm_style.lower()
+        if style not in {"pre", "post"}:
+            raise ValueError("layer_norm_style must be either 'pre' or 'post' (case-insensitive)")
+
+        self.layer_norm_style = style
         self.layers = nn.ModuleList(
-            [EncoderLayer(d_model, num_heads, d_ff, dropout_rate) for _ in range(num_layers)]
+            [
+                EncoderLayer(
+                    d_model,
+                    num_heads,
+                    d_ff,
+                    dropout_rate,
+                    layer_norm_style=style,
+                )
+                for _ in range(num_layers)
+            ]
         )
         self.use_ckpt = False
 

src/{transformer → mini_transformer}/transformer.py

@@ -9,9 +9,9 @@ import torch
 import torch.nn as nn
 from torch import Tensor
 
-from transformer import modules
-from transformer.configs import ModelCfg
-from transformer.utils import broadcast_padding_mask, create_causal_mask, sample_from_logits
+from . import modules
+from .configs import ModelCfg
+from .utils import broadcast_padding_mask, create_causal_mask, sample_from_logits
 
 __all__ = ["BasicEncoderDecoderTransformer"]
 
@@ -57,14 +57,41 @@ class BasicEncoderDecoderTransformer(nn.Module):
 
         self.cfg = cfg
 
+        style_raw = getattr(cfg, "layer_norm_style", None)
+        if style_raw is None:
+            style = "post" if cfg.num_layers <= 4 else "pre"
+        elif not isinstance(style_raw, str):
+            raise TypeError(
+                f"layer_norm_style must be a string when provided, got {type(style_raw)}"
+            )
+        else:
+            style = style_raw.lower()
+        if style not in {"pre", "post"}:
+            raise ValueError(
+                "layer_norm_style must be either 'pre' or 'post'"
+                f" (case-insensitive); got {style_raw!r}"
+            )
+        self.layer_norm_style = style
+        cfg.layer_norm_style = style
+
         self.embed = modules.InputEmbedding(
             cfg.vocab_size, cfg.d_model, cfg.max_seq_len, cfg.pad_id, cfg.dropout_rate
         )
         self.encoder = modules.TransformerEncoder(
-            cfg.d_model, cfg.num_heads, cfg.d_ff, cfg.num_layers, cfg.dropout_rate
+            cfg.d_model,
+            cfg.num_heads,
+            cfg.d_ff,
+            cfg.num_layers,
+            cfg.dropout_rate,
+            layer_norm_style=style,
         )
         self.decoder = modules.TransformerDecoder(
-            cfg.d_model, cfg.num_heads, cfg.d_ff, cfg.num_layers, cfg.dropout_rate
+            cfg.d_model,
+            cfg.num_heads,
+            cfg.d_ff,
+            cfg.num_layers,
+            cfg.dropout_rate,
+            layer_norm_style=style,
         )
 
         self.lm_head = modules.LMHead(cfg.d_model, cfg.vocab_size)
@@ -306,7 +333,7 @@ class BasicEncoderDecoderTransformer(nn.Module):
                 rng = _make_generator()
             for step in range(max_new_tokens):
                 # decode on current tgt_ids; take last-step logits
-                tgt_padding_mask = tgt_ids != self.pad_id
+                tgt_padding_mask = tgt_ids == self.pad_id
                 tgt_padding_mask = broadcast_padding_mask(tgt_padding_mask, self.cfg.num_heads)
                 step_hidden = self.decode(
                     hidden_states, tgt_ids, src_padding_mask, tgt_padding_mask

src/{transformer → mini_transformer}/utils.py

@@ -6,12 +6,20 @@ from collections.abc import Callable, Iterable, Sequence
 from pathlib import Path
 from typing import Any
 
-import matplotlib.pyplot as plt
 import numpy as np
 import torch
 import torch.nn.functional as F
 from torch import Tensor
 
+
+def _require_matplotlib() -> Any:
+    try:
+        import matplotlib.pyplot as plt  # type: ignore
+    except ImportError as exc:  # pragma: no cover - optional dependency
+        raise ImportError("Install mini-transformer[viz] to enable plotting utilities.") from exc
+    return plt
+
+
 # -----------------------------------------------------------------------------
 # Global seeding utilities
 # -----------------------------------------------------------------------------
@@ -564,6 +572,8 @@ def plot_layer_heads_grid(
 ):
     """Plot a single layer as a 3-row (enc self / dec self / dec cross) grid."""
 
+    plt = _require_matplotlib()
+
     enc_layers = maps_dict["enc_self"]
     dec_self_layers = maps_dict["dec_self"]
     dec_cross_layers = maps_dict["dec_cross"]
@@ -628,7 +638,8 @@ def plot_all_layers_all_heads(
 ):
     """Render attention grids for every layer (optionally saving a multi-page PDF)."""
 
-    from matplotlib.backends.backend_pdf import PdfPages
+    plt = _require_matplotlib()
+    from matplotlib.backends.backend_pdf import PdfPages  # type: ignore
 
     enc_layers = len(maps_dict["enc_self"])
     dec_layers = len(maps_dict["dec_self"])
@@ -877,6 +888,7 @@ def debug_transformer_forward(
     }
 
     if attention_requested:
+        plt = _require_matplotlib()
         if not hasattr(model, "embed"):
             raise AttributeError("model must expose an embed() method for debugging")
         with torch.no_grad():

src/transformer/__init__.py

@@ -1,3 +0,0 @@
-from .transformer import BasicEncoderDecoderTransformer
-
-__all__ = ["BasicEncoderDecoderTransformer"]

src/transformer/modules/__init__.py

@@ -1,6 +0,0 @@
-from transformer.modules.decoder import TransformerDecoder
-from transformer.modules.embedding import InputEmbedding
-from transformer.modules.encoder import TransformerEncoder
-from transformer.modules.lm_head import LMHead
-
-__all__ = ["TransformerEncoder", "TransformerDecoder", "InputEmbedding", "LMHead"]

tests/units/modules/test_attention.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from transformer.modules.attention import MultiHeadAttention
+from mini_transformer.modules.attention import MultiHeadAttention
 
 
 # -----------------------

tests/units/modules/test_decoder.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from transformer.modules.decoder import DecoderLayer, TransformerDecoder
+from mini_transformer.modules.decoder import DecoderLayer, TransformerDecoder
 
 
 def _dev():
@@ -24,6 +24,8 @@ def test_decoder_layer_param_checks():
         DecoderLayer(32, 4, "64", 0.1)
     with pytest.raises(TypeError):
         DecoderLayer(32, 4, 64, "0.1")
+    with pytest.raises(TypeError):
+        DecoderLayer(32, 4, 64, 0.1, layer_norm_style=123)
 
     with pytest.raises(ValueError):
         DecoderLayer(0, 4, 64, 0.1)
@@ -33,6 +35,8 @@ def test_decoder_layer_param_checks():
         DecoderLayer(32, 4, 0, 0.1)
     with pytest.raises(ValueError):
         DecoderLayer(32, 4, 64, 1.0)  # upper bound excluded
+    with pytest.raises(ValueError):
+        DecoderLayer(32, 4, 64, 0.1, layer_norm_style="weird")
 
 
 def test_transformer_decoder_layers_count_checks():
@@ -64,6 +68,25 @@ def test_decoder_forward_happy_path(B, Sx, Sy, D, H, FF, L):
     assert out.device == device
 
 
+def test_transformer_decoder_pre_norm_layers_flag():
+    dec = TransformerDecoder(24, 3, 48, 2, 0.1, layer_norm_style="pre")
+    assert all(layer.pre_norm for layer in dec.layers)
+
+
+def test_decoder_layer_pre_norm_forward_shapes():
+    device = _dev()
+    layer = DecoderLayer(24, 3, 48, 0.1, layer_norm_style="pre").to(device)
+    x = torch.randn(2, 5, 24, device=device)
+    y = torch.randn(2, 6, 24, device=device)
+    heads = layer.self_attention_layer.num_heads
+    src_mask = torch.zeros(2, heads, 1, 5, dtype=torch.bool, device=device)
+    tgt_mask = torch.zeros(2, heads, 1, 6, dtype=torch.bool, device=device)
+    causal = torch.ones(2, heads, 6, 6, dtype=torch.bool, device=device).triu(1)
+    out = layer(x, y, src_mask, tgt_mask, causal)
+    assert out.shape == y.shape
+    assert layer.pre_norm is True
+
+
 def test_decoder_forward_input_checks_and_message_format():
     layer = DecoderLayer(24, 3, 48, 0.1)
     heads = layer.self_attention_layer.num_heads

tests/units/modules/test_embedding.py

@@ -3,7 +3,7 @@ import math
 import pytest
 import torch
 
-from transformer.modules.embedding import InputEmbedding, PositionalEmbedding
+from mini_transformer.modules.embedding import InputEmbedding, PositionalEmbedding
 
 
 # -----------------------