model.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPVisionModel

PAD_TOKEN_ID = 50256
IGNORE_INDEX = -100
NUM_IMAGE_PATCHES = 256
VISION_PROJECTOR_HIDDEN_DIM = 2048
VISION_ENCODER = "openai/clip-vit-large-patch14"
VOCAB_SIZE = 50257


class Config:
    input_sequence_length: int = 1024
    embedding_dim: int = 1280
    hidden_dim: int = 5120
    num_attention_heads: int = 10
    layer_count: int = 20
    num_relative_positions: int = 1024
    device_type: str = "cuda" if torch.cuda.is_available() else "cpu"


def config_from_checkpoint(checkpoint_config=None) -> Config:
    config = Config()
    if isinstance(checkpoint_config, dict):
        for key in [
            "input_sequence_length",
            "embedding_dim",
            "hidden_dim",
            "num_attention_heads",
            "layer_count",
            "num_relative_positions",
        ]:
            if key in checkpoint_config:
                setattr(config, key, checkpoint_config[key])
    config.device_type = "cuda" if torch.cuda.is_available() else "cpu"
    return config


class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size, embedding_dim):
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, embedding_dim)

    def embed(self, input_indices):
        return self.token_embedding_table(input_indices)


class RelativePositionEmbedding(nn.Module):
    def __init__(self, num_relative_positions: int):
        super().__init__()
        self.num_relative_positions = num_relative_positions
        self.bias_embedding_table = nn.Embedding(num_relative_positions, 1)

    def forward(self, query_len, key_len, device_type=None):
        query_positions = torch.arange(query_len, device=device_type)[:, None]
        key_positions = torch.arange(key_len, device=device_type)[None, :]
        relative_position_matrix = query_positions - key_positions
        clamped_relative_position_matrix = relative_position_matrix.clamp(
            min=0, max=self.num_relative_positions - 1
        )
        return self.bias_embedding_table(clamped_relative_position_matrix).squeeze(-1)


class AttentionHead(nn.Module):
    def __init__(self, head_size, config):
        super().__init__()
        self.key_fc = nn.Linear(config.embedding_dim, head_size, bias=False)
        self.query_fc = nn.Linear(config.embedding_dim, head_size, bias=False)
        self.value_fc = nn.Linear(config.embedding_dim, head_size, bias=False)
        self.head_size = head_size
        self.relative_position_embedding_layer = RelativePositionEmbedding(
            num_relative_positions=config.num_relative_positions
        )

    def forward(self, input_tensor):
        _, token_len, _ = input_tensor.shape
        key = self.key_fc(input_tensor)
        query = self.query_fc(input_tensor)
        value = self.value_fc(input_tensor)
        attention_scores = query @ key.transpose(-2, -1) * self.head_size ** (-0.5)
        attention_scores = attention_scores + self.relative_position_embedding_layer(
            token_len, token_len, device_type=input_tensor.device
        )
        mask = torch.triu(torch.ones(token_len, token_len, device=input_tensor.device), diagonal=1)
        attention_scores = attention_scores.masked_fill(mask == 1, float("-inf"))
        attention_weights = F.softmax(attention_scores, dim=-1)
        return attention_weights @ value


class MultiHeadAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.num_attention_heads = config.num_attention_heads
        self.embedding_dim = config.embedding_dim
        self.head_size = int(config.embedding_dim / config.num_attention_heads)
        self.attention_heads = nn.ModuleList(
            [AttentionHead(self.head_size, config) for _ in range(config.num_attention_heads)]
        )
        self.output_projection = nn.Linear(config.embedding_dim, config.embedding_dim)

    def forward(self, input_tensor):
        head_outputs = [head(input_tensor) for head in self.attention_heads]
        return self.output_projection(torch.cat(head_outputs, dim=-1))


class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(config.embedding_dim, config.hidden_dim),
            nn.ReLU(),
            nn.Linear(config.hidden_dim, config.embedding_dim),
        )

    def forward(self, input_tensor):
        return self.net(input_tensor)


class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.layer_norm1 = nn.LayerNorm(config.embedding_dim)
        self.layer_norm2 = nn.LayerNorm(config.embedding_dim)
        self.multihead_attention = MultiHeadAttention(config)
        self.feed_forward = FeedForward(config)

    def forward(self, input_tensor):
        attention_output = self.multihead_attention(self.layer_norm1(input_tensor))
        residual_attention = attention_output + input_tensor
        feedforward_output = self.feed_forward(self.layer_norm2(residual_attention))
        return feedforward_output + residual_attention


class VocabularyLogits(nn.Module):
    def __init__(self, vocab_size, config):
        super().__init__()
        self.output_norm = nn.LayerNorm(config.embedding_dim)
        self.vocab_projection = nn.Linear(config.embedding_dim, vocab_size)

    def forward(self, transformer_block_output):
        return self.vocab_projection(self.output_norm(transformer_block_output))


class EveryonesGPT(nn.Module):
    def __init__(self, vocab_size, config):
        super().__init__()
        self.config = config
        self.token_embedding_layer = TokenEmbedding(vocab_size, config.embedding_dim)
        self.blocks = nn.Sequential(*[TransformerBlock(config) for _ in range(config.layer_count)])
        self.vocab_projection = VocabularyLogits(vocab_size, config)
        self.criterion = nn.CrossEntropyLoss(ignore_index=IGNORE_INDEX)

    def forward(self, input_indices, target_indices=None):
        token_embeddings = self.token_embedding_layer.embed(input_indices)
        blocks_output = self.blocks(token_embeddings)
        logits = self.vocab_projection(blocks_output)
        if target_indices is None:
            return logits, None
        batch_size, token_len, vocab_size = logits.shape
        loss = self.criterion(logits.view(batch_size * token_len, vocab_size), target_indices.view(batch_size * token_len))
        return logits, loss


class VLM(nn.Module):
    def __init__(self, llm, vision_encoder=VISION_ENCODER, projector_hidden_dim=VISION_PROJECTOR_HIDDEN_DIM):
        super().__init__()
        self.llm = llm
        self.vision = CLIPVisionModel.from_pretrained(vision_encoder)
        self.projector = nn.Sequential(
            nn.Linear(self.vision.config.hidden_size, projector_hidden_dim),
            nn.ReLU(),
            nn.Linear(projector_hidden_dim, llm.config.embedding_dim),
        )
        self.loss_fn = nn.CrossEntropyLoss(ignore_index=IGNORE_INDEX)

    def make_multimodal_embeddings(self, image_batch, input_token_ids):
        text_embeddings = self.llm.token_embedding_layer.embed(input_token_ids)
        vision_outputs = self.vision(image_batch, output_hidden_states=True)
        image_features = vision_outputs.hidden_states[-1][:, 1:]
        image_embeddings = self.projector(image_features)
        return torch.cat([image_embeddings, text_embeddings[:, NUM_IMAGE_PATCHES:]], dim=1)

    def forward_logits(self, image_batch, input_token_ids):
        multimodal_embeddings = self.make_multimodal_embeddings(image_batch, input_token_ids)
        blocks_output = self.llm.blocks(multimodal_embeddings)
        return self.llm.vocab_projection(blocks_output)

    def forward(self, image_batch, input_token_ids, target_labels):
        logits = self.forward_logits(image_batch, input_token_ids)
        batch_size, token_len, vocab_size = logits.shape
        loss = self.loss_fn(logits.view(batch_size * token_len, vocab_size), target_labels.view(batch_size * token_len))
        return logits, loss


def build_model(config=None, vision_encoder=VISION_ENCODER, projector_hidden_dim=VISION_PROJECTOR_HIDDEN_DIM):
    config = config or Config()
    llm = EveryonesGPT(vocab_size=VOCAB_SIZE, config=config)
    return VLM(llm, vision_encoder=vision_encoder, projector_hidden_dim=projector_hidden_dim)