diff --git a/.gitignore b/.gitignore index f33e8ab9a7c0ab792da1023e2a6f3e342f230ece..401fde1974a978fec4dddf45cc933da6ef35faae 100644 --- a/.gitignore +++ b/.gitignore @@ -54,6 +54,9 @@ data/ outputs/ runs/ artifacts/ +trained_models/ +!trained_models/README.md +.chainlit/ *.pt *.pth *.ckpt diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000000000000000000000000000000000000..0a9de7ebc9e8570bcd55dee7b7e8e4fc9325adbb --- /dev/null +++ b/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. diff --git a/Dockerfile.dev b/Dockerfile.dev index febe389bfeea0d35d0c6f07e0db0567f7e7b8318..5e1f3ff7328939f010cb05f3cadeb3f4ddc6af61 100644 --- a/Dockerfile.dev +++ b/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"] diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..f17299efac9d3537eebaf60f02dd2bacec8cafbc --- /dev/null +++ b/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. diff --git a/Makefile b/Makefile index 24ab064544f6111a6eb935ff4cdda1a88daf55ca..a430bcc7f31c9f697f31929e9045473480335ceb 100644 --- a/Makefile +++ b/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 diff --git a/README.md b/README.md index dc863d7198c3f06df3a7ec3b89bdab9fd5d7bc68..5c071ff114bb8a7b9db6e86767f607ba8dc7ab12 100755 --- a/README.md +++ b/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 `__` (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//` 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//`: +``` +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. diff --git a/architecture.txt b/architecture.txt deleted file mode 100755 index 1c0c61c25e036d8cd5b4cb1416dbf500e1e40700..0000000000000000000000000000000000000000 --- a/architecture.txt +++ /dev/null @@ -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 diff --git a/chainlit.md b/chainlit.md new file mode 100644 index 0000000000000000000000000000000000000000..4507ac4676a6387c4b52a0d1111e94753a102b32 --- /dev/null +++ b/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. diff --git a/configs/generation/generation.yaml b/configs/generation/generation.yaml index ae02e30ebcc3e253dc60d3baa40cc747806abc07..239928c2471d9c2d4caa425ca70f7bb05e58e9ae 100644 --- a/configs/generation/generation.yaml +++ b/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 diff --git a/configs/model/large_model.yaml b/configs/model/large_model.yaml index 57f43fdeb73537b30cd6209efeb07dd8f827ff5c..d64d2d3a5e0206b562916f21bbceacb930b63414 100644 --- a/configs/model/large_model.yaml +++ b/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 diff --git a/configs/model/medium_model.yaml b/configs/model/medium_model.yaml index d38be493a14cf6fb9e8e24f19deb69ecd3e2091a..df3d64c563e5f1d7a896ba4c95a2a7150dc066a9 100644 --- a/configs/model/medium_model.yaml +++ b/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 diff --git a/configs/model/small_model.yaml b/configs/model/small_model.yaml index 8938d5b8522cbdbbdeb23c5bdfbc34342eeadb68..d4e7a6db1db4ad50ac3d97b9cbc015822861072d 100644 --- a/configs/model/small_model.yaml +++ b/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 diff --git a/configs/model/test_model.yaml b/configs/model/test_model.yaml index 5582f8913e841560ba8423a2b9d48e9f076361dd..80e615a909b35670c45f0adb0302031f25a36844 100644 --- a/configs/model/test_model.yaml +++ b/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 diff --git a/docker-compose.yml b/docker-compose.yml index 118fb1a75f116591b5638b9c62790d3ea9f5cffa..027772d2fbf0c00c2913b8454b93c01c5dcc515a 100644 --- a/docker-compose.yml +++ b/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" diff --git a/environment.yml b/environment.yml index eacb0884ce4d9f482e5e951521533543dd321c7a..71e72f83308103d0abd8d8df2f4f012efa22cbb2 100644 --- a/environment.yml +++ b/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 diff --git a/scripts/load_data.ipynb b/notebooks/load_data.ipynb similarity index 97% rename from scripts/load_data.ipynb rename to notebooks/load_data.ipynb index f815f041ed913fdefd760424ed52a4c4a19d2f72..8851bf53ba0260b1ff81c45572afcff601c1eb6f 100644 --- a/scripts/load_data.ipynb +++ b/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", diff --git a/scripts/tokenizer.ipynb b/notebooks/tokenizer.ipynb similarity index 94% rename from scripts/tokenizer.ipynb rename to notebooks/tokenizer.ipynb index bbab23877a523bc938a803de8a7d8a51660a21a3..6793c73b26b69688d01e291ab0166f874793a7d9 100644 --- a/scripts/tokenizer.ipynb +++ b/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" ] }, { diff --git a/scripts/train.ipynb b/notebooks/train.ipynb similarity index 99% rename from scripts/train.ipynb rename to notebooks/train.ipynb index 36ceceac77697ca5e926da5aba6b901a7a296342..90103c3bb012e378a2b877ac2909ae2c4e1ba679 100644 --- a/scripts/train.ipynb +++ b/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", diff --git a/pyproject.toml b/pyproject.toml index 7c92ab194863d9c533757a98d68a22b24daec40a..93be9b9e3cfb566b244c3f265e5183baab32adaf 100755 --- a/pyproject.toml +++ b/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 diff --git a/requirements.txt b/requirements.txt deleted file mode 100755 index 3593391e48afdb14ba280bc72003d13ba476e976..0000000000000000000000000000000000000000 Binary files a/requirements.txt and /dev/null differ diff --git a/scripts/inference.py b/scripts/inference.py deleted file mode 100644 index 6993f959d629b8314a03c98c9541aac4fa1fe106..0000000000000000000000000000000000000000 --- a/scripts/inference.py +++ /dev/null @@ -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() diff --git a/scripts/server.py b/scripts/server.py deleted file mode 100644 index ce4db0db2feec74158fe038a0aa6d4ba7f431a59..0000000000000000000000000000000000000000 --- a/scripts/server.py +++ /dev/null @@ -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} diff --git a/src/mini_transformer/__init__.py b/src/mini_transformer/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..226309d300f18e53fad146b5b391bd57cb90614c --- /dev/null +++ b/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__"] diff --git a/src/mini_transformer/apps/__init__.py b/src/mini_transformer/apps/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e64de3960cfa5d17738d73acf1bb1ea2543a9a7a --- /dev/null +++ b/src/mini_transformer/apps/__init__.py @@ -0,0 +1,3 @@ +"""Interactive applications (Chainlit UI & FastAPI server).""" + +__all__ = ["chainlit_app", "server"] diff --git a/src/mini_transformer/apps/chainlit_app.py b/src/mini_transformer/apps/chainlit_app.py new file mode 100644 index 0000000000000000000000000000000000000000..2d4222531abfb7a2e60ae392a801fe4b42105bf4 --- /dev/null +++ b/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() diff --git a/src/mini_transformer/apps/server.py b/src/mini_transformer/apps/server.py new file mode 100644 index 0000000000000000000000000000000000000000..0ac5dfdbce6bd34e1ea26c75398ab420c04ede73 --- /dev/null +++ b/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} diff --git a/src/mini_transformer/cli.py b/src/mini_transformer/cli.py new file mode 100644 index 0000000000000000000000000000000000000000..5ea1b91be38ec6bfabec0c959eaea7cb0a460e2f --- /dev/null +++ b/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()) diff --git a/src/mini_transformer/conf/__init__.py b/src/mini_transformer/conf/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f46efa6345b60462827003c4d4c6b16b7f876201 --- /dev/null +++ b/src/mini_transformer/conf/__init__.py @@ -0,0 +1 @@ +"""Hydra configuration package included with mini_transformer.""" diff --git a/src/mini_transformer/conf/generation/generation.yaml b/src/mini_transformer/conf/generation/generation.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7625bc9c7379ab1ec962c015d49f1ca893a47ff1 --- /dev/null +++ b/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 diff --git a/src/mini_transformer/conf/infer_mode.yaml b/src/mini_transformer/conf/infer_mode.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c27fc00670d6221a5acb14852f943a31c721b554 --- /dev/null +++ b/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 diff --git a/src/mini_transformer/conf/model/small.yaml b/src/mini_transformer/conf/model/small.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7247e78d42920328608ae5a56701d86c17dc629d --- /dev/null +++ b/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 diff --git a/src/mini_transformer/conf/runtime/local.yaml b/src/mini_transformer/conf/runtime/local.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f726b741393bbd283218987a4def98ce9e1802e6 --- /dev/null +++ b/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"} diff --git a/src/mini_transformer/conf/tokenizer/bpe_8k.yaml b/src/mini_transformer/conf/tokenizer/bpe_8k.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0c94a05bccb30b3541e938e5a230942a32c802d2 --- /dev/null +++ b/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: +bos_token: +eos_token: +unk_token: +special_tokens: ["", "", "", ""] +pad_id: 0 +bos_id: 1 +eos_id: 2 +unk_id: 3 diff --git a/src/transformer/configs.py b/src/mini_transformer/configs.py similarity index 98% rename from src/transformer/configs.py rename to src/mini_transformer/configs.py index 419df268b83751388f57482328ac666ffc47af9d..c1da6e2f69eaf917a8e3b60e3d3d1b7a2c65ef33 100644 --- a/src/transformer/configs.py +++ b/src/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 diff --git a/src/mini_transformer/inference.py b/src/mini_transformer/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..c98473ac4a8e56b20cd869d9169f0e3ada67c5f0 --- /dev/null +++ b/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) diff --git a/src/mini_transformer/model_loader.py b/src/mini_transformer/model_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..e7b8915e32433c345e0fe764cd586842e2012b8e --- /dev/null +++ b/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 diff --git a/src/mini_transformer/modules/__init__.py b/src/mini_transformer/modules/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..1fcfbdf6141eae86c03ad44fac066ded29fe7b81 --- /dev/null +++ b/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"] diff --git a/src/transformer/modules/attention.py b/src/mini_transformer/modules/attention.py similarity index 99% rename from src/transformer/modules/attention.py rename to src/mini_transformer/modules/attention.py index 7e8af8f20cbae80840f708e33e5128208ff2a2f8..dd67b4cb8e6093c72391ac20c414e479d25f6a5f 100755 --- a/src/transformer/modules/attention.py +++ b/src/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, diff --git a/src/transformer/modules/decoder.py b/src/mini_transformer/modules/decoder.py similarity index 72% rename from src/transformer/modules/decoder.py rename to src/mini_transformer/modules/decoder.py index 435c08f07466c377b3409b94c8974e24ded22702..52d7bd5083fe1bbe72808d574509334347896816 100755 --- a/src/transformer/modules/decoder.py +++ b/src/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 diff --git a/src/transformer/modules/embedding.py b/src/mini_transformer/modules/embedding.py similarity index 98% rename from src/transformer/modules/embedding.py rename to src/mini_transformer/modules/embedding.py index 59c5b7ff784b628e7c40680f462de0bf2f09bfbf..c3ed221bf25c978fc19f7ec36dd760f6d0ea72c4 100755 --- a/src/transformer/modules/embedding.py +++ b/src/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"] diff --git a/src/transformer/modules/encoder.py b/src/mini_transformer/modules/encoder.py similarity index 67% rename from src/transformer/modules/encoder.py rename to src/mini_transformer/modules/encoder.py index 00317faafd468b6597d53a16a5e77c3fb1865713..821a7aa13f2d2e22445515105bc55dc5f9ec9344 100755 --- a/src/transformer/modules/encoder.py +++ b/src/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 diff --git a/src/transformer/modules/feedforward.py b/src/mini_transformer/modules/feedforward.py similarity index 100% rename from src/transformer/modules/feedforward.py rename to src/mini_transformer/modules/feedforward.py diff --git a/src/transformer/modules/lm_head.py b/src/mini_transformer/modules/lm_head.py similarity index 100% rename from src/transformer/modules/lm_head.py rename to src/mini_transformer/modules/lm_head.py diff --git a/src/transformer/transformer.py b/src/mini_transformer/transformer.py similarity index 94% rename from src/transformer/transformer.py rename to src/mini_transformer/transformer.py index 2b6759ea6bf1a306c4d8307054310c6c0b9f6a15..ca1f6a55a584bdda246e497d5b54c0f3ed3960a1 100755 --- a/src/transformer/transformer.py +++ b/src/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 diff --git a/src/transformer/utils.py b/src/mini_transformer/utils.py similarity index 98% rename from src/transformer/utils.py rename to src/mini_transformer/utils.py index 5bc96a7804036ecf4da5c2e8973695fabae4e70e..ece1606e78f193a5a85a701f99bd77eab0ec9d5c 100755 --- a/src/transformer/utils.py +++ b/src/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(): diff --git a/src/transformer/__init__.py b/src/transformer/__init__.py deleted file mode 100755 index 33e5981a9c7a16147f75c7c0a66a0ec616f6a195..0000000000000000000000000000000000000000 --- a/src/transformer/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .transformer import BasicEncoderDecoderTransformer - -__all__ = ["BasicEncoderDecoderTransformer"] diff --git a/src/transformer/modules/__init__.py b/src/transformer/modules/__init__.py deleted file mode 100755 index 311f78156f42ab2b5d73b948af9ce53ce71017f0..0000000000000000000000000000000000000000 --- a/src/transformer/modules/__init__.py +++ /dev/null @@ -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"] diff --git a/tests/units/modules/test_attention.py b/tests/units/modules/test_attention.py index 6a05be3a8a90c8e59afd58c494eb91bea1823369..fda6bfe4db44090fd8c094728bb7363d0599e55e 100644 --- a/tests/units/modules/test_attention.py +++ b/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 # ----------------------- diff --git a/tests/units/modules/test_decoder.py b/tests/units/modules/test_decoder.py index 2052fa3d190a602de6bf7b637c242d4b9addbbe2..834815f37fdd2584df010975a1ddb2f89c236abe 100644 --- a/tests/units/modules/test_decoder.py +++ b/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 diff --git a/tests/units/modules/test_embedding.py b/tests/units/modules/test_embedding.py index a64e220da0e2597f294bc26ff360ac1ab0714270..037f5f0519546375c0140ea54d189a5b962b5319 100755 --- a/tests/units/modules/test_embedding.py +++ b/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 # ----------------------- diff --git a/tests/units/modules/test_encoder.py b/tests/units/modules/test_encoder.py index 69c8f23542d4475e309676eecf6a1ee82a281ace..5dfd2e0ea6247ff4610c2d7327fb2b80895405b9 100644 --- a/tests/units/modules/test_encoder.py +++ b/tests/units/modules/test_encoder.py @@ -1,7 +1,7 @@ import pytest import torch -from transformer.modules.encoder import EncoderLayer, TransformerEncoder +from mini_transformer.modules.encoder import EncoderLayer, TransformerEncoder def _dev(): @@ -24,6 +24,8 @@ def test_encoder_layer_param_checks(): EncoderLayer(32, 4, "64", 0.1) with pytest.raises(TypeError): EncoderLayer(32, 4, 64, "0.1") + with pytest.raises(TypeError): + EncoderLayer(32, 4, 64, 0.1, layer_norm_style=123) with pytest.raises(ValueError): EncoderLayer(0, 4, 64, 0.1) @@ -33,6 +35,8 @@ def test_encoder_layer_param_checks(): EncoderLayer(32, 4, 0, 0.1) with pytest.raises(ValueError): EncoderLayer(32, 4, 64, 1.0) # upper bound excluded + with pytest.raises(ValueError): + EncoderLayer(32, 4, 64, 0.1, layer_norm_style="middle") def test_transformer_encoder_layers_count_checks(): @@ -59,6 +63,22 @@ def test_encoder_forward_happy_path(B, S, D, H, FF, L): assert out.device == device +def test_transformer_encoder_pre_norm_layers_flag(): + enc = TransformerEncoder(24, 3, 48, 2, 0.1, layer_norm_style="pre") + assert all(layer.pre_norm for layer in enc.layers) + + +def test_encoder_pre_norm_forward_matches_shapes(): + device = _dev() + layer = EncoderLayer(24, 3, 48, 0.1, layer_norm_style="pre").to(device) + + x = torch.randn(2, 5, 24, device=device) + mask = torch.zeros(2, 3, 1, 5, dtype=torch.bool, device=device) + out = layer(x, mask) + assert out.shape == x.shape + assert layer.pre_norm is True + + def test_encoder_forward_input_checks_and_message_format(): layer = EncoderLayer(24, 3, 48, 0.1) diff --git a/tests/units/modules/test_feedforward.py b/tests/units/modules/test_feedforward.py index 8e479c46bad71db0d2f438ba80df221a09eb447b..83c65b9ba200f17a57e035595c8282b9615ccf2c 100644 --- a/tests/units/modules/test_feedforward.py +++ b/tests/units/modules/test_feedforward.py @@ -1,7 +1,7 @@ import pytest import torch -from transformer.modules.feedforward import FeedForwardLayer +from mini_transformer.modules.feedforward import FeedForwardLayer # ======================= # Constructor checks diff --git a/tests/units/modules/test_lm_head.py b/tests/units/modules/test_lm_head.py index 3819ad2e0ec04d6a54e12611a8f3557729db9954..93051c018d17c1a7a93e1fdc8c179f9681f4f821 100644 --- a/tests/units/modules/test_lm_head.py +++ b/tests/units/modules/test_lm_head.py @@ -1,7 +1,7 @@ import pytest import torch -from transformer.modules.lm_head import LMHead +from mini_transformer.modules.lm_head import LMHead # ----------------------- # Constructor checks diff --git a/tests/units/test_cli.py b/tests/units/test_cli.py new file mode 100644 index 0000000000000000000000000000000000000000..fe5a3a64224dba44040ff18a4707fabb20a77961 --- /dev/null +++ b/tests/units/test_cli.py @@ -0,0 +1,118 @@ +from pathlib import Path +from unittest import mock + +import pytest + +from mini_transformer import cli + +MINIMAL_CFG = """ +model: + name: demo + best_checkpoint_path: ./checkpoints/best.pt + latest_checkpoint_path: null + tokenizer: demo_tok + d_model: 16 + num_layers: 1 + num_heads: 2 + d_ff: 32 + dropout_rate: 0.0 + vocab_size: 32 + max_seq_len: 32 + pad_id: 0 + bos_id: 1 + eos_id: 2 + +tokenizer: + name: demo_tok + path: ./tokenizer/tokenizer.json + corpus: null + dataset: null + vocab_size: 32 + max_seq_len: 32 + pad_token: "" + bos_token: "" + eos_token: "" + unk_token: "" + special_tokens: ["", "", "", ""] + pad_id: 0 + bos_id: 1 + eos_id: 2 + unk_id: 3 + +generation: + max_new_tokens: 4 + 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 + +runtime: + seed: 42 + device: cpu + output_dir: ./outputs + data_dir: ./data + tokenizer_dir: ./tokenizer + cache_dir: ./cache + checkpoint_path: ./checkpoints + +input_text: "" +""" + + +def _make_demo_model(tmp_path: Path) -> Path: + models_root = tmp_path / "trained_models" + model_dir = models_root / "demo" + config_dir = model_dir / "configs" + config_dir.mkdir(parents=True, exist_ok=True) + (model_dir / "checkpoints").mkdir(exist_ok=True) + (model_dir / "tokenizer").mkdir(exist_ok=True) + (config_dir / "config_inference.yaml").write_text(MINIMAL_CFG) + return models_root + + +def test_infer_main_uses_local_model(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, capsys): + models_root = _make_demo_model(tmp_path) + monkeypatch.setenv(cli.MODELS_ENV, str(models_root)) + + with mock.patch.object(cli, "run_inference", return_value=["hello"]) as run_mock: + rc = cli.infer_main(["-m", "demo", "-t", "hi there"]) + + assert rc == 0 + cfg_passed = run_mock.call_args[0][0] + assert Path(cfg_passed.model.best_checkpoint_path).is_absolute() + assert cfg_passed.input_text == "hi there" + out = capsys.readouterr().out + assert "hello" in out + + +def test_infer_main_defaults_to_first_model(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + models_root = _make_demo_model(tmp_path) + monkeypatch.setenv(cli.MODELS_ENV, str(models_root)) + + with mock.patch.object(cli, "run_inference", return_value=["out"]) as run_mock: + cli.infer_main(["-t", "text"]) + assert run_mock.called + + +def test_fetch_main_invokes_snapshot(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + monkeypatch.setenv(cli.MODELS_ENV, str(tmp_path / "trained_models")) + + with mock.patch("mini_transformer.cli.snapshot_to_local", return_value=tmp_path) as snap: + rc = cli.fetch_main( + [ + "AlaBoussoffara/transformer_test", + "--name", + "demo", + "--force", + ] + ) + + assert rc == 0 + snap.assert_called_once() + kwargs = snap.call_args.kwargs + assert kwargs["local_name"] == "demo" + assert kwargs["force"] is True diff --git a/tests/units/test_server.py b/tests/units/test_server.py new file mode 100644 index 0000000000000000000000000000000000000000..81f1a278d2a2537e81aeca990c865cfda2281913 --- /dev/null +++ b/tests/units/test_server.py @@ -0,0 +1,86 @@ +from pathlib import Path + +from fastapi.testclient import TestClient + +from mini_transformer.apps import server +from mini_transformer.model_loader import MODEL_NAME_ENV, MODELS_ENV + +MINIMAL_CFG = """ +model: + name: demo + best_checkpoint_path: ./checkpoints/best.pt + latest_checkpoint_path: null + tokenizer: demo_tok + d_model: 16 + num_layers: 1 + num_heads: 2 + d_ff: 32 + dropout_rate: 0.0 + vocab_size: 32 + max_seq_len: 32 + pad_id: 0 + bos_id: 1 + eos_id: 2 + +tokenizer: + name: demo_tok + path: ./tokenizer/tokenizer.json + corpus: null + dataset: null + vocab_size: 32 + max_seq_len: 32 + pad_token: "" + bos_token: "" + eos_token: "" + unk_token: "" + special_tokens: ["", "", "", ""] + pad_id: 0 + bos_id: 1 + eos_id: 2 + unk_id: 3 + +generation: + max_new_tokens: 4 + 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 + +runtime: + seed: 42 + device: cpu + output_dir: ./outputs + data_dir: ./data + tokenizer_dir: ./tokenizer + cache_dir: ./cache + checkpoint_path: ./checkpoints + +input_text: "" +""" + + +def _make_demo_model(tmp_path: Path) -> Path: + models_root = tmp_path / "trained_models" + model_dir = models_root / "demo" + config_dir = model_dir / "configs" + config_dir.mkdir(parents=True, exist_ok=True) + (model_dir / "checkpoints").mkdir(exist_ok=True) + (model_dir / "tokenizer").mkdir(exist_ok=True) + (config_dir / "config_inference.yaml").write_text(MINIMAL_CFG) + return models_root + + +def test_generate_endpoint_uses_local_model(monkeypatch, tmp_path): + models_root = _make_demo_model(tmp_path) + monkeypatch.setenv(MODELS_ENV, str(models_root)) + monkeypatch.setenv(MODEL_NAME_ENV, "demo") + monkeypatch.setattr(server, "run_inference", lambda cfg: ["ok"]) + + client = TestClient(server.app) + response = client.get("/generate", params={"text": "hello"}) + assert response.status_code == 200 + assert response.json()["outputs"] == ["ok"] diff --git a/tests/units/test_transformer.py b/tests/units/test_transformer.py index 9d5028c4e72f2674c1d47b7456da59bd89f41fdb..ad17e1e84afa5a3bc45b0e1889a3c1fc1cf8a287 100644 --- a/tests/units/test_transformer.py +++ b/tests/units/test_transformer.py @@ -1,9 +1,9 @@ import pytest import torch -from transformer.configs import ModelCfg -from transformer.transformer import BasicEncoderDecoderTransformer -from transformer.utils import broadcast_padding_mask +from mini_transformer.configs import ModelCfg +from mini_transformer.transformer import BasicEncoderDecoderTransformer +from mini_transformer.utils import broadcast_padding_mask def _dev(): @@ -72,6 +72,26 @@ def test_ctor_cfg_type_and_values(): assert model.embed.token_embed.weight.data_ptr() == model.lm_head.fc.weight.data_ptr() +def test_layer_norm_style_auto_selection(): + shallow = BasicEncoderDecoderTransformer(_tiny_cfg(num_layers=4)) + assert shallow.layer_norm_style == "post" + assert all(not layer.pre_norm for layer in shallow.encoder.layers) + assert all(not layer.pre_norm for layer in shallow.decoder.layers) + + deep = BasicEncoderDecoderTransformer(_tiny_cfg(num_layers=6)) + assert deep.layer_norm_style == "pre" + assert all(layer.pre_norm for layer in deep.encoder.layers) + assert all(layer.pre_norm for layer in deep.decoder.layers) + + +def test_layer_norm_style_override_respected(): + cfg = _tiny_cfg(num_layers=6, layer_norm_style="post") + model = BasicEncoderDecoderTransformer(cfg) + assert model.layer_norm_style == "post" + assert all(not layer.pre_norm for layer in model.encoder.layers) + assert all(not layer.pre_norm for layer in model.decoder.layers) + + # ----------------------- # Encode/Decode/Forward — shapes, dtypes, devices, zero-length # ----------------------- diff --git a/tests/units/test_utils.py b/tests/units/test_utils.py index a45c811f1a5ac7d9504c8f5866f9d1700e6d33ea..409d5625d783a902114526327e61c4ba369d224e 100755 --- a/tests/units/test_utils.py +++ b/tests/units/test_utils.py @@ -4,7 +4,7 @@ import pytest import torch import torch.nn.functional as F -from transformer.utils import ( +from mini_transformer.utils import ( calculate_attention, combine_masks, create_causal_mask,