fixes

CONTRIBUTING.md

@@ -4,7 +4,7 @@ Thanks for your interest in contributing! Follow these steps to get started.
 
 1. **Create the development environment**
    ```bash
-   make conda-setup
+   make create-env
    ```
 
 2. **Run quality checks** before opening a pull request:

Makefile

@@ -61,10 +61,10 @@ cov:
 	$(CONDA_RUN) pytest --cov=mini_transformer --cov-report=term-missing
 
 fetch-test:
-	$(CONDA_RUN) mini-transformer-fetch AlaBoussoffara/transformer_test --force
+	$(CONDA_RUN) mini-transformer-fetch AlaBoussoffara/transformer_test --name test_model_v1 --force
 
 fetch-small:
-	$(CONDA_RUN) mini-transformer-fetch AlaBoussoffara/transformer_small --force
+	$(CONDA_RUN) mini-transformer-fetch AlaBoussoffara/transformer_small --name small_model_v1 --force
 
 infer:
 	$(CONDA_RUN) mini-transformer-infer --model small_model_v1 -t "translate this sentence."

README.md

@@ -1,4 +1,4 @@
-# mini-transformer
+# 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.
 
@@ -30,12 +30,11 @@ pip install "mini-transformer[server,viz]"
 ## 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
+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/`), extract or move those contents up into the root of `trained_models/<local-name>/` so the structure contains `configs/`, `checkpoints/`, and `tokenizer/` directly.
-In other words, after fetching `AlaBoussoffara/transformer_small`, move the `transformer_small/small_model_v1/` directory into `trained_models/AlaBoussoffara__transformer_small/`.
+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>/`:
 ```
@@ -69,7 +68,7 @@ Follow these steps to spin up the bundled Chainlit chat UI for a local inference
    ```
 2. Download a demo checkpoint (skip if you already have one under `trained_models/`):
    ```bash
-   mini-transformer-fetch --model-id AlaBoussoffara/transformer_small
+   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
@@ -96,7 +95,7 @@ print(run_inference(cfg)[0])
 
 ## Development Quickstart
 ```bash
-make conda-setup
+make create-env
 make lint
 make type
 make test
@@ -122,7 +121,7 @@ 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`
+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`

src/mini_transformer/modules/attention.py

@@ -124,7 +124,7 @@ class MultiHeadAttention(nn.Module):
         k = split_heads(self.key_linear(key), self.num_heads)
         v = split_heads(self.value_linear(value), self.num_heads)
         pad_mask = create_qk_padding_mask(q_mask, k_mask)
-        combined_mask = combine_masks(pad_mask, causal_mask)
+        combined_mask = combine_masks(causal_mask, pad_mask)
 
         p = self.dropout_rate if self.training else 0.0
         attn = cast(Tensor, calculate_attention(q, k, v, combined_mask, dropout_p=p))

src/mini_transformer/transformer.py

@@ -127,7 +127,7 @@ class BasicEncoderDecoderTransformer(nn.Module):
                 f"src_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(src_ids.shape)}"
             )
         if src_ids.dtype != torch.long:
-            raise TypeError(f"src_idsmust be torch.long (int64), got {src_ids.dtype}")
+            raise TypeError(f"src_ids must be torch.long (int64), got {src_ids.dtype}")
 
         if not isinstance(tgt_ids, Tensor):
             raise TypeError(f"tgt_ids must be a torch.Tensor, got {type(tgt_ids)}")

src/mini_transformer/utils.py

@@ -511,7 +511,7 @@ def extract_all_attention_maps(
         k = split_heads(self_mha.key_linear(y_in), self_mha.num_heads)
         v = split_heads(self_mha.value_linear(y_in), self_mha.num_heads)
         pad = create_qk_padding_mask(tgt_pad_b, tgt_pad_b)
-        combined_mask = combine_masks(pad, tgt_causal)
+        combined_mask = combine_masks(tgt_causal, pad)
         self_ctx, self_probs = calculate_attention(
             q, k, v, combined_mask, dropout_p=0.0, return_probs=True
         )