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)