added evaluation metrics & evaluated small model

README.md

@@ -9,6 +9,24 @@ A compact, encoder–decoder Transformer packaged so it can be used both as a Py
 - 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

environment.yml

@@ -26,3 +26,8 @@ dependencies:
       - mypy
       - pre-commit
       - httpx
+  - evaluate
+  - sacrebleu
+  - rouge-score
+  - bert-score
+  - tqdm

notebooks/train.ipynb

@@ -103,12 +103,14 @@
     "cfg = compose(\n",
     "    config_name=\"train_mode\",\n",
     "    overrides=[\n",
-    "        \"model=medium_model\",\n",
-    "        \"tokenizer=bpe_16k\",\n",
+    "        \"model=small_model\",\n",
+    "        \"tokenizer=bpe_8k\",\n",
     "        \"dataset=medium_dataset\",\n",
     "        \"trainer.batch_size=32\",\n",
     "        \"trainer.gradient_accumulation=16\",\n",
     "        \"trainer.epochs=10\",\n",
+    "        \"trainer.resume=best\",\n",
+    "        \"model.best_checkpoint_path=trained_models/small_model_v1/checkpoints/best/checkpoint.pt\",\n",
     "    ],\n",
     ")\n",
     "scfg_temp = OmegaConf.merge(OmegaConf.structured(TrainAppCfg), cfg)\n",
@@ -221,7 +223,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Total number of parameters: 184705024\n"
+      "Total number of parameters: 25204736\n"
      ]
     }
    ],
@@ -444,7 +446,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Starting fresh (no resume checkpoint provided).\n"
+      "Loading checkpoint from /mnt/shared/Local/Projets/Mini-Transformer/trained_models/small_model_v1/checkpoints/best/checkpoint.pt\n"
      ]
     }
    ],
@@ -757,7 +759,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 13,
    "id": "439d4094",
    "metadata": {},
    "outputs": [
@@ -765,7 +767,22 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Epoch 1 [train]:   2%|▏         | 1380/57063 [10:36<6:56:33,  2.23it/s]"
+      "Epoch 10 [train]:   0%|          | 33/57063 [00:03<1:47:43,  8.82it/s]\n",
+      "Epoch 10 [train]:   0%|          | 33/57063 [00:03<1:47:43,  8.82it/s]\n"
+     ]
+    },
+    {
+     "ename": "KeyboardInterrupt",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
+      "\u001b[31mKeyboardInterrupt\u001b[39m                         Traceback (most recent call last)",
+      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[13]\u001b[39m\u001b[32m, line 59\u001b[39m\n\u001b[32m     57\u001b[39m loss = loss / grad_accum\n\u001b[32m     58\u001b[39m scaled_loss = scaler.scale(loss)\n\u001b[32m---> \u001b[39m\u001b[32m59\u001b[39m \u001b[43mscaled_loss\u001b[49m\u001b[43m.\u001b[49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m     60\u001b[39m accum_counter += \u001b[32m1\u001b[39m\n\u001b[32m     62\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m accum_counter == grad_accum:\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/miniconda3/envs/mini-transformer/lib/python3.11/site-packages/torch/_tensor.py:525\u001b[39m, in \u001b[36mTensor.backward\u001b[39m\u001b[34m(self, gradient, retain_graph, create_graph, inputs)\u001b[39m\n\u001b[32m    515\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[32m    516\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[32m    517\u001b[39m         Tensor.backward,\n\u001b[32m    518\u001b[39m         (\u001b[38;5;28mself\u001b[39m,),\n\u001b[32m   (...)\u001b[39m\u001b[32m    523\u001b[39m         inputs=inputs,\n\u001b[32m    524\u001b[39m     )\n\u001b[32m--> \u001b[39m\u001b[32m525\u001b[39m \u001b[43mtorch\u001b[49m\u001b[43m.\u001b[49m\u001b[43mautograd\u001b[49m\u001b[43m.\u001b[49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m    526\u001b[39m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgradient\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m=\u001b[49m\u001b[43minputs\u001b[49m\n\u001b[32m    527\u001b[39m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/miniconda3/envs/mini-transformer/lib/python3.11/site-packages/torch/autograd/__init__.py:267\u001b[39m, in \u001b[36mbackward\u001b[39m\u001b[34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[39m\n\u001b[32m    262\u001b[39m     retain_graph = create_graph\n\u001b[32m    264\u001b[39m \u001b[38;5;66;03m# The reason we repeat the same comment below is that\u001b[39;00m\n\u001b[32m    265\u001b[39m \u001b[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[32m    266\u001b[39m \u001b[38;5;66;03m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m267\u001b[39m \u001b[43m_engine_run_backward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m    268\u001b[39m \u001b[43m    \u001b[49m\u001b[43mtensors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    269\u001b[39m \u001b[43m    \u001b[49m\u001b[43mgrad_tensors_\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    270\u001b[39m \u001b[43m    \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    271\u001b[39m \u001b[43m    \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    272\u001b[39m \u001b[43m    \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    273\u001b[39m \u001b[43m    \u001b[49m\u001b[43mallow_unreachable\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[32m    274\u001b[39m \u001b[43m    \u001b[49m\u001b[43maccumulate_grad\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[32m    275\u001b[39m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/miniconda3/envs/mini-transformer/lib/python3.11/site-packages/torch/autograd/graph.py:744\u001b[39m, in \u001b[36m_engine_run_backward\u001b[39m\u001b[34m(t_outputs, *args, **kwargs)\u001b[39m\n\u001b[32m    742\u001b[39m     unregister_hooks = _register_logging_hooks_on_whole_graph(t_outputs)\n\u001b[32m    743\u001b[39m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[32m--> \u001b[39m\u001b[32m744\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mVariable\u001b[49m\u001b[43m.\u001b[49m\u001b[43m_execution_engine\u001b[49m\u001b[43m.\u001b[49m\u001b[43mrun_backward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# Calls into the C++ engine to run the backward pass\u001b[39;49;00m\n\u001b[32m    745\u001b[39m \u001b[43m        \u001b[49m\u001b[43mt_outputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\n\u001b[32m    746\u001b[39m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# Calls into the C++ engine to run the backward pass\u001b[39;00m\n\u001b[32m    747\u001b[39m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[32m    748\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m attach_logging_hooks:\n",
+      "\u001b[31mKeyboardInterrupt\u001b[39m: "
      ]
     }
    ],
@@ -1041,13 +1058,60 @@
    "execution_count": null,
    "id": "39f61da8",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Testing:  32%|███▏      | 101/313 [02:04<03:10,  1.11it/s]"
+     ]
+    }
+   ],
    "source": [
-    "# test the model\n",
+    "# test the model (extended: compute BLEU/chrF/ROUGE, perplexity, exact/token accuracy)\n",
+    "import time\n",
+    "\n",
+    "import evaluate\n",
+    "import torch.nn.functional as F\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "# select some samples to test\n",
+    "test_dataset = test_dataset.select(range(10000))\n",
+    "test_dataloader = DataLoader(\n",
+    "    test_dataset,\n",
+    "    batch_size=scfg.trainer.batch_size,\n",
+    "    shuffle=False,\n",
+    "    collate_fn=tuple_collate_fn,\n",
+    "    generator=dataloader_generator,\n",
+    "    worker_init_fn=worker_init,\n",
+    "    num_workers=num_workers,\n",
+    "    pin_memory=pin_memory,\n",
+    "    persistent_workers=persistent_workers,\n",
+    ")\n",
     "model.eval()\n",
     "test_loss = 0.0\n",
+    "all_hyps = []\n",
+    "all_refs = []\n",
+    "latencies = []\n",
+    "total_nll = 0.0\n",
+    "total_tokens = 0\n",
+    "exact_matches = 0\n",
+    "token_matches = 0\n",
+    "token_total = 0\n",
+    "\n",
+    "bleu_metric = evaluate.load(\"sacrebleu\")\n",
+    "chrf_metric = evaluate.load(\"chrf\")\n",
+    "rouge_metric = evaluate.load(\"rouge\")\n",
+    "\n",
+    "pad_id = (\n",
+    "    getattr(tokenizer_cfg, \"pad_id\", None)\n",
+    "    if \"tokenizer_cfg\" in globals()\n",
+    "    else getattr(model_cfg, \"pad_id\", None)\n",
+    ")\n",
+    "\n",
     "with torch.no_grad(), autocast_ctx:\n",
     "    for src, tgt in tqdm(test_dataloader, desc=\"Testing\"):\n",
+    "        # tokenization (src and tgt are lists of strings per tuple_collate_fn)\n",
     "        encoded_src = tokenizer(\n",
     "            src,\n",
     "            padding=True,\n",
@@ -1067,18 +1131,96 @@
     "        tgt_ids = encoded_tgt[\"input_ids\"].to(device)\n",
     "        decoder_in = tgt_ids[:, :-1]\n",
     "        labels = tgt_ids[:, 1:]\n",
-    "        tgt_padd_mask = decoder_in.eq(scfg.tokenizer.pad_id)\n",
+    "        tgt_padd_mask = decoder_in.eq(pad_id)\n",
     "\n",
+    "        # forward for test loss (uses configured criterion)\n",
     "        logits = model(src_ids, decoder_in, src_padd_mask, tgt_padd_mask)\n",
     "        loss = criterion(logits.reshape(-1, logits.size(-1)), labels.reshape(-1))\n",
     "        test_loss += loss.item()\n",
+    "\n",
+    "        # compute token-summed NLL for perplexity\n",
+    "        V = logits.size(-1)\n",
+    "        logits_flat = logits.reshape(-1, V)  # Changed from view to reshape\n",
+    "        labels_flat = labels.reshape(-1)  # Changed from view to reshape\n",
+    "        nll_sum = F.cross_entropy(logits_flat, labels_flat, ignore_index=pad_id, reduction=\"sum\")\n",
+    "        non_pad = (labels_flat != pad_id).sum().item()\n",
+    "        total_nll += nll_sum.item()\n",
+    "        total_tokens += non_pad\n",
+    "\n",
+    "        # generation and latency\n",
+    "        t0 = time.perf_counter()\n",
+    "        out_ids = model.generate(\n",
+    "            src_ids, src_padd_mask, max_new_tokens=128, temperature=1.0, do_sample=False\n",
+    "        )\n",
+    "        t1 = time.perf_counter()\n",
+    "        latencies.append((t1 - t0) / max(1, src_ids.size(0)))\n",
+    "\n",
+    "        hyps = tokenizer.batch_decode(out_ids.cpu(), skip_special_tokens=True)\n",
+    "        all_hyps.extend(hyps)\n",
+    "        all_refs.extend(tgt)\n",
+    "\n",
+    "        # exact match and token-level accuracy\n",
+    "        for h, r in zip(hyps, tgt, strict=False):\n",
+    "            if h.strip().lower() == r.strip().lower():\n",
+    "                exact_matches += 1\n",
+    "\n",
+    "        hyp_tok = tokenizer(hyps, padding=True, truncation=True, return_tensors=\"pt\")[\"input_ids\"]\n",
+    "        ref_tok = tgt_ids[:, : hyp_tok.size(1)].cpu()\n",
+    "        matches = (hyp_tok[:, : ref_tok.size(1)] == ref_tok).cpu()\n",
+    "        token_matches += int(matches.sum().item())\n",
+    "        token_total += int(ref_tok.numel())\n",
+    "\n",
+    "# aggregate and report\n",
     "avg_test_loss = test_loss / max(1, len(test_dataloader))\n",
+    "avg_latency = sum(latencies) / len(latencies) if latencies else None\n",
+    "p95_latency = sorted(latencies)[int(0.95 * len(latencies))] if latencies else None\n",
+    "avg_nll = (total_nll / total_tokens) if total_tokens > 0 else float(\"nan\")\n",
+    "perplexity = math.exp(avg_nll) if total_tokens > 0 and avg_nll < 100 else float(\"inf\")\n",
+    "exact_match_rate = exact_matches / len(all_hyps) if all_hyps else 0.0\n",
+    "token_accuracy = token_matches / token_total if token_total > 0 else 0.0\n",
+    "\n",
+    "# corpus-level metrics\n",
+    "refs_for_metric = [[r] for r in all_refs]\n",
+    "bleu_res = bleu_metric.compute(predictions=all_hyps, references=refs_for_metric)\n",
+    "chrf_res = chrf_metric.compute(predictions=all_hyps, references=refs_for_metric)\n",
+    "rouge_res = rouge_metric.compute(predictions=all_hyps, references=[r[0] for r in refs_for_metric])\n",
+    "\n",
     "print(f\"Test Loss: {avg_test_loss:.4f}\")\n",
-    "# save test loss to the manifest\n",
-    "manifest_data[\"test\"] = {\n",
+    "print(f\"Generated {len(all_hyps)} hypotheses\")\n",
+    "print(f\"Avg latency per sentence: {avg_latency:.4f} s\")\n",
+    "print(f\"p95 latency per sentence: {p95_latency:.4f} s\")\n",
+    "print(f\"Perplexity (corpus): {perplexity:.3f}\")\n",
+    "print(f\"Exact match rate: {exact_match_rate:.3%}\")\n",
+    "print(f\"Token accuracy: {token_accuracy:.3%}\")\n",
+    "print(\"BLEU:\", bleu_res)\n",
+    "print(\"chrF:\", chrf_res)\n",
+    "print(\"ROUGE:\", rouge_res)\n",
+    "\n",
+    "# Prepare test metrics\n",
+    "test_metrics = {\n",
     "    \"test_loss\": avg_test_loss,\n",
+    "    \"perplexity\": perplexity,\n",
+    "    \"avg_latency_s\": avg_latency,\n",
+    "    \"p95_latency_s\": p95_latency,\n",
+    "    \"exact_match_rate\": exact_match_rate,\n",
+    "    \"token_accuracy\": token_accuracy,\n",
+    "    \"bleu\": bleu_res,\n",
+    "    \"chrf\": chrf_res,\n",
+    "    \"rouge\": rouge_res,\n",
     "    \"timestamp\": datetime.now().isoformat(),\n",
+    "    \"model_name\": model_cfg.name,\n",
+    "    \"dataset_size\": len(test_dataset),\n",
+    "    \"checkpoint_path\": str(resume_checkpoint_path) if resume_checkpoint_path else None,\n",
     "}\n",
+    "\n",
+    "# Save to manifest\n",
+    "manifest_data[\"test\"] = test_metrics\n",
+    "\n",
+    "# Save metrics to a separate file\n",
+    "metrics_file = run_dir / \"test_metrics.json\"\n",
+    "metrics_file.write_text(json.dumps(test_metrics, indent=2))\n",
+    "print(f\"[saved] Test metrics -> {metrics_file}\")\n",
+    "\n",
     "manifest_path.write_text(json.dumps(manifest_data, indent=2))"
    ]
   },

pyproject.toml

@@ -46,6 +46,14 @@ notebook = [
   "jupyterlab",
 ]
 
+evaluation = [
+  "evaluate",
+  "sacrebleu",
+  "rouge-score",
+  "bert-score",
+  "tqdm",
+]
+
 
 [project.scripts]
 mini-transformer-infer = "mini_transformer.cli:infer_main"

requirements.txt

@@ -9,3 +9,9 @@ huggingface_hub
 chainlit
 uvicorn[standard]
 fastapi
+
+evaluate
+sacrebleu
+rouge-score
+bert-score
+tqdm