File size: 12,342 Bytes
1727faf 2c4010c 1727faf f7e9c38 1727faf 2c4010c 1727faf 2c4010c f7e9c38 1727faf b66671e 1727faf b66671e 1727faf b66671e 1727faf f7e9c38 1727faf f7e9c38 1727faf f7e9c38 1727faf f7e9c38 9b486d0 f7e9c38 1727faf f7e9c38 1727faf 9b486d0 1727faf 9b486d0 1727faf 9b486d0 1727faf 9b486d0 1727faf f7e9c38 2c4010c 1727faf f7e9c38 1727faf f7e9c38 1727faf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 | """BaguettotronVLM model — self-contained for HuggingFace Hub."""
from __future__ import annotations
import torch
import torch.nn as nn
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.
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."""
def __init__(self, in_dim: int, out_dim: int, factor: int):
super().__init__()
self.factor = factor
self.unshuffle = nn.PixelUnshuffle(factor)
self.mlp = nn.Sequential(
nn.Linear(in_dim * factor * factor, out_dim),
nn.GELU(),
nn.Linear(out_dim, out_dim),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
B, N, D = x.shape
spatial = int(N ** 0.5)
x = x.reshape(B, spatial, spatial, D).permute(0, 3, 1, 2)
x = self.unshuffle(x)
x = x.flatten(2).transpose(1, 2)
return self.mlp(x)
class BaguettotronVLMForConditionalGeneration(PreTrainedModel):
"""
BaguettotronVLM: InternViT-300M + PixelUnshuffle projector + Baguettotron-321M.
~628M total parameters. The same architecture is shipped for every
checkpoint; only the weights and `config.chat_style` differ.
Load with:
from transformers import AutoModelForImageTextToText
model = AutoModelForImageTextToText.from_pretrained(
"andreagemelli/baguettotron-vision-vqa",
trust_remote_code=True,
dtype=torch.bfloat16,
)
"""
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)
# 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
# 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,
)
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, 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
def _init_weights(self, module: nn.Module) -> None:
# Pretrained components are initialised from their respective hubs;
# the projector weights come from the saved checkpoint — skip random init.
pass
# ------------------------------------------------------------------
# Training interface
# ------------------------------------------------------------------
def forward(
self,
input_ids: torch.Tensor,
attention_mask: torch.Tensor,
labels: torch.Tensor | None = None,
pixel_values: torch.Tensor | None = None,
**kwargs,
) -> CausalLMOutputWithPast:
inputs_embeds = self.llm.get_input_embeddings()(input_ids)
if pixel_values is not None:
pixel_values = pixel_values.to(dtype=self.vit.dtype, device=self.vit.device)
with torch.no_grad():
vit_out = self.vit(pixel_values)
image_features = vit_out.last_hidden_state
if image_features.shape[1] == self.config.vit_tokens + 1:
image_features = image_features[:, 1:, :]
visual_tokens = self.projector(image_features.to(inputs_embeds.dtype))
image_mask = input_ids == self.image_token_id
inputs_embeds[image_mask] = visual_tokens.reshape(
-1, self.config.llm_hidden
).to(inputs_embeds.dtype)
return self.llm(
inputs_embeds=inputs_embeds,
attention_mask=attention_mask,
labels=labels,
return_dict=True,
use_cache=False,
)
# ------------------------------------------------------------------
# Inference interface
# ------------------------------------------------------------------
@torch.no_grad()
def chat(
self,
input_ids: torch.Tensor,
attention_mask: torch.Tensor,
pixel_values: torch.Tensor | None = None,
max_new_tokens: int = 256,
repetition_penalty: float = 1.3,
**generate_kwargs,
) -> str:
"""Inject visual tokens, generate autoregressively, return decoded string."""
inputs_embeds = self.llm.get_input_embeddings()(input_ids)
if pixel_values is not None:
pixel_values = pixel_values.to(dtype=self.vit.dtype, device=self.vit.device)
vit_out = self.vit(pixel_values)
image_features = vit_out.last_hidden_state
if image_features.shape[1] == self.config.vit_tokens + 1:
image_features = image_features[:, 1:, :]
visual_tokens = self.projector(image_features.to(inputs_embeds.dtype))
image_mask = input_ids == self.image_token_id
inputs_embeds[image_mask] = visual_tokens.reshape(
-1, self.config.llm_hidden
).to(inputs_embeds.dtype)
# 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
)
)
# 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,
eos_token_id=eos_ids,
**generate_kwargs,
)
decoded = self._tokenizer.decode(output_ids[0], skip_special_tokens=False)
for marker in ("<|im_end|>", "<|im_start|>", "<|end_of_text|>"):
if marker in decoded:
decoded = decoded[: decoded.index(marker)]
return decoded.strip()
|