Update README with model loading code

README.md

@@ -13,9 +13,16 @@ Part of the [LAPVQA collection](https://huggingface.co/collections/dmusingu/lapv
 
 ## Description
 
-RRG decoder heads trained **end-to-end** alongside their vision encoders, providing
-a fine-tuned baseline for comparison with the frozen-encoder setup in
-[`lapvqa-rrg`](https://huggingface.co/dmusingu/lapvqa-rrg).
+RRG decoders trained **end-to-end** alongside their vision encoders.
+Each checkpoint is a dict: `{state_dict, vis_dim, d_model, num_layers, nhead, encoder, epoch, val_bleu4}`.
+
+| File | Encoder | vis_dim |
+|---|---|---|
+| `clip-vit-l14.pt` | CLIP ViT-L/14 (fine-tuned) | 1024 |
+| `siglip.pt` | SigLIP (fine-tuned) | 1152 |
+| `florence2.pt` | Florence-2 (fine-tuned) | 1024 |
+| `coca.pt` | CoCa (fine-tuned) | 768 |
+| `mae-vit-l16.pt` | MAE ViT-L/16 (fine-tuned) | 1024 |
 
 ## Results (MIMIC-CXR test set, MAE-ViT-L/16)
 
@@ -23,12 +30,19 @@ a fine-tuned baseline for comparison with the frozen-encoder setup in
 |---|---|---|
 | 0.032 | 0.164 | 0.195 |
 
-## Files
-
-| File | Encoder backbone |
-|---|---|
-| `clip-vit-l14.pt` | CLIP ViT-L/14 (fine-tuned) |
-| `siglip.pt` | SigLIP (fine-tuned) |
-| `florence2.pt` | Florence-2 (fine-tuned) |
-| `coca.pt` | CoCa (fine-tuned) |
-| `mae-vit-l16.pt` | MAE ViT-L/16 (fine-tuned) |
+## Loading
+
+```python
+import torch
+from lapvqa.rrg.heads import ReportGenerationHead
+
+ckpt = torch.load("mae-vit-l16.pt", map_location="cpu")
+head = ReportGenerationHead(
+    vis_dim    = ckpt["vis_dim"],
+    d_model    = ckpt["d_model"],
+    num_layers = ckpt["num_layers"],
+    nhead      = ckpt["nhead"],
+)
+head.load_state_dict(ckpt["state_dict"])
+head.eval()
+```