model.py

import torch
from transformers import AutoTokenizer, CLIPProcessor

from medpalm.transformer import (
    AutoregressiveWrapper,
    Decoder,
    Encoder,
    Transformer,
    ViTransformerWrapper,
)


class MedPalmTokenizer:
    def __init__(self):
        try:
            self.processor = CLIPProcessor.from_pretrained(
                "laion/CLIP-ViT-L-14-laion2B-s32B-b82K"
            )
            self.tokenizer = AutoTokenizer.from_pretrained(
                "EleutherAI/gpt-neox-20b",
                additional_special_tokens=["<image>", "</image>"],
                eos_token="<eos>",
                pad_token="<pad>",
                extra_ids=0,
                model_max_length=8192,
            )

            self.im_idx, self.im_end_idx = (
                self.tokenizer.convert_tokens_to_ids(
                    ["<image>", "</image>"]
                )
            )
        except Exception as e:
            print(f"Error init tokenizer: {e}")

    def tokenize_texts(self, texts):
        try:
            texts = self.tokenizer(
                texts,
                return_tensors="pt",
                padding=True,
                truncation=True,
            ).input_ids
            image_tokens = torch.tensor(
                [[self.im_idx, self.im_end_idx]] * texts.shape[0]
            )
            return (
                torch.cat(
                    [texts[:, 0:1], image_tokens, texts[:, 1:]], dim=1
                ),
                texts,
            )
        except Exception as e:
            print(f"Error tokenizing texts: {e}")

    def tokenize_images(self, images):
        try:
            tokenized_images = self.processor(
                images=images, return_tensors="pt"
            ).pixel_values
            print(f"Tokenized image: {tokenized_images.shape}")
            return tokenized_images

        except Exception as e:
            print(f"Error tokenizing texts: {e}")

    def tokenize(self, sample):
        try:
            text_tokens, only_text_tokens = self.tokenize_texts(
                sample["target_text"]
            )
            attention_mask = (
                text_tokens != self.tokenizer.pad_token_id
            )
            dummy_image_features = torch.ones(
                (text_tokens.shape[0], 64)
            )
            attention_mask = torch.cat(
                [dummy_image_features, attention_mask], dim=1
            )
            return {
                "text_tokens": text_tokens,
                "images": self.tokenize_images(sample["image"]),
                "labels": only_text_tokens,
                "attention_mask": attention_mask,
            }

        except Exception as e:
            print(f"Error during tokenization {e}")


class MedPalm(torch.nn.Module):
    """
    MedPalm model for medical image and text processing.

    Args:
        image_size (int): Size of the input image (default: 256).
        patch_size (int): Size of each image patch (default: 32).
        encoder_dim (int): Dimensionality of the encoder (default: 512).
        encoder_depth (int): Number of encoder layers (default: 6).
        encoder_heads (int): Number of attention heads in the encoder (default: 8).
        num_tokens (int): Number of tokens in the decoder (default: 20000).
        max_seq_len (int): Maximum sequence length in the decoder (default: 1024).
        decoder_dim (int): Dimensionality of the decoder (default: 512).
        decoder_depth (int): Number of decoder layers (default: 6).
        decoder_heads (int): Number of attention heads in the decoder (default: 8).
        alibi_num_heads (int): Number of attention heads in the alibi mechanism (default: 4).
        use_abs_pos_emb (bool): Whether to use absolute positional embeddings (default: False).
        cross_attend (bool): Whether to enable cross-attention in the decoder (default: True).
        alibi_pos_bias (bool): Whether to use positional bias in the alibi mechanism (default: True).
        rotary_xpos (bool): Whether to use rotary positional embeddings (default: True).
        attn_flash (bool): Whether to use attention flash in the decoder (default: True).
        qk_norm (bool): Whether to normalize the query-key vectors in attention (default: True).
    """

    def __init__(
        self,
        image_size=256,
        patch_size=32,
        encoder_dim=512,
        encoder_depth=6,
        encoder_heads=8,
        num_tokens=20000,
        max_seq_len=1024,
        decoder_dim=512,
        decoder_depth=6,
        decoder_heads=8,
        alibi_num_heads=4,
        use_abs_pos_emb=False,
        cross_attend=True,
        alibi_pos_bias=True,
        rotary_xpos=True,
        attn_flash=True,
        qk_norm=True,
    ):
        super(MedPalm, self).__init__()

        self.encoder = ViTransformerWrapper(
            image_size=image_size,
            patch_size=patch_size,
            attn_layers=Encoder(
                dim=encoder_dim,
                depth=encoder_depth,
                heads=encoder_heads,
            ),
        )

        self.decoder = Transformer(
            num_tokens=num_tokens,
            max_seq_len=max_seq_len,
            use_abs_pos_emb=use_abs_pos_emb,
            attn_layers=Decoder(
                dim=decoder_dim,
                depth=decoder_depth,
                heads=decoder_heads,
                cross_attend=cross_attend,
                alibi_pos_bias=alibi_pos_bias,
                alibi_num_heads=alibi_num_heads,
                rotary_xpos=rotary_xpos,
                attn_flash=attn_flash,
                qk_norm=qk_norm,
            ),
        )

        self.decoder = AutoregressiveWrapper(self.decoder)

    def forward(self, img, text):
        """
        Forward pass of the MedPalm model.

        Args:
            img (torch.Tensor): Input image tensor.
            text (torch.Tensor): Input text tensor.

        Returns:
            torch.Tensor: Output tensor from the decoder.
        """
        try:
            encoded = self.encoder(img, return_embeddings=True)
            return self.decoder(text, context=encoded)
        except Exception as error:
            print(f"Failed in forward method: {error}")
            raise