Building a nanoVLM — Solutions

def __init__(self, cfg):
    ...
    self.proj = nn.Linear(self.input_dim, self.output_dim, bias=False)

def pixel_shuffle(self, x):
    bsz, seq, embed_dim = x.size()
    seq_root = int(seq ** 0.5)
    assert seq_root ** 2 == seq
    assert seq_root % self.scale_factor == 0

    height = width = seq_root
    x = x.view(bsz, height, width, embed_dim)
    h_out = height // self.scale_factor
    w_out = width // self.scale_factor

    x = x.reshape(bsz, h_out, self.scale_factor, w_out, self.scale_factor, embed_dim)
    x = x.permute(0, 1, 3, 2, 4, 5).contiguous()
    x = x.reshape(bsz, h_out * w_out, embed_dim * self.scale_factor ** 2)
    return x

def forward(self, x):
    x = self.pixel_shuffle(x)
    x = self.proj(x)
    return x

def __call__(self, batch):
    images  = [item["image"]     for item in batch]
    texts   = [item["text_data"] for item in batch]
    answers = [item["answer"]    for item in batch]

    # Step 1
    images = torch.stack(images)

    # Step 2
    input_sequences = []
    for i in range(len(texts)):
        input_sequences.append(f"{texts[i]}{answers[i]}")

    # Step 3
    encoded_full_sequences = self.tokenizer.batch_encode_plus(
        input_sequences,
        padding="max_length",
        padding_side="left",
        return_tensors="pt",
        truncation=True,
        max_length=self.max_length,
    )
    input_ids      = encoded_full_sequences["input_ids"]
    attention_mask = encoded_full_sequences["attention_mask"]

    # Step 4
    labels = input_ids.clone()
    labels[:, :-1] = input_ids[:, 1:].clone()
    labels[:, -1]  = -100

    # Step 5
    original_lengths = [len(self.tokenizer.encode(seq)) for seq in input_sequences]

    for i in range(len(batch)):
        question_length = len(self.tokenizer.encode(texts[i], add_special_tokens=False))

        if original_lengths[i] > self.max_length:   # Case A: truncated
            labels[i, :] = -100
            continue

        first_token_pos = attention_mask[i].nonzero(as_tuple=True)[0][0].item()
        question_end    = first_token_pos + question_length - 1   # -1 for shift
        labels[i, :question_end] = -100

    return {"image": images, "input_ids": input_ids,
            "attention_mask": attention_mask, "labels": labels}

def forward(self, input_ids, image, attention_mask=None, targets=None):

    image_embeds = self.vision_encoder(image)            # (B, 196, 768)
    image_embeds = self.MP(image_embeds)                 # (B, 49,  576)

    text_embeds = self.decoder.token_embedding(input_ids)  # (B, T, 576)

    combined_embeds = torch.cat((image_embeds, text_embeds), dim=1)  # (B, 49+T, 576)

    if attention_mask is not None:
        B, img_len = image_embeds.size(0), image_embeds.size(1)
        image_attention = torch.ones((B, img_len),
                                      device=attention_mask.device,
                                      dtype=attention_mask.dtype)
        attention_mask = torch.cat((image_attention, attention_mask), dim=1)

    output_token_embeddings = self.decoder(combined_embeds, attention_mask)

    loss = None
    if targets is not None:
        logits = self.decoder.head(output_token_embeddings)   # (B, 49+T, vocab)
        logits = logits[:, image_embeds.size(1):, :]          # (B, T,    vocab)
        loss   = F.cross_entropy(logits.reshape(-1, logits.size(-1)),
                                  targets.reshape(-1), ignore_index=-100)

    return logits, loss

@torch.no_grad()
def generate(self, input_ids, image, attention_mask=None, max_new_tokens=5):

    image_embd = self.vision_encoder(image)          # (B, 196, 768)
    image_embd = self.MP(image_embd)                 # (B, 49,  576)

    token_embd     = self.decoder.token_embedding(input_ids)
    combined_embed = torch.cat((image_embd, token_embd), dim=1)

    batch_size, img_seq_len = image_embd.size(0), image_embd.size(1)
    if attention_mask is not None:
        image_attention_mask = torch.ones((batch_size, img_seq_len),
                                           device=attention_mask.device, dtype=attention_mask.dtype)
        attention_mask = torch.cat((image_attention_mask, attention_mask), dim=1)

    outputs          = combined_embed
    generated_tokens = torch.zeros((batch_size, max_new_tokens),
                                    device=input_ids.device, dtype=input_ids.dtype)

    for i in range(max_new_tokens):
        model_out         = self.decoder(outputs, attention_mask)
        last_token_logits = model_out[:, -1, :]

        if not self.decoder.lm_use_tokens:
            last_token_logits = self.decoder.head(last_token_logits)

        probs      = torch.softmax(last_token_logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)
        generated_tokens[:, i] = next_token.squeeze(-1)

        generated_embed = self.decoder.token_embedding(next_token)
        outputs         = torch.cat((outputs, generated_embed), dim=1)

        if attention_mask is not None:
            attention_mask = torch.cat(
                (attention_mask, torch.ones((batch_size, 1), device=attention_mask.device)), dim=1)

        if next_token.item() == 2:   # EOS
            break

    return generated_tokens