Fix generation from the published artifact; rewrite model card

README.md

@@ -16,11 +16,9 @@ tags:
 base_model:
 - OpenGVLab/InternViT-300M-448px-V2_5
 - PleIAs/Baguettotron
-datasets:
-- liuhaotian/LLaVA-CC3M-Pretrain-595K
 ---
 
-# Baguettotron-VLM — Stage 1 (Projector Warmup)
+# baguettotron-internvit-alignment
 
 **Baguettotron-VLM** is an open, fully-reproducible, multilingual Vision-Language
 Model in the **sub-1B parameter** class. It extends
@@ -29,26 +27,24 @@ text-only reasoning SLM — with visual capabilities via the
 [InternViT-300M-448px-V2.5](https://huggingface.co/OpenGVLab/InternViT-300M-448px-V2_5)
 vision encoder and a lightweight MLP projector, for a total of ~628M parameters.
 
-It supports **six European languages** (EN, FR, DE, ES, IT, PL, inherited from
-the Baguettotron backbone), natively emits `<think>` reasoning traces, and was
-trained end-to-end on a single H100 SXM — code, data pipeline, weights, and
-training logs are all public. Everything is released under **Apache 2.0** with
-a proof-of-concept goal: exploring what an open sub-1B multilingual VLM can do,
-not chasing SOTA.
+Trained end-to-end in **~72h on a single H100 (~€219)**. It inherits six European
+languages (EN, FR, DE, ES, IT, PL) from the Baguettotron backbone, and — unusually —
+its entire LM pretraining corpus is auditable: Baguettotron was trained on
+[SYNTH](https://huggingface.co/datasets/PleIAs/SYNTH), ~58K Wikipedia articles expanded
+synthetically, with no scraped web.
 
-The project ships three checkpoints from a staged training curriculum:
+Two checkpoints are published:
 
-- **Stage 1 — projector warmup**: [andreagemelli/Baguettotron-VLM-Stage1](https://huggingface.co/andreagemelli/Baguettotron-VLM-Stage1)
-  (MLP projector only, LLaVA-CC3M-Pretrain-595K — published for reproducibility).
-- **Stage 2 — instruction tuning**: [andreagemelli/Baguettotron-VLM-Stage2](https://huggingface.co/andreagemelli/Baguettotron-VLM-Stage2)
-  (full fine-tune on The Cauldron + SYNTH, direct-answer mode).
-- **Stage 3 — reasoning SFT (flagship)**: [andreagemelli/Baguettotron-VLM](https://huggingface.co/andreagemelli/Baguettotron-VLM)
-  (reasoning traces via R1-Vision-Reasoning-Instructions, dynamic `<think>` toggle).
+- [**baguettotron-vision-vqa**](https://huggingface.co/andreagemelli/baguettotron-vision-vqa)
+  — instruction-tuned on The Cauldron. **This is the model you want.**
+- [baguettotron-internvit-alignment](https://huggingface.co/andreagemelli/baguettotron-internvit-alignment)
+  — projector-only warmup, published so the alignment stage can be reproduced.
 
-Source & training pipeline: [github.com/andreagemelli/baguettotron-vlm](https://github.com/andreagemelli/baguettotron-vlm).
+Apache 2.0. This is a proof of concept and a costed recipe, not a SOTA attempt.
+Source: [github.com/andreagemelli/baguettotron-vlm](https://github.com/andreagemelli/baguettotron-vlm).
 
 
-> **Stage 1 checkpoint.** Only the MLP projector has been trained (on LLaVA-CC3M-Pretrain-595K); the ViT and the Baguettotron LLM are the unmodified base weights. Published for reproducibility — for actual use prefer the Stage 2 or Stage 3 checkpoints.
+> **Alignment checkpoint.** Only the MLP projector was trained (on LLaVA-CC3M-Pretrain-595K) — the ViT and the Baguettotron LLM are unmodified base weights. It writes short captions and nothing more. Published so the alignment stage can be reproduced; for actual use take [baguettotron-vision-vqa](https://huggingface.co/andreagemelli/baguettotron-vision-vqa).
 
 ## Architecture
 
@@ -59,32 +55,35 @@ Image (448×448)
   → MLP projector (2-layer, ~2.7M)            → 256 tokens × 576d
   → Interleave with text tokens
   → Baguettotron (321M, Llama arch, 80L, h=576)
-  → Text output with <think> reasoning traces
 
 Total: ~628M parameters
 ```
 
 ## Usage
 
+```bash
+pip install "transformers>=4.56,<5" torch pillow timm einops accelerate
+```
+
 ```python
 import torch
 from transformers import AutoModelForImageTextToText, AutoProcessor
 from PIL import Image
 
 model = AutoModelForImageTextToText.from_pretrained(
-    "andreagemelli/Baguettotron-VLM-Stage1",
+    "andreagemelli/baguettotron-internvit-alignment",
     trust_remote_code=True,
-    torch_dtype=torch.bfloat16,
+    dtype=torch.bfloat16,
     device_map="auto",
 )
 processor = AutoProcessor.from_pretrained(
-    "andreagemelli/Baguettotron-VLM-Stage1",
+    "andreagemelli/baguettotron-internvit-alignment",
     trust_remote_code=True,
 )
 
 image = Image.open("photo.jpg").convert("RGB")
 inputs = processor(
-    messages=[{"role": "user", "content": "<image>\nDescribe the image"}],
+    messages=[{"role": "user", "content": "<image>\nDescribe the image concisely."}],
     image=image,
 )
 inputs = {k: v.to(model.device) for k, v in inputs.items() if v is not None}
@@ -92,13 +91,61 @@ inputs = {k: v.to(model.device) for k, v in inputs.items() if v is not None}
 print(model.chat(**inputs))
 ```
 
+`chat()` returns the answer as a plain string, already trimmed at the end of the
+assistant turn. `max_new_tokens=...` bounds the length.
+
+### Apple Silicon and CPU
+
+The snippet above runs unchanged on `mps` and `cpu` — keep `device_map="auto"`, or drop
+it and call `.to("mps")` yourself. `dtype=torch.float32` works too. A short answer takes
+~1-3s on `mps` against ~20-25s on CPU, so prefer `mps` on a Mac.
+
+One caveat handled for you: transformers' repetition-penalty processor corrupts the
+first decoding step on MPS, which produced a garbage first token. `chat()` disables the
+penalty on that backend automatically.
+
+### Example output
+
+Four COCO images, greedy, `Describe the image concisely.` — verbatim:
+
+| image | output |
+|---|---|
+| two cats on a couch | `a cat is sleeping on the couch` |
+| close-up of a brown bear | `the bear is a good friend.` |
+| an upside-down STOP sign | `a sign for a stop sign` |
+| a red double-decker bus | `the bus is a red double - decoration` |
+
+That is the honest range of this checkpoint: it locates the subject and then drifts, which is what 2.7M trained parameters buys. It does not read the text on the sign or the bus.
+
 ### Chat template
 
-Stage 1 was trained on short image captions with no `<think>` traces. The processor emits a bare assistant prefix (`<|im_start|>assistant\n`) and the model completes the caption directly. Keep prompts simple ("Describe the image").
+Trained on short image captions with no `<think>` traces. The processor emits a bare assistant prefix (`<|im_start|>assistant\n`) and the model completes the caption directly. Keep prompts simple ("Describe the image").
+
+## What the evaluation found
+
+Measured on 6 held-out [RealWorldQA](https://huggingface.co/datasets/xai-org/RealworldQA)
+images, judged 1–5 by a vision judge, 244 of 288 cells completed. Thin, and reported as
+such — full tables and caveats in
+[RESULTS.md](https://github.com/andreagemelli/baguettotron-vlm/blob/main/.planning/RESULTS.md).
+
+- **Short factual questions are its strength** (2.5–3.5), open description its weakest
+  (~1.5). That inverted our own assumption that this would be a describer.
+- **A third checkpoint made things worse.** A reasoning-SFT stage on top of the
+  instruction-tuned model scored 1.25 against its 2.17, roughly 3 SE. The reasoning
+  traces are undertrained for imagery and mostly add fluent hallucination, so that
+  checkpoint is **not published**. It is kept as a recorded negative result.
+- **Decoding presets are inside noise** — greedy 2.17 vs 1.88 for the other three.
+
+### Decoding
+
+Use **greedy**, which is the default. It is recommended for determinism, not because
+it beat anything. Pass `do_sample=True, temperature=...` to `chat()` if you want
+sampling anyway.
+
 
-## Training details
+## Training
 
-| | Stage 1 |
+| | |
 |---|---|
 | Data | LLaVA-CC3M-Pretrain-595K (595K image-caption pairs) |
 | Trainable params | ~2.7M (projector only) |
@@ -107,29 +154,26 @@ Stage 1 was trained on short image captions with no `<think>` traces. The proces
 | Learning rate | 1e-3, cosine, 250-step warmup |
 | Precision | bf16 |
 | Hardware | 1× H100 SXM (RunPod) |
-| Duration | ~5h |
-
-## Project scope & limitations
-
-Baguettotron-VLM is a **self-educational, open-source proof of concept** trained
-on a limited compute budget. The goal has always been to explore what an open,
-fully-reproducible, multilingual sub-1B VLM can do when every design choice has
-to justify itself against a hard budget cap.
-
-Concretely, this means the model may have limitations:
-
-- **Likely undertrained for VQA alignment.** The Stage 2 instruction-tuning
-  budget (~60h on a single H100) is small relative to the data mixture; the
-  model has seen far fewer tokens than typical VLMs of comparable size.
-- **Hallucinations.** The Baguettotron backbone is pretrained on synthetic
-  corpora and reasoning traces, which gives it strong `<think>` behaviour but
-  also a tendency to confabulate details that aren't grounded in the image,
-  especially on fine-grained or text-heavy questions.
-- **Best on simple, scene-level prompts.** In informal testing the model is
-  most reliable at open-ended description and scene understanding
-  ("what is happening in this image?") and less reliable at OCR, counting,
-  chart reading, or multi-hop visual reasoning. Quantitative benchmark
-  numbers will be published once the evaluation harness lands.
+| Duration | ~5h (~€15) |
+
+## Limitations
+
+- **No stop token.** `<|im_end|>` was masked out of the training loss, so the model
+  never learned to emit one — it spells the turn marker out as ordinary text instead.
+  `chat()` works around this by stopping on that text (a 512-token call drops from
+  ~21s to ~1.5s), but the defect is in the weights and needs a retrain to fix.
+- **Resolution ceiling.** One 448×448 crop → 256 visual tokens puts document text at
+  roughly 2–4 px/char. OCR, charts and documents are out of reach by architecture, not
+  by budget. Neither published checkpoint reads text in an image.
+- **Hallucinations**, especially on fine-grained or text-heavy questions.
+- **Undertrained for VQA alignment** — ~60h of instruction tuning on one H100 is small
+  relative to the data mixture.
+- **Multilingual capability is inherited, not verified.** The backbone covers six
+  languages; the VLM was never evaluated on non-English benchmarks.
+- **Evaluation is thin**: 6 held-out images, one judge, one rubric, 244 judged cells.
+  Every number here is directional.
+- **`transformers` 5.x cannot load this model.** The blocker is upstream in InternViT's
+  remote code, which predates v5's tied-weight API. Pin `transformers<5`.
 
 **Contributions and suggestions are very welcome** — issues, PRs, and ideas for
 better data mixes, training recipes, or evaluation setups are all appreciated.
@@ -146,7 +190,7 @@ If you use or extend Baguettotron-VLM in your research, please cite it:
   title        = {Baguettotron-VLM: An Open, Reproducible, Multilingual Sub-1B Vision-Language Model},
   author       = {Gemelli, Andrea},
   year         = {2026},
-  howpublished = {\url{https://huggingface.co/andreagemelli/Baguettotron-VLM}},
+  howpublished = {\url{https://huggingface.co/andreagemelli/baguettotron-vision-vqa}},
   note         = {Code: \url{https://github.com/andreagemelli/baguettotron-vlm}}
 }
 ```
@@ -154,4 +198,4 @@ If you use or extend Baguettotron-VLM in your research, please cite it:
 
 ## License
 
-Apache 2.0 — see the [GitHub repo](https://github.com/andreagemelli/baguettotron-vlm).
+Apache 2.0 — see the [GitHub repo](https://github.com/andreagemelli/baguettotron-vlm).

chat_template.jinja

@@ -0,0 +1,5 @@
+{% for m in messages %}<|im_start|>{{ m['role'] }}
+{{ m['content'] }}<|im_end|>
+{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant
+<think>
+{% endif %}

config.json

@@ -11,6 +11,13 @@
   "chat_style": "base",
   "stage": 1,
   "torch_dtype": "bfloat16",
+  "bos_token_id": 1,
+  "eos_token_id": [
+    2,
+    65492
+  ],
+  "pad_token_id": 3,
+  "tie_word_embeddings": false,
   "auto_map": {
     "AutoConfig": "configuration_baguettotron_vlm.BaguettotronVLMConfig",
     "AutoModelForImageTextToText": "modeling_baguettotron_vlm.BaguettotronVLMForConditionalGeneration",

generation_config.json

@@ -0,0 +1,12 @@
+{
+  "_from_model_config": false,
+  "bos_token_id": 1,
+  "eos_token_id": [
+    2,
+    65492
+  ],
+  "pad_token_id": 3,
+  "do_sample": false,
+  "max_new_tokens": 256,
+  "transformers_version": "4.57.0"
+}

modeling_baguettotron_vlm.py

@@ -8,12 +8,39 @@ from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     PreTrainedModel,
+    StoppingCriteria,
 )
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
 from .configuration_baguettotron_vlm import BaguettotronVLMConfig
 
 
+class StopOnTurnEnd(StoppingCriteria):
+    """Stop when the decoded tail contains a turn marker.
+
+    These weights never emit a stop *token*: <|im_end|> was masked out of the
+    training loss, so the model learned to spell the marker out as ordinary
+    text ("<|", "im", "_", "end", "|>") and then start a fresh turn. Passing
+    eos_token_id therefore does nothing — generation runs to max_new_tokens
+    every time, which is ~20x more tokens than the answer needs.
+
+    Matching on decoded text rather than token ids is deliberate: how the
+    marker splits depends on what precedes it (".<|" merges into a single
+    token after a period), so an id-sequence match misses many cases.
+    """
+
+    def __init__(self, tokenizer, markers=("<|im_end|>", "<|im_start|>"), window=8):
+        self.tokenizer = tokenizer
+        self.markers = markers
+        self.window = window
+
+    def __call__(self, input_ids: torch.Tensor, scores, **kwargs) -> bool:
+        tail = self.tokenizer.decode(
+            input_ids[0, -self.window:], skip_special_tokens=False
+        )
+        return any(marker in tail for marker in self.markers)
+
+
 class PixelUnshuffleProjector(nn.Module):
     """Reduces ViT tokens 4× via PixelUnshuffle then projects to LLM dim."""
 
@@ -55,37 +82,69 @@ class BaguettotronVLMForConditionalGeneration(PreTrainedModel):
 
     config_class = BaguettotronVLMConfig
     _no_split_modules = ["InternVisionEncoderLayer", "LlamaDecoderLayer"]
+    # Nothing is tied in this model — see __init__. Without this, the inherited
+    # Llama tied-weight bookkeeping drops llm.lm_head.weight during loading.
+    _tied_weights_keys: list[str] = []
     # Tell HF Trainer not to pass num_items_in_batch (loss handled internally)
     model_accepts_loss_kwargs: bool = False
 
     def __init__(self, config: BaguettotronVLMConfig):
         super().__init__(config)
 
-        self.vit = AutoModel.from_pretrained(
-            config.vit_model_id,
-            dtype=torch.bfloat16,
-            low_cpu_mem_usage=True,
-            trust_remote_code=True,
-        )
-        self.projector = PixelUnshuffleProjector(
-            in_dim=config.vit_hidden,
-            out_dim=config.llm_hidden,
-            factor=config.unshuffle_factor,
-        )
+        # Honour the dtype the caller asked for. from_pretrained(dtype=X) makes X
+        # the default dtype for the duration of __init__, and the projector picks
+        # it up automatically. Hard-coding bfloat16 for the submodules while the
+        # projector follows the default is what made dtype=torch.float32 fail with
+        # "mat1 and mat2 must have the same dtype".
+        dtype = torch.get_default_dtype()
+        if dtype not in (torch.float32, torch.float16, torch.bfloat16):
+            dtype = torch.bfloat16
 
-        tokenizer = AutoTokenizer.from_pretrained(config.llm_model_id)
-        tokenizer.add_special_tokens(
-            {"additional_special_tokens": [config.image_token, "</image>"]}
-        )
-        raw_id = tokenizer.convert_tokens_to_ids(config.image_token)
-        self.image_token_id: int = raw_id if isinstance(raw_id, int) else int(raw_id[0])
+        # The explicit CPU device context shields these nested from_pretrained
+        # calls from an outer meta-device init context (accelerate's device_map,
+        # and transformers >= 5), which otherwise aborts with "You are using
+        # from_pretrained with a meta device context manager".
+        with torch.device("cpu"):
+            self.vit = AutoModel.from_pretrained(
+                config.vit_model_id,
+                dtype=dtype,
+                low_cpu_mem_usage=True,
+                trust_remote_code=True,
+            )
+            self.projector = PixelUnshuffleProjector(
+                in_dim=config.vit_hidden,
+                out_dim=config.llm_hidden,
+                factor=config.unshuffle_factor,
+            )
 
-        self.llm = AutoModelForCausalLM.from_pretrained(
-            config.llm_model_id, dtype=torch.bfloat16
-        )
+            tokenizer = AutoTokenizer.from_pretrained(config.llm_model_id)
+            tokenizer.add_special_tokens(
+                {"additional_special_tokens": [config.image_token, "</image>"]}
+            )
+            raw_id = tokenizer.convert_tokens_to_ids(config.image_token)
+            self.image_token_id: int = (
+                raw_id if isinstance(raw_id, int) else int(raw_id[0])
+            )
+
+            self.llm = AutoModelForCausalLM.from_pretrained(
+                config.llm_model_id, dtype=dtype
+            )
         self.llm.resize_token_embeddings(len(tokenizer), mean_resizing=False)
-        # Break weight tying — safetensors rejects shared-storage tensors
+        # Break weight tying — safetensors rejects shared-storage tensors, and
+        # training updates lm_head independently of the input embeddings.
         self.llm.lm_head.weight = nn.Parameter(self.llm.lm_head.weight.data.clone())
+        # Keep it broken. Llama declares lm_head.weight as a tied key, so
+        # from_pretrained would skip it while loading and then re-tie it to
+        # embed_tokens — silently discarding the trained output head.
+        self.llm.config.tie_word_embeddings = False
+        self.llm._tied_weights_keys = []
+
+        # NOTE: transformers >= 5 cannot load this model, and the blocker is
+        # upstream: OpenGVLab's InternViT remote code predates v5's tied-weight
+        # API, so v5 aborts inside the nested AutoModel.from_pretrained above
+        # with "'InternVisionModel' object has no attribute
+        # 'all_tied_weights_keys'". Nothing can be patched from here — the model
+        # cards pin transformers<5 until InternViT is updated.
 
         self._tokenizer = tokenizer
 
@@ -158,17 +217,51 @@ class BaguettotronVLMForConditionalGeneration(PreTrainedModel):
                 -1, self.config.llm_hidden
             ).to(inputs_embeds.dtype)
 
-        im_end_id = int(self._tokenizer.convert_tokens_to_ids("<|im_end|>"))
+        # These weights end a turn by spelling out the literal text "<|im_end|>"
+        # as ordinary tokens and then emitting <|end_of_text|> — the special
+        # <|im_end|> token is never produced. <|end_of_text|> is therefore the id
+        # that actually stops generation; the others are listed for safety.
+        # PleIAs/Baguettotron's tokenizer declares no eos_token, so eos_token_id
+        # is None there and must be filtered out — an unfiltered [None, ...] makes
+        # generate() raise "'NoneType' object cannot be interpreted as an integer".
+        unk_id = self._tokenizer.unk_token_id
+        candidates = (
+            self._tokenizer.eos_token_id,
+            self._tokenizer.convert_tokens_to_ids("<|end_of_text|>"),
+            self._tokenizer.convert_tokens_to_ids("<|im_end|>"),
+        )
+        eos_ids = list(
+            dict.fromkeys(
+                tok_id
+                for tok_id in candidates
+                if tok_id is not None and tok_id != unk_id
+            )
+        )
+
+        # transformers' repetition-penalty processor corrupts the very first step
+        # on MPS: generate(inputs_embeds=...) starts from an empty input_ids, and
+        # the empty-index gather/scatter zeroes the entire logits row on Metal
+        # (it is a no-op on CPU and CUDA, as it should be). The result is a
+        # garbage first token that derails the whole answer.
+        if inputs_embeds.device.type == "mps" and repetition_penalty != 1.0:
+            repetition_penalty = 1.0
+
+        # do_sample defaults to greedy but can be overridden by callers
+        # (e.g. the inference sweep) without colliding on the keyword.
+        generate_kwargs.setdefault("do_sample", False)
+        generate_kwargs.setdefault(
+            "stopping_criteria", [StopOnTurnEnd(self._tokenizer)]
+        )
         output_ids = self.llm.generate(
             inputs_embeds=inputs_embeds,
             attention_mask=attention_mask,
             max_new_tokens=max_new_tokens,
-            do_sample=False,
             repetition_penalty=repetition_penalty,
-            eos_token_id=[self._tokenizer.eos_token_id, im_end_id],
+            eos_token_id=eos_ids,
             **generate_kwargs,
         )
         decoded = self._tokenizer.decode(output_ids[0], skip_special_tokens=False)
-        if "<|im_end|>" in decoded:
-            decoded = decoded[: decoded.index("<|im_end|>")]
+        for marker in ("<|im_end|>", "<|im_start|>", "<|end_of_text|>"):
+            if marker in decoded:
+                decoded = decoded[: decoded.index(marker)]
         return decoded.strip()

special_tokens_map.json

@@ -15,32 +15,5 @@
       "single_word": false
     }
   ],
-  "bos_token": {
-    "content": "<|begin_of_text|>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "eos_token": {
-    "content": "<|end_of_text|>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "pad_token": {
-    "content": "[PAD]",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "unk_token": {
-    "content": "[UNK]",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  }
+  "eos_token": "<|end_of_text|>"
 }

tokenizer.json

@@ -50,7 +50,7 @@
     },
     {
       "id": 65492,
-      "content": "<|im_end>",
+      "content": "<|im_end|>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -69,11 +69,11 @@
     {
       "id": 65494,
       "content": "</think>",
-      "single_word": false,
+      "single_word": true,
       "lstrip": false,
       "rstrip": false,
-      "normalized": false,
-      "special": true
+      "normalized": true,
+      "special": false
     },
     {
       "id": 65495,

tokenizer_config.json

@@ -41,7 +41,7 @@
       "special": true
     },
     "65492": {
-      "content": "<|im_end>",
+      "content": "<|im_end|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -59,10 +59,10 @@
     "65494": {
       "content": "</think>",
       "lstrip": false,
-      "normalized": false,
+      "normalized": true,
       "rstrip": false,
-      "single_word": false,
-      "special": true
+      "single_word": true,
+      "special": false
     },
     "65495": {
       "content": "source_1",
@@ -413,12 +413,9 @@
     "<image>",
     "</image>"
   ],
-  "bos_token": "<|begin_of_text|>",
   "clean_up_tokenization_spaces": true,
   "eos_token": "<|end_of_text|>",
   "extra_special_tokens": {},
   "model_max_length": 1000000000000000019884624838656,
-  "pad_token": "[PAD]",
-  "tokenizer_class": "PreTrainedTokenizerFast",
-  "unk_token": "[UNK]"
+  "tokenizer_class": "PreTrainedTokenizerFast"
 }