Fix generation from the published artifact; rewrite model card

README.md

@@ -16,8 +16,6 @@ tags:
 base_model:
 - OpenGVLab/InternViT-300M-448px-V2_5
 - PleIAs/Baguettotron
-datasets:
-- liuhaotian/LLaVA-CC3M-Pretrain-595K
 ---
 
 # baguettotron-internvit-alignment
@@ -38,9 +36,9 @@ Two checkpoints are published:
   — instruction-tuned on top of it. It goes past plain description and follows visual
   instructions, so prefer it for a richer chat experience.
 
-Resorces:
-- Code: [github.com/andreagemelli/baguettotron-vlm](https://github.com/andreagemelli/baguettotron-vlm)
-- Blogpost: [andreagemelli.me/posts/baguettotron-vlm](https://andreagemelli.me/posts/baguettotron-vlm/)
+Apache 2.0. Both live in the [Baguettotron-VLM collection](https://huggingface.co/collections/andreagemelli/baguettotron-vlm-69de37b4cab1960226e9c1f7).
+Source: [github.com/andreagemelli/baguettotron-vlm](https://github.com/andreagemelli/baguettotron-vlm) ·
+Write-up: [andreagemelli.me/posts/baguettotron-vlm](https://andreagemelli.me/posts/baguettotron-vlm/)
 
 
 > **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).
@@ -98,13 +96,22 @@ Greedy, prompt `Describe the image concisely.` — verbatim output. Images from
 |---|---|
 | <img src="examples/cats.jpg" width="180"> | `a cat is sleeping on the couch` |
 | <img src="examples/bear.jpg" width="180"> | `the bear is a good friend.` |
-| <img src="examples/stop.jpg" width="180"> | `a sign for a stop sign` |
+| <img src="examples/stop.jpg" width="180"> | `a sign for a stop` |
 | <img src="examples/bus.jpg" width="180"> | `the bus is a red double - decoration` |
 
 ### Chat template
 
 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").
 
+## Limitations
+
+- **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.
+- **Multilingual capability is inherited, not verified.** The backbone covers six
+  languages; the VLM was never evaluated on non-English benchmarks.
+
 > Tested against `transformers` 4.57. Newer major versions may need adjustments.
 
 **Contributions and suggestions are very welcome** — issues, PRs, and ideas for
@@ -130,4 +137,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).

modeling_baguettotron_vlm.py

@@ -7,14 +7,36 @@ from transformers import (
     AutoModel,
     AutoModelForCausalLM,
     AutoTokenizer,
+    LogitsProcessor,
     PreTrainedModel,
     StoppingCriteria,
 )
+from transformers.generation.logits_process import RepetitionPenaltyLogitsProcessor
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
 from .configuration_baguettotron_vlm import BaguettotronVLMConfig
 
 
+class SafeRepetitionPenalty(LogitsProcessor):
+    """Repetition penalty that is a no-op on the first decoding step.
+
+    generate(inputs_embeds=...) begins with an empty input_ids tensor. On MPS
+    the empty-index gather inside transformers' RepetitionPenaltyLogitsProcessor
+    zeroes the entire logits row rather than leaving it untouched (it is a
+    correct no-op on CPU and CUDA), which corrupts the first token and derails
+    the answer. There is nothing to penalise on that step anyway, so skipping it
+    removes the corruption while keeping the penalty for every later step.
+    """
+
+    def __init__(self, penalty: float):
+        self.inner = RepetitionPenaltyLogitsProcessor(penalty)
+
+    def __call__(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
+        if input_ids.shape[-1] == 0:
+            return scores
+        return self.inner(input_ids, scores)
+
+
 class StopOnTurnEnd(StoppingCriteria):
     """Stop when the decoded tail contains a turn marker.
 
@@ -237,25 +259,22 @@ class BaguettotronVLMForConditionalGeneration(PreTrainedModel):
             )
         )
 
-        # 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)]
         )
+        # The penalty goes through SafeRepetitionPenalty rather than generate()'s
+        # repetition_penalty kwarg, so the first decoding step is skipped.
+        if repetition_penalty != 1.0:
+            generate_kwargs.setdefault(
+                "logits_processor", [SafeRepetitionPenalty(repetition_penalty)]
+            )
         output_ids = self.llm.generate(
             inputs_embeds=inputs_embeds,
             attention_mask=attention_mask,
             max_new_tokens=max_new_tokens,
-            repetition_penalty=repetition_penalty,
             eos_token_id=eos_ids,
             **generate_kwargs,
         )