---
title: Mini Transformer Demo
emoji: 🤖
colorFrom: yellow
colorTo: purple
sdk: docker
app_port: 7860
---

# Mini Transformer Demo

Chainlit-powered dialogue interface that loads a mini encoder–decoder Transformer from `trained_models/` and serves it via Hugging Face Spaces. If you want the full project write-up, see [README.project.md](README.project.md).
# Mini-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.

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

## Models

### Small Model
The small model is a compact encoder-decoder transformer designed for efficient machine translation. Architecture details:
- Parameters: ~25M
- Architecture: 4 encoder/decoder layers, 4 attention heads, d_model=512, d_ff=1024
- Vocabulary: BPE tokenizer with 8K tokens
- Maximum sequence length: 128 tokens

Performance metrics (tested on 10K samples):
- BLEU score: 25.13 (EN-FR translation)
- Token accuracy: 63.43%
- Perplexity: 3.532
- Latency: 38.5ms average, 48.1ms p95 per sentence
- Additional metrics: chrF: 56.17, ROUGE-1: 59.19%, ROUGE-2: 40.03%

The model demonstrates good balance between performance and computational efficiency, making it suitable for deployment in resource-constrained environments while maintaining reasonable translation quality.

## Hugging Face Spaces DEMO
Here you can find a Huggingface Spaces for a quick demo
https://huggingface.co/spaces/AlaBoussoffara/Mini-Transformer
## 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.

## Training Notebook
`notebooks/train.ipynb` demonstrates the Hydra-driven training loop used during development.
- All hyper-parameters come from the composed config (see `configs/train_mode.yaml`); the key knobs live under `trainer.*`.
- Checkpoint cadence is controlled by `trainer.save_interval` (in optimizer steps) and always saves on epoch boundaries.
- Dataloader behaviour (workers, pinned memory) can be tuned via `trainer.num_workers` and `trainer.pin_memory`.
- Gradient accumulation is respected even for partial micro-batch sets, so you can safely mix different batch counts.
For long-running jobs consider exporting the notebook to a script (`jupyter nbconvert --to script`) or reusing the same logic inside a CLI tool.

## Preparing Models
Download the demo models hosted on Hugging Face:
```bash
mini-transformer-fetch AlaBoussoffara/transformer_test
mini-transformer-fetch AlaBoussoffara/transformer_small
```
Use `--name` to customise the local directory and `--force` to refresh an existing download.
The local folder name defaults to `<author>__<repo>` (e.g. `AlaBoussoffara__transformer_small`). The Hugging Face repo is downloaded as a whole, so if the model files live inside a subfolder (for example `transformer_small/small_model_v1/`), move that entire inner folder—the one that already contains `configs/`, `checkpoints/`, and `tokenizer/`—into `trained_models/` so it becomes your model directory. In practice, after fetching `AlaBoussoffara/transformer_small`, move the `transformer_small/small_model_v1/` directory into `trained_models/` and rename it to the directory name you want to use (for example `AlaBoussoffara__transformer_small/`) so that the final layout is `trained_models/<model-name>/configs/...`.

You can also place your own model manually under `trained_models/<model-name>/`:
```
trained_models/
  my-model/
    configs/
      config_inference.yaml
    checkpoints/
      best.pt
    tokenizer/
      tokenizer.json
```
Relative paths in `config_inference.yaml` should stay inside the model folder. Set `MINI_TRANSFORMER_MODELS=/path/to/trained_models` if you store models elsewhere.

## CLI Usage
Install the optional extras (see Installation) and use the commands below.

```bash
mini-transformer-infer --model AlaBoussoffara__transformer_small -t "Once upon a time"
mini-transformer-serve --model AlaBoussoffara__transformer_small --reload
mini-transformer-ui --model AlaBoussoffara__transformer_small --host 0.0.0.0 --port 8000
```
If you see “Tokenizer file not found”, update the model config or set `MINI_TRANSFORMER_TOKENIZER_PATH` to the correct JSON file.

## Chainlit Demo
Follow these steps to spin up the bundled Chainlit chat UI for a local inference demo.

1. Install the project with the server extras (or run `make create-env`):
   ```bash
   pip install -e .[server]
   ```
2. Download a demo checkpoint (skip if you already have one under `trained_models/`):
   ```bash
   mini-transformer-fetch 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.
   Use `/config temperature=0.7 top_k=50` (or `/config reset`) to tweak generation settings on the fly.

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

## Docker Inference UI
Build the lean inference image and launch the Chainlit UI in a container:
```bash
make docker-build-infer
make docker-run-infer  # press Ctrl+C to stop
```
Models stay outside the container under `trained_models/`, mounted at runtime with docker compose.

## 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 create-env
make lint
make type
make test
pre-commit run --all-files   # optional: run all hooks locally
```

## Testing & QA
- Unit tests live under `tests/units/`; run them with `make test` or `python -m pytest`.
- The suite covers core building blocks (attention math, masking, sampling), CLI flows, and attention-debug utilities—including both pre- and post-layernorm configurations.
- Add tests alongside new features; keeping coverage high ensures `mini-transformer` behaves the same whether it runs from the repo or as an installed package.

## 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 precommit     # run the configured pre-commit hooks (auto-fixes where possible)
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-dev   # build the development container
make docker-run-dev     # open an interactive shell in the development container
make docker-build-infer # build the inference/UI container
make docker-run-infer   # run the inference service (Ctrl+C to stop)
```

## Full Workflow At A Glance
1. Fetch a model: `mini-transformer-fetch 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.