README.md

metadata

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.

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

# inside the repository (editable install with optional extras)
pip install -e .[server,viz]

Or install the published package:

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:

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.

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):

   pip install -e .[server]
   

2. Download a demo checkpoint (skip if you already have one under trained_models/):

   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):

   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:

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

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

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

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.