this code was adapted from https github comlucidrainsflamingopytorch licensed under the mit license mit license c 2020 the google ai language team s the huggingface inc team and githublonepatient permission is hereby granted free of charge to any person obtaining a copy of this software and associated documentation files the software to deal in the software without restriction including without limitation the rights to use copy modify merge publish distribute sublicense andor sell copies of the software and to permit persons to whom the software is furnished to do so subject to the following conditions the above notice and this permission notice shall be included in all copies or substantial portions of the software the software is provided as is without warranty of any kind express or implied including but not limited to the warranties of merchantability fitness for a particular purpose and noninfringement in no event shall the s or holders be liable for any claim damages or other liability whether in an action of contract tort or otherwise arising from out of or in connection with the software or the use or other dealings in the software generic interface to various configurations of the perceiver resampler that simply takes in a series of potentially timeindexed contextual embeddings and resamples compresses them down to a prespecified number of latents note that the perceiver in general resamples based solely off the longrange context there s a nice opportunity here to prime the perceiver resampler with say a single layer s worth of language embeddings the target domain and use that to softly retrieve compress what we need this would be a novel contribution we should explore references deepmind s flamingo https www deepmind comblogtacklingmultipletaskswithasinglevisuallanguagemodel code borrowed w love from https github comlucidrainsflamingopytorch instantiates a perceiver resampler that operates over a sequence of embeddings say from a resnet or vit or mae of a given dimension performs depth blocks of crossattention with a fixed nlatents inputs then returns a tensor of shape bsz nlatents embeddim param embeddim dimensionality of embeddings being fed to the perceiver resampler also dimensionality of latent embeddings returned by the perceiver resampler could be e g vit embeddim resnet pool dim and so on args config ideficsconfig config object embeddim int the size of each embedding vector depth int depth of the perceiver resampler transformer w cross attention should be shallow 3 nheads int number of heads in each transformer block for multiheaded selfattention headdim int dimensionality of each head projection in the transformer block nlatents int number of latent embeddings to resample compress the input sequence to usually 128 create latents for perceiver create transformer blocks resample arbitrary length context compress down to self nlatents latent embeddings einsum repeatself latents seq embed bsz seq embed bszcontext shape0 latents self latents repeatcontext shape0 1 1 feed through perceiver attention blocks for attn ff in self blocks latents attncontext latents latents latents fflatents latents return self layernormlatents class ideficsperceiverattentionnn module def initself embeddim int nheads int headdim int qklayernorms bool none normalization scaling q k v projection no bias detail from perceiverflamingo papers runs perceiver selfattention with special context latents appended along the seq dimension args context torch tensor tensor of shape bsz seq embeddim representing longform context to resample latents torch tensor tensor of shape bsz nlatents embeddim representing fixed length latents to compress to returns torch tensor tensor of shape bsz nlatents embeddim representing attention over latents w cross from context query key value projections note that in flamingo latents are concatenated with context prior to attn note this results in queries w seq nlatents and keys values with seq lencontext nlatents multiheaded selfattention w stable softmax subtract perrow max amax before softmax call attn should be a 2d matrix of shape nlatents x context nlatents einsum rearrangex bsz seq heads embed bsz heads seq embed headsself nheads attend project back to output einsum rearrangeresampled bsz heads seq embed bsz seq heads embed headsself nheads simple mlp block with intermediatesize and embedding size super init self embeddim config visionconfig embeddim self ln nn layernormself embeddim self fc nn linearself embeddim intermediatesize biasfalse self act nn relu self cproj nn linearintermediatesize self embeddim biasfalse def forwardself hiddenstates optionaltupletorch floattensor torch floattensor hiddenstates self lnhiddenstates hiddenstates self fchiddenstates hiddenstates self acthiddenstates hiddenstates self cprojhiddenstates return hiddenstates this code was adapted from https github com lucidrains flamingo pytorch licensed under the mit license mit license c 2020 the google ai language team s the huggingface inc team and github lonepatient permission is hereby granted free of charge to any person obtaining a copy of this software and associated documentation files the software to deal in the software without restriction including without limitation the rights to use copy modify merge publish distribute sublicense and or sell copies of the software and to permit persons to whom the software is furnished to do so subject to the following conditions the above notice and this permission notice shall be included in all copies or substantial portions of the software the software is provided as is without warranty of any kind express or implied including but not limited to the warranties of merchantability fitness for a particular purpose and noninfringement in no event shall the s or holders be liable for any claim damages or other liability whether in an action of contract tort or otherwise arising from out of or in connection with the software or the use or other dealings in the software generic interface to various configurations of the perceiver resampler that simply takes in a series of potentially time indexed contextual embeddings and resamples compresses them down to a pre specified number of latents note that the perceiver in general resamples based solely off the long range context there s a nice opportunity here to prime the perceiver resampler with say a single layer s worth of language embeddings the target domain and use that to softly retrieve compress what we need this would be a novel contribution we should explore references deepmind s flamingo https www deepmind com blog tackling multiple tasks with a single visual language model code borrowed w love from https github com lucidrains flamingo pytorch instantiates a perceiver resampler that operates over a sequence of embeddings say from a resnet or vit or mae of a given dimension performs depth blocks of cross attention with a fixed n_latents inputs then returns a tensor of shape bsz n_latents embed_dim param embed_dim dimensionality of embeddings being fed to the perceiver resampler also dimensionality of latent embeddings returned by the perceiver resampler could be e g vit embed_dim resnet pool dim and so on args config ideficsconfig config object embed_dim int the size of each embedding vector depth int depth of the perceiver resampler transformer w cross attention should be shallow 3 n_heads int number of heads in each transformer block for multi headed self attention head_dim int dimensionality of each head projection in the transformer block n_latents int number of latent embeddings to resample compress the input sequence to usually 128 create latents for perceiver create transformer blocks resample arbitrary length context compress down to self n_latents latent embeddings einsum repeat self latents seq embed bsz seq embed bsz context shape 0 feed through perceiver attention blocks perceiver cross attention module let long form inputs be context resampled embeddings be latents normalization scaling q k v projection no bias detail from perceiver flamingo papers runs perceiver self attention with special context latents appended along the seq dimension args context torch tensor tensor of shape bsz seq embed_dim representing long form context to resample latents torch tensor tensor of shape bsz n_latents embed_dim representing fixed length latents to compress to returns torch tensor tensor of shape bsz n_latents embed_dim representing attention over latents w cross from context query key value projections note that in flamingo latents are concatenated with context prior to attn note this results in queries w seq n_latents and keys values with seq len context n_latents multiheaded self attention w stable softmax subtract per row max amax before softmax call attn should be a 2d matrix of shape n_latents x context n_latents einsum rearrange x bsz seq heads embed bsz heads seq embed heads self n_heads attend project back to output einsum rearrange resampled bsz heads seq embed bsz seq heads embed heads self n_heads simple mlp block with intermediate_size and embedding size
from typing import Optional, Tuple

import torch
import torch.nn as nn

from .configuration_idefics import IdeficsConfig


class IdeficsPerceiverResampler(nn.Module):
    def __init__(
        self, config: IdeficsConfig, embed_dim: int, depth: int, n_heads: int, head_dim: int, n_latents: int
    ) -> None:
        
        super().__init__()
        self.embed_dim, self.n_heads, self.head_dim, self.n_latents = embed_dim, n_heads, head_dim, n_latents
        self.qk_layer_norms = config.perceiver_config.qk_layer_norms_perceiver

        
        self.latents = nn.Parameter(torch.randn(self.n_latents, self.embed_dim), requires_grad=True)

        self.intermediate_dim = (
            self.embed_dim * 4
            if not hasattr(config.vision_config, "embed_dim")
            else config.vision_config.embed_dim * 4
        )
        
        self.blocks = nn.ModuleList(
            [
                nn.ModuleList(
                    [
                        IdeficsPerceiverAttention(self.embed_dim, self.n_heads, self.head_dim, self.qk_layer_norms),
                        IdeficsMLP(self.intermediate_dim, config),
                    ]
                )
                for _ in range(depth)
            ]
        )
        self.layer_norm = nn.LayerNorm(self.embed_dim)

    def forward(self, context: torch.Tensor) -> torch.Tensor:
        
        
        latents = self.latents.repeat(context.shape[0], 1, 1)

        
        for attn, ff in self.blocks:
            latents = attn(context, latents) + latents
            latents = ff(latents) + latents

        return self.layer_norm(latents)


class IdeficsPerceiverAttention(nn.Module):
    def __init__(self, embed_dim: int, n_heads: int, head_dim: int, qk_layer_norms: bool) -> None:
        
        super().__init__()
        self.embed_dim, self.n_heads, self.head_dim = embed_dim, n_heads, head_dim
        self.qk_layer_norms = qk_layer_norms
        
        self.context_layer_norm = nn.LayerNorm(self.embed_dim)
        self.latents_layer_norm = nn.LayerNorm(self.embed_dim)
        if self.qk_layer_norms:
            self.q_layer_norm = nn.LayerNorm(self.head_dim)
            self.k_layer_norm = nn.LayerNorm(self.head_dim)

        self.qk_scale = self.head_dim**-0.5

        
        self.q_proj = nn.Linear(self.embed_dim, self.n_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.embed_dim, self.n_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.embed_dim, self.n_heads * self.head_dim, bias=False)

        self.output_proj = nn.Linear(self.n_heads * self.head_dim, embed_dim, bias=False)

    def forward(self, context: torch.Tensor, latents: torch.Tensor) -> torch.Tensor:
        
        context = self.context_layer_norm(context)
        latents = self.latents_layer_norm(latents)
        batch_size, seq_length, embed_dim = context.shape[:3]

        
        
        q = self.q_proj(latents)
        k = self.k_proj(torch.cat([context, latents], dim=-2))
        v = self.v_proj(torch.cat([context, latents], dim=-2))

        
        
        
        q, k, v = [x.reshape(batch_size, x.shape[1], self.n_heads, self.head_dim).transpose(1, 2) for x in (q, k, v)]

        if self.qk_layer_norms:
            q = self.q_layer_norm(q)
            k = self.k_layer_norm(k)

        scores = torch.einsum("... i d, ... j d -> ... i j", q * self.qk_scale, k)
        stabilized_scores = scores - (scores.amax(dim=-1, keepdim=True).detach())
        attn = stabilized_scores.softmax(dim=-1)

        
        resampled = torch.einsum("... i j, ... j d -> ... i d", attn, v)
        
        return self.output_proj(resampled.transpose(1, 2).flatten(-2))


class IdeficsMLP(nn.Module):
    def __init__(self, intermediate_size, config: IdeficsConfig):
        
        super().__init__()
        self.embed_dim = config.vision_config.embed_dim
        self.ln = nn.LayerNorm(self.embed_dim)
        self.fc = nn.Linear(self.embed_dim, intermediate_size, bias=False)
        self.act = nn.ReLU()
        self.c_proj = nn.Linear(intermediate_size, self.embed_dim, bias=False)

    def forward(self, hidden_states: Optional[Tuple[torch.FloatTensor]]) -> torch.FloatTensor:
        hidden_states = self.ln(hidden_states)
        hidden_states = self.fc(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states = self.c_proj(hidden_states)

        return hidden_states