Pushing deployment to space

CNN_encoder.py

@@ -0,0 +1,70 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from utility import  load_model,get_layers
+import tensorflow as tf
+from tensorflow.keras.models import Model # type: ignore
+import sys
+
+
+class CNN_Encoder(nn.Module):
+    def __init__(self,model_path,model_name,pop_conv_layers,encoder_layers,tags_threshold,tags_embeddings=None,finetune_visual_model=False,num_tags=105):
+        super(CNN_Encoder,self).__init__()
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        if tags_embeddings is not None:
+            # Initialize embeddings and move them to the device
+            self.tags_embeddings = nn.Parameter(torch.tensor(tags_embeddings, dtype=torch.float32).to(self.device), requires_grad=True)
+        else:
+            # Initialize embeddings with ones and move them to the device
+            self.tags_embeddings = nn.Parameter(torch.ones((num_tags, 400), dtype=torch.float32).to(self.device), requires_grad=True)
+
+
+        self.tags_threshold=tags_threshold
+        # visual_model.children() gets an iterator over child modules(layers) of the model (pretrained chexnet)
+        #list* => converts this iterator into a list of layers so it is easier to manipulate
+        # if pop_conv_layers is, it removes the last layer from the model
+        # nn.sequentail=> creates a new model that consists  of all the layers up to the pop_conv_layers, stacks layers in a linear sequence
+        visual_model=load_model(model_path,model_name)
+        self.visual_model = Model(inputs=visual_model.input,
+                                        outputs=[visual_model.output, visual_model.layers[-pop_conv_layers - 1].output],
+                                        trainable=finetune_visual_model)
+
+        self.encoder_layers=get_layers(encoder_layers,'relu')
+
+    def get_visual_features(self, images):
+        images_np = images.cpu().numpy()
+        images_tf = tf.convert_to_tensor(images_np)
+        images_tf = tf.transpose(images_tf, perm=[0,2,3,1 ])
+        predictions, visual_features = self.visual_model(images_tf)
+        predictions = torch.tensor(predictions.numpy(), device=self.device, requires_grad=True)
+        visual_features = torch.tensor(visual_features.numpy(), device=self.device, requires_grad=True)
+
+        predictions = predictions.view(predictions.size(0), predictions.size(-1), -1)
+        visual_features = visual_features.view(visual_features.size(0), -1, visual_features.size(-1))
+        
+        if self.tags_threshold >= 0:
+            predictions = (predictions >= self.tags_threshold).float()
+        
+        return predictions, visual_features
+
+    def forward(self, images):
+        # print python version
+        print("Python version:", sys.version)
+
+        images=images.to(self.device)
+        tags_predictions, visual_features = self.get_visual_features(images)
+        if tags_predictions is not None:
+            tags_predictions=tags_predictions.to(self.device)
+            self.tags_embeddings = self.tags_embeddings.to(self.device)
+            tags_embed = tags_predictions * self.tags_embeddings 
+        else:
+            tags_embed=None
+
+        for layer in self.encoder_layers:
+            visual_features = layer(visual_features)
+            if tags_embed is not None:
+                tags_embed = layer(tags_embed)
+
+        return visual_features, tags_embed
+
+

Dockerfile

@@ -0,0 +1,13 @@
+FROM python:3.8.8-slim
+
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+WORKDIR /app
+
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user . /app
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

configs.py

@@ -0,0 +1,72 @@
+from tags import tags
+
+
+class argHandler(dict):
+    __getattr__ = dict.get
+    __setattr__ = dict.__setitem__
+    __delattr__ = dict.__delitem__
+    _descriptions = {'help, --h, -h': 'show this super helpful message and exit'}
+    def setDefaults(self):
+        self.define('train_csv', './IU-XRay/training_set.csv',
+                    'path to training csv containing the images names and the labels')
+        self.define('test_csv', './IU-XRay/testing_set.csv',
+                    'path to testing csv containing the images names and the labels')
+        self.define('all_data_csv', './IU-XRay/all_data.csv',
+                    'path to all data csv containing the images names and the labels')
+        self.define('image_directory', './IU-XRay/images',
+                    'path to folder containing the patient folders which containing the images')
+        self.define('output_images_folder', './outputs/CDGPT2',
+                    'path to folder containing output images')
+        self.define('data_dir', './IU-XRay',
+                    'path to folder containing the patient folders which containing the images')
+        # self.define('visual_model_name', 'fine_tuned_chexnet',
+        self.define('visual_model_name', 'fine_tuned_chexnet',
+                    'path to folder containing the patient folders which containing the images')
+        self.define('finetune_visual_model', False,
+                    'option if you want to finetune the visual model')
+        self.define('visual_model_pop_layers', 2,
+                    'number of conv layers to pop to get visual features')
+        self.define('csv_label_columns', ['Caption'], 'the name of the label columns in the csv')
+        self.define('image_target_size', (224, 224, 3), 'the target size to resize the image')
+
+        self.define('max_sequence_length', 200,
+                    'Maximum number of words in a sentence')
+        self.define('num_epochs', 100, 'maximum number of epochs')
+        self.define('encoder_layers', [0.4], 'a list describing the hidden layers of the encoder. Example [10,0.4,5] will create a hidden layer with size 10 then dropout wth drop prob 0.4, then hidden layer with size 5. If empty it will connect to output nodes directly.')
+        self.define('classifier_layers', [0.4], 'a list describing the hidden layers of the encoder. Example [10,0.4,5] will create a hidden layer with size 10 then dropout wth drop prob 0.4, then hidden layer with size 5. If empty it will connect to output nodes directly.')
+        self.define('tags_threshold', -1,
+                    'The threshold from which to detect a tag. -1 will multiply the tags embeddings according to prediction')
+        self.define('tokenizer_vocab_size', 1001,
+                    'The number of words to tokinze, the rest will be set as <unk>')
+        self.define('batch_size', 16, 'batch size for training and testing')
+        self.define('generator_workers', 8, 'The number of cpu workers generating batches.')
+        self.define('beam_width', 7, 'The beam search width during evaluation')
+        self.define('epochs_to_evaluate', 3, 'The number of epochs to train before evaluating on the test set.')
+
+        self.define('generator_queue_length', 24, 'The maximum number of batches in the queue to be trained on.')
+
+        self.define('ckpt_path', './checkpoints/CDGPT2/',
+                    'where to save the checkpoints. The path will be created if it does not exist. The system saves every epoch by default')
+        self.define('continue_from_last_ckpt', True,
+                    'continue training from last ckpt or not')
+        self.define('calculate_loss_after_epoch', False,
+                    'if True it will calculate the train and test loss by passing over the data again and append it to losses_csv')
+        self.define('learning_rate', 1e-3, 'The optimizer learning rate')
+        self.define('optimizer_type', 'Adam', 'Choose from (Adam, SGD, RMSprop, Adagrad, Adadelta, Adamax, Nadam)')
+        self.define('tags', tags,
+                    'the names of the tags')
+
+    def define(self, argName, default, description):
+        self[argName] = default
+        self._descriptions[argName] = description
+
+    def help(self):
+        print('Arguments:')
+        spacing = max([len(i) for i in self._descriptions.keys()]) + 2
+        for item in self._descriptions:
+            currentSpacing = spacing - len(item)
+            print('  --' + item + (' ' * currentSpacing) + self._descriptions[item])
+        exit()
+
+
+        

distil_gpt2.py

@@ -0,0 +1,680 @@
+from dataclasses import dataclass
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import math
+from transformers import GPT2Tokenizer
+import tiktoken
+from transformers import GPT2LMHeadModel
+from transformers import PretrainedConfig
+
+
+@dataclass
+class GPTConfig(PretrainedConfig):
+    visual_size: int = 1024
+    vocab_size: int = 50257
+    block_size: int = 1024
+    tags_embd: int = 400
+    n_embd: int = 768
+    n_layer: int = 6
+    n_head: int = 12
+
+    def __init__(self,**kwargs):
+        super().__init__(**kwargs)
+        self.hidden_size = self.n_embd
+
+
+class CasualSelfAttention(nn.Module):
+    def __init__(self, config: GPTConfig):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        self.c_attn = nn.Linear(config.n_embd, config.n_embd * 3)
+        self.visual_attn = nn.Linear(config.visual_size, config.n_embd * 2)
+        self.tags_attn = nn.Linear(config.tags_embd, config.n_embd * 2)
+    
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
+        self.n_head = config.n_head
+        self.n_embed = config.n_embd 
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        self.register_buffer(
+            'bias', torch.tril(torch.ones(199, 353))
+            .view(1, 1, 199, 353)
+        )
+    
+    def forward(self, x: torch.Tensor, visual_features: torch.Tensor = None, tags_embedding: torch.Tensor = None) -> torch.Tensor:
+
+        B, T, C = x.size() 
+        visual_features=visual_features.to(self.device)
+        tags_embedding=tags_embedding.to(self.device)
+
+
+        qkv = self.c_attn(x) # the error happens here
+        
+        q, k, v = qkv.split(self.n_embed, dim=2) 
+        q = q.view(B, T, self.n_head, self.n_embed // self.n_head).transpose(1, 2) 
+        k = k.view(B, T, self.n_head, self.n_embed // self.n_head).transpose(1, 2) 
+        v = v.view(B, T, self.n_head, self.n_embed // self.n_head).transpose(1, 2)
+        # Handle visual input if provided
+        if visual_features is not None: 
+            visual_kv = self.visual_attn(visual_features) 
+            visual_k, visual_v = visual_kv.split(self.n_embed, dim=2)
+            visual_k = visual_k.view(B, visual_features.size(1), self.n_head, self.n_embed // self.n_head).transpose(1, 2) 
+            visual_v = visual_v.view(B, visual_features.size(1), self.n_head, self.n_embed // self.n_head).transpose(1, 2)
+
+            k = torch.cat([k, visual_k], dim=-2) 
+            v = torch.cat([v, visual_v], dim=-2)
+        
+        if tags_embedding is not None:
+            tags_kv = self.tags_attn(tags_embedding)  
+            tags_k, tags_v = tags_kv.split(self.n_embed, dim=2)
+            tags_k = tags_k.view(B, tags_embedding.size(1), self.n_head, self.n_embed // self.n_head).transpose(1, 2) 
+            tags_v = tags_v.view(B, tags_embedding.size(1), self.n_head, self.n_embed // self.n_head).transpose(1, 2)
+
+            k = torch.cat([k, tags_k], dim=-2) 
+            v = torch.cat([v, tags_v], dim=-2)
+
+        # Causal self-attention computation
+        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) 
+        device = att.device
+        query_seq_len, key_seq_len = T, k.size(-2)
+        
+        # Text can attend to: previous text + all visual/tag tokens
+        text_mask = torch.tril(torch.ones(T, T, device=device))  # Text-to-text causal
+        non_text_mask = torch.ones(T, key_seq_len - T, device=device)  # Text-to-other full
+        combined_mask = torch.cat([text_mask, non_text_mask], dim=1)
+        
+        # Reshape for broadcasting
+        combined_mask = combined_mask.view(1, 1, T, key_seq_len)
+        att = att.masked_fill(combined_mask == 0, float('-inf'))
+
+
+
+        att = F.softmax(att, dim=-1)
+        visual_att = att[..., :T, T:].mean().item()  # Text → Visual attention
+        y = att @ v 
+        y = y.transpose(1, 2).contiguous().view(B, T, self.n_head * (self.n_embed // self.n_head))
+        y = self.c_proj(y)
+
+        return y
+
+class MLP(nn.Module):
+    def __init__(self, config: GPTConfig):
+        super(MLP, self).__init__()
+        self.c_fc = nn.Linear(config.n_embd, config.n_embd * 4) # c_fc means fully connected layer and c is for context
+        self.gelu = nn.GELU()
+        self.c_proj = nn.Linear(config.n_embd * 4, config.n_embd)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.c_fc(x)
+        x = self.gelu(x)
+        x = self.c_proj(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, config: GPTConfig):
+        super(Block, self).__init__()
+        self.ln_1 = nn.LayerNorm(config.n_embd) 
+        self.attn = CasualSelfAttention(config)
+        self.ln_2 = nn.LayerNorm(config.n_embd)
+        self.mlp = MLP(config)
+
+    def forward(self, x: torch.Tensor,visual_features: torch.Tensor, tags_embedding: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.ln_1(x),visual_features, tags_embedding) 
+        x = x + self.mlp(self.ln_2(x)) 
+        return x
+
+
+class DistilGPT2(GPT2LMHeadModel):
+    def __init__(self, config: GPTConfig):
+        super(DistilGPT2, self).__init__(config)
+        self.config = config 
+
+
+        self.transformer = nn.ModuleDict(
+            {
+                'wte': nn.Embedding(config.vocab_size, config.n_embd), 
+                'wpe': nn.Embedding(config.block_size, config.n_embd), 
+                'h': nn.ModuleList(
+                    [
+                        Block(config) for _ in range(config.n_layer)
+                    ]
+                ), # transformer blocks
+                'ln_f': nn.LayerNorm(config.n_embd) # final layer normalization
+            }
+        )
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) # linear layer for projection from embedding to vocab size
+
+    def forward(self, idx: torch.Tensor, visual_features: torch.Tensor = None, tags_embedding: torch.Tensor = None, return_dict: bool = False) -> torch.Tensor:
+        idx=idx.to(self.device)
+        B, T = idx.size()
+        assert T <= self.config.block_size, f"Cannot forward sequence of length {T}, Block size is {self.config.block_size}"
+
+        # forward the token and positional embeddings
+        pos = torch.arange(0, T, dtype=torch.long, device=idx.device) 
+        pos_emb = self.transformer['wpe'](pos)  
+        tok_emb = self.transformer['wte'](idx)  
+        x = tok_emb + pos_emb  
+
+        # forward the transformer
+        for block in self.transformer['h']:
+            x = block(x, visual_features=visual_features, tags_embedding=tags_embedding)
+
+        # forward the head
+        x = self.transformer['ln_f'](x)
+        logits = self.lm_head(x)
+
+        if return_dict:
+            return {'logits': logits}
+        else:
+            
+            return logits
+   
+    
+    @classmethod
+    def from_pretrained(cls, model_type: str):
+        """Loads pre-trained GPT-2 model weights from Hugging Face and handles custom layers."""
+        from transformers import GPT2LMHeadModel
+        print(f"Loading weights from pre-trained GPT: {model_type}")
+
+        # Ensure the model type is supported
+        assert model_type in {'distilgpt2', 'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
+
+        # Define configurations based on the model type
+        config_args = {
+            'distilgpt2':   dict(n_layer=6, n_head=12, n_embd=768),
+            'gpt2':         dict(n_layer=12, n_head=12, n_embd=768),
+            'gpt2-medium':  dict(n_layer=24, n_head=16, n_embd=1024),
+            'gpt2-large':   dict(n_layer=36, n_head=20, n_embd=1280),
+            'gpt2-xl':      dict(n_layer=48, n_head=25, n_embd=1600),
+        }[model_type]
+        config_args['vocab_size'] = 50257
+        config_args['block_size'] = 1024
+
+        # Initialize the custom model with the given configuration
+        config = GPTConfig(**config_args)
+        from transformers import GPT2Config
+
+        config = GPT2Config.from_pretrained('distilgpt2')
+
+        config.visual_size=1024
+        config.block_size=1024
+        config.tags_embd=400
+        config.n_embd=768
+        config.n_layer=6
+        config.n_head=12
+
+        model = cls(config)
+
+        # Load state dictionary from Hugging Face model
+        model_hf = GPT2LMHeadModel.from_pretrained(model_type)
+        sd_hf = model_hf.state_dict()
+
+        # State dictionary of the custom model
+        sd = model.state_dict()
+
+        # Filter out custom keys that are not in the pre-trained model
+        custom_keys = {k for k in sd if 'visual_attn' in k or 'tags_attn' in k}
+        sd_keys_filtered = [k for k in sd if k not in custom_keys]
+
+        # Load matching keys
+        for k in sd_keys_filtered:
+            if k in sd_hf and sd_hf[k].shape == sd[k].shape:
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k])
+
+        # Initialize custom layers separately
+        for k in custom_keys:
+            with torch.no_grad():
+                print(f"Initializing custom layer: {k}")
+                sd[k].normal_(0.0, 0.02)  # Adjust initialization method as needed
+
+        # Update the model's state dictionary
+        model.load_state_dict(sd, strict=False)
+
+        return model
+    
+    def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
+        # Prepare inputs for autoregressive generation
+        inputs = {"idx": input_ids}
+        if past:
+            inputs["past_key_values"] = past  # Include past key values for caching
+
+        # Include additional features like visual and tags if provided
+        if "visual_features" in kwargs:
+            inputs["visual_features"] = kwargs["visual_features"]
+        if "tags_embedding" in kwargs:
+            inputs["tags_embedding"] = kwargs["tags_embedding"]
+
+        return inputs
+    def generate(
+        self,
+        input_ids: torch.Tensor = None,
+        max_length: int = None,
+        min_length: int = None,
+        do_sample: bool = None,
+        early_stopping: bool = None,
+        num_beams: int = None,
+        temperature: float = None,
+        top_k: int = None,
+        top_p: float = None,
+        repetition_penalty: float = None,
+        bos_token_id: int = None,
+        pad_token_id: int = None,
+        eos_token_ids: int = None,
+        length_penalty: float = None,
+        no_repeat_ngram_size: int = None,
+        num_return_sequences: int = None,
+        attention_mask: torch.Tensor = None,
+        visual_features: torch.Tensor = None,
+        tags_embedding: torch.Tensor = None,
+    ):
+        """
+        Generate sequences using autoregressive decoding.
+
+        Args:
+            input_ids (torch.Tensor): Input tensor of token IDs.
+            max_length (int): Maximum length of the generated sequence.
+            min_length (int): Minimum length of the generated sequence.
+            do_sample (bool): Whether to use sampling; if False, uses greedy decoding.
+            early_stopping (bool): Whether to stop when all beams have finished.
+            num_beams (int): Number of beams for beam search.
+            temperature (float): Sampling temperature.
+            top_k (int): Top-k sampling.
+            top_p (float): Top-p (nucleus) sampling.
+            repetition_penalty (float): Penalty for repeated n-grams.
+            bos_token_id (int): Beginning of sequence token ID.
+            pad_token_id (int): Padding token ID.
+            eos_token_ids (int): End of sequence token ID.
+            length_penalty (float): Beam search length penalty.
+            no_repeat_ngram_size (int): Size of n-grams not to repeat.
+            num_return_sequences (int): Number of sequences to return.
+            attention_mask (torch.Tensor): Attention mask for padding tokens.
+            visual_features (torch.Tensor): Visual features for the transformer.
+            tags_embedding (torch.Tensor): Tags embeddings for the transformer.
+
+        Returns:
+            torch.Tensor: Generated sequences of token IDs.
+        """
+        # Default values for unspecified parameters
+        max_length = max_length or self.config.block_size
+        min_length = min_length or 0
+        do_sample = do_sample or False
+        early_stopping = early_stopping or False
+        num_beams = num_beams or 1
+        temperature = temperature or 1.0
+        top_k = top_k or 0
+        top_p = top_p or 1.0
+        repetition_penalty = repetition_penalty or 1.0
+        bos_token_id = bos_token_id or self.config.bos_token_id
+        pad_token_id = pad_token_id or self.config.pad_token_id
+        eos_token_ids = eos_token_ids or self.config.eos_token_ids
+        length_penalty = length_penalty or 1.0
+        no_repeat_ngram_size = no_repeat_ngram_size or 0
+        num_return_sequences = num_return_sequences or 1
+
+        if input_ids is not None:
+            batch_size=input_ids.shape[0]
+        else:
+            batch_size=1
+
+        if input_ids is None:
+            assert isinstance(bos_token_id, int) and bos_token_id >= 0, (
+                "You should either supply a context to complete as `input_ids` input "
+                "or a `bos_token_id` (integer >= 0) as a first token to start the generation."
+            )
+            input_ids = torch.full((batch_size, 1), bos_token_id, dtype=torch.long)
+
+        else:
+            assert input_ids.dim() == 2, "Input prompt should be of shape (batch_size, sequence length)."
+
+        # Avoid duplicate outputs when greedy decoding
+        if not do_sample:
+            if num_beams == 1:
+                assert num_return_sequences == 1, (
+                    "Greedy decoding will always produce the same output for num_beams == 1 "
+                    "and num_return_sequences > 1. Please set num_return_sequences = 1."
+                )
+            else:
+                assert num_beams >= num_return_sequences, (
+                    "Greedy beam search decoding cannot return more sequences than it has beams. "
+                    "Please set num_beams >= num_return_sequences."
+                )
+
+        # Create attention mask if necessary
+        if attention_mask is None:
+            if pad_token_id is not None and pad_token_id in input_ids:
+                attention_mask = (input_ids != pad_token_id).long()
+            else:
+                attention_mask = torch.ones_like(input_ids)
+
+        # Set pad_token_id if not provided and eos_token_ids is available
+        if pad_token_id is None and eos_token_ids is not None:
+            pad_token_id = eos_token_ids
+            print(f"Setting `pad_token_id` to {pad_token_id} (first `eos_token_ids`) to generate sequence.")
+
+        # Current sequence length and vocabulary size
+        cur_len = input_ids.size(1)
+        vocab_size = self.config.vocab_size
+
+        # Adjust effective batch size and multiplier for sampling
+        if do_sample:
+            effective_batch_size = batch_size * num_return_sequences
+            effective_batch_mult = num_return_sequences
+        else:
+            effective_batch_size = batch_size
+            effective_batch_mult = 1
+
+        # Expand input_ids and attention_mask for beam search or multiple return sequences
+        if num_return_sequences > 1 or num_beams > 1:
+            input_ids_len = input_ids.size(-1)
+
+            # Expand dimensions and repeat for each beam and return sequence
+            input_ids = input_ids.unsqueeze(1).expand(batch_size, effective_batch_mult * num_beams, input_ids_len)
+            attention_mask = attention_mask.unsqueeze(1).expand(batch_size, effective_batch_mult * num_beams, input_ids_len)
+
+            # Reshape to combine batch and beam dimensions
+            input_ids = input_ids.reshape(effective_batch_size * num_beams, input_ids_len)
+            attention_mask = attention_mask.reshape(effective_batch_size * num_beams, input_ids_len)
+
+        if num_beams > 1:
+            output = self._generate_beam_search(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                visual_features=visual_features,
+                tags_embedding=tags_embedding,
+                cur_len=input_ids.size(1),
+                max_length=max_length,
+                min_length=min_length,
+                do_sample=do_sample,
+                early_stopping=early_stopping,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+                no_repeat_ngram_size=no_repeat_ngram_size,
+                pad_token_id=pad_token_id,
+                eos_token_ids=eos_token_ids,
+                length_penalty=length_penalty,
+                num_return_sequences=num_return_sequences,
+                num_beams=num_beams,
+            )
+        else:
+            output = self._generate_no_beam_search(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                visual_features=visual_features,
+                tags_embedding=tags_embedding,
+                cur_len=input_ids.size(1),
+                max_length=max_length,
+                min_length=min_length,
+                do_sample=do_sample,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+                no_repeat_ngram_size=no_repeat_ngram_size,
+                pad_token_id=pad_token_id,
+                eos_token_ids=eos_token_ids,
+                batch_size=batch_size,
+                vocab_size=vocab_size,
+            )
+
+        return output
+    
+
+    def _generate_no_beam_search(
+    self,
+    input_ids,
+    visual_features,
+    tags_embedding,
+    cur_len,
+    max_length,
+    min_length,
+    do_sample,
+    temperature,
+    top_k,
+    top_p,
+    repetition_penalty,
+    no_repeat_ngram_size,
+    pad_token_id,
+    eos_token_ids,
+    batch_size,
+    vocab_size,
+    attention_mask,
+):
+        """
+        Generate sequences for each example without beam search (num_beams == 1).
+        All returned sequences are generated independently.
+        """
+        # Track unfinished sentences and their lengths
+        unfinished_sents=torch.ones_like(input_ids[:,0])
+        sent_lengths=torch.ones_like(input_ids[:,0])*max_length
+
+        past=None
+        
+
+        while cur_len < max_length:
+            if past is None:
+                inputs = input_ids
+            else:
+                inputs = input_ids[:, -1].unsqueeze(1)
+
+            model_inputs = self.prepare_inputs_for_generation(
+                inputs, past=past, visual_features=visual_features, tags_embedding=tags_embedding
+            )
+            outputs = self(**model_inputs)
+            # next_token_logits = outputs[0][-1, :]  # Extract logits for the last token, shape: [batch_size, vocab_size]
+            next_token_logits = outputs[:, -1, :]
+
+            # next_token_logits = next_token_logits.unsqueeze(0)  # Add a new dimension: [1, batch_size, vocab_size]
+            next_token_logits = next_token_logits.expand(batch_size, vocab_size)  # Expand to match batch size: [batch_size, vocab_size]
+
+
+            # if self._do_output_past(outputs): # we dont have this function implemented
+            #     past = outputs[1]
+
+            # Apply repetition penalty
+            if repetition_penalty != 1.0:
+                next_token_logits_penalties=self._create_next_token_logits_penalties(input_ids,next_token_logits,repetition_penalty)
+                next_token_logits=next_token_logits @ next_token_logits_penalties.T # .T de mn 3ndy 
+               
+
+            # Prevent repetition of n-grams
+            if no_repeat_ngram_size > 0:   # not checked generated by chat
+                banned_tokens=self.calc_banned_ngram_tokens(input_ids,batch_size,no_repeat_ngram_size,cur_len) # not checked generated by chat
+                banned_tokens_indices_mask=[]
+
+                for banned_tokens_slice in banned_tokens:
+                        banned_tokens_indices_mask.append(
+                            [True if token in banned_tokens_slice else False for token in range(vocab_size)]
+                        )   
+                
+                banned_tokens_indices_mask=torch.tensor(banned_tokens_indices_mask,dtype=bool)  
+
+                next_token_logits[banned_tokens_indices_mask]= -float('inf')                         
+
+            # Min length constraint for EOS
+            if eos_token_ids is not None and cur_len < min_length:
+                # create eos_token_id boolean mask
+                is_token_logit_eos_token = torch.arange(vocab_size, device=next_token_logits.device) == eos_token_ids
+                eos_token_indices_mask = is_token_logit_eos_token.unsqueeze(0).expand(batch_size, -1)
+                # next_token_logits=next_token_logits.unsqueeze(0).expand(batch_size,vocab_size)
+
+                next_token_logits = next_token_logits.masked_fill(eos_token_indices_mask, -float("inf"))
+
+
+
+               
+
+            # Sampling or greedy decoding
+            if do_sample:
+                if temperature != 1.0:
+                    next_token_logits = next_token_logits / temperature
+                
+                next_token_logits=self.top_k_top_p_filtering(next_token_logits,top_k=top_k,top_p=top_p)
+
+                next_token = torch.multinomial(torch.softmax(next_token_logits, dim=-1), num_samples=1).squeeze(1)
+
+                
+            else:
+                next_token=torch.argmax(next_token_logits,dim=-1)
+            
+
+            if eos_token_ids is not None:
+                unfinished_sents=unfinished_sents.to(self.device)
+                tokens_to_add = next_token * unfinished_sents + pad_token_id * (1 - unfinished_sents)
+
+            else:
+                tokens_to_add = next_token
+            input_ids=input_ids.to(self.device)
+            input_ids = torch.cat([input_ids, tokens_to_add.unsqueeze(-1)], dim=1)
+
+            if eos_token_ids is not None:
+                eos_in_sents = tokens_to_add == eos_token_ids
+                # If sentence is unfinished and the token to add is eos, sent_lengths is filled with current length
+                is_sents_unfinished_and_token_to_add_is_eos = unfinished_sents * eos_in_sents.int()
+                sent_lengths=sent_lengths.to(self.device)
+                sent_lengths = (
+                    sent_lengths * (1 - is_sents_unfinished_and_token_to_add_is_eos)
+                    + cur_len * is_sents_unfinished_and_token_to_add_is_eos
+                )
+
+                # Unfinished sentences are set to zero if eos is in the sentence
+                unfinished_sents -= is_sents_unfinished_and_token_to_add_is_eos
+
+            # Stop if there is a </s> in each sentence, or if we exceed the maximum length
+            if torch.max(unfinished_sents) == 0:  # => this line is what keeps it stopping at 57 etc..
+                break
+
+            cur_len += 1
+
+        # Pad sequences if necessary
+        min_sent_length = sent_lengths.min()
+        max_sent_length = sent_lengths.max()
+
+        if min_sent_length != max_sent_length:
+            assert pad_token_id is not None, "`Pad_token_id` has to be defined if batches have different lengths"
+            padding = torch.ones((batch_size, max_sent_length), dtype=torch.int) * pad_token_id
+            broad_casted_sent_lengths = sent_lengths.unsqueeze(-1).expand(batch_size, max_sent_length)
+            broad_casted_range = torch.arange(max_sent_length).unsqueeze(0).expand(batch_size, max_sent_length).T
+
+            # Use torch.where to apply padding where necessary
+            decoded = torch.where(broad_casted_range < broad_casted_sent_lengths, input_ids, padding)
+        else:
+            decoded = input_ids
+
+        return decoded
+    
+    def _create_next_token_logits_penalties(self,input_ids, logits, repetition_penalty):
+        """
+        Create logit penalties for already seen input_ids based on repetition penalty.
+
+        Args:
+            input_ids (torch.Tensor): Tensor of shape (batch_size, seq_len) containing input token IDs.
+            logits (torch.Tensor): Tensor of shape (batch_size, vocab_size) containing next-token logits.
+            repetition_penalty (float): The penalty to apply for repeated tokens.
+
+        Returns:
+            torch.Tensor: Tensor of shape (batch_size, vocab_size) with applied penalties.
+        """
+        token_penalties=torch.ones_like(logits)
+        prev_input_ids=[torch.unique(input_id) for input_id in input_ids]
+
+        for i, prev_input_id in enumerate(prev_input_ids):
+            logits_penalized=logits[i][prev_input_ids]
+            logit_penalties=torch.zeros_like(logits_penalized)
+
+            logit_penalties[logits_penalized<0]=repetition_penalty
+            logit_penalties[logits_penalized>0]=1/repetition_penalty
+
+            token_penalties[i].scatter_(0,prev_input_id,logit_penalties)
+        return token_penalties
+
+    
+    def top_k_top_p_filtering(self,logits, top_k=0, top_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1):
+        """
+        Filter a distribution of logits using top-k and/or nucleus (top-p) filtering.
+        
+        Args:
+            logits: Logits distribution of shape (batch size, vocabulary size).
+            top_k (int): Keep only top k tokens with the highest probability.
+            top_p (float): Keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+            filter_value (float): Value to assign to filtered logits.
+            min_tokens_to_keep (int): Ensure at least this many tokens are kept.
+
+        Returns:
+            torch.Tensor: Filtered logits.
+        """
+        logits_shape = logits.size()
+
+        # Top-k filtering
+        if top_k > 0:
+            top_k = min(max(top_k, min_tokens_to_keep), logits_shape[-1])  # Safety check
+            # Remove all tokens with a probability less than the last token of the top-k
+            top_k_values, _ = torch.topk(logits, top_k, dim=-1)
+            min_top_k_values = top_k_values[:, -1].unsqueeze(-1)  # Minimum logit in top-k
+            logits = torch.where(logits < min_top_k_values, torch.full_like(logits, filter_value), logits)
+
+        # Top-p (nucleus) filtering
+        if top_p < 1.0:
+            sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+            cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+
+            # Remove tokens with cumulative probability above the threshold
+            sorted_indices_to_remove = cumulative_probs > top_p
+
+            if min_tokens_to_keep > 1:
+                # Ensure we keep at least min_tokens_to_keep tokens
+                sorted_indices_to_remove[:, :min_tokens_to_keep] = 0
+
+            # Shift the indices to the right to keep also the first token above the threshold
+            sorted_indices_to_remove = sorted_indices_to_remove.roll(1, dims=-1)
+            sorted_indices_to_remove[:, 0] = 0
+
+            # Scatter sorted indices back to original indexing
+            indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+            logits = torch.where(indices_to_remove, torch.full_like(logits, filter_value), logits)
+
+        return logits
+    
+    def calc_banned_ngram_tokens(self,prev_input_ids, num_hypos, no_repeat_ngram_size, cur_len):
+        """
+        Calculate banned n-gram tokens for no-repeat n-gram constraints.
+        
+        Args:
+            prev_input_ids (torch.Tensor): Tensor of shape (num_hypos, seq_len) containing token sequences.
+            num_hypos (int): Number of hypotheses in the batch.
+            no_repeat_ngram_size (int): Size of the n-grams to avoid repeating.
+            cur_len (int): Current length of the sequence being generated.
+        
+        Returns:
+            List[List[int]]: List of banned tokens for each hypothesis.
+        """
+        if cur_len + 1 < no_repeat_ngram_size:
+            # Return no banned tokens if not enough tokens have been generated
+            return [[] for _ in range(num_hypos)]
+
+        # Dictionary to store generated n-grams for each hypothesis
+        generated_ngrams = [{} for _ in range(num_hypos)]
+
+        # Populate the n-grams
+        for idx in range(num_hypos):
+            gen_tokens = prev_input_ids[idx].tolist()
+            generated_ngram = generated_ngrams[idx]
+            for ngram in zip(*[gen_tokens[i:] for i in range(no_repeat_ngram_size)]):
+                prev_ngram_tuple = tuple(ngram[:-1])
+                generated_ngram[prev_ngram_tuple] = generated_ngram.get(prev_ngram_tuple, []) + [ngram[-1]]
+
+        def _get_generated_ngrams(hypo_idx):
+            # Get n-grams that have already appeared
+            start_idx = cur_len + 1 - no_repeat_ngram_size
+            ngram_idx = tuple(prev_input_ids[hypo_idx, start_idx:cur_len].tolist())
+            return generated_ngrams[hypo_idx].get(ngram_idx, [])
+
+        # Calculate banned tokens for each hypothesis
+        banned_tokens = [_get_generated_ngrams(hypo_idx) for hypo_idx in range(num_hypos)]
+        return banned_tokens
+
+

generate_report.py

@@ -0,0 +1,69 @@
+from PIL import Image
+import io
+import torch
+import os
+import numpy as np
+from CNN_encoder import CNN_Encoder
+from distil_gpt2 import DistilGPT2
+from configs import argHandler
+from utils import load_image
+from tokenizer_wrapper import TokenizerWrapper
+from huggingface_hub import hf_hub_download
+# from src.models.cnn_encoder import 
+# from src.models.distil_gpt2 import DistilGPT2
+# from src.configs import argHandler
+
+FLAGS = argHandler()
+FLAGS.setDefaults()
+tokenizer_wrapper = TokenizerWrapper( FLAGS.csv_label_columns[0], FLAGS.max_sequence_length, FLAGS.tokenizer_vocab_size)
+
+encoder = CNN_Encoder('pretrained_visual_model', FLAGS.visual_model_name, FLAGS.visual_model_pop_layers,
+                      FLAGS.encoder_layers, FLAGS.tags_threshold, num_tags=len(FLAGS.tags))
+decoder = DistilGPT2.from_pretrained('distilgpt2')
+
+optimizer = torch.optim.Adam(decoder.parameters(), lr=FLAGS.learning_rate)
+# checkpoint_path = os.path.join(FLAGS.ckpt_path, "checkpoint.pth")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+encoder.to(device)
+decoder.to(device)
+
+checkpoint_path = hf_hub_download(repo_id="TransformingBerry/CDGPT2_checkpoint", filename="checkpoint.pth")
+
+
+if os.path.exists(checkpoint_path):
+    print(f"Restoring from checkpoint: {checkpoint_path}")
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)
+    encoder.load_state_dict(checkpoint['encoder_state_dict'])
+    decoder.load_state_dict(checkpoint['decoder_state_dict'])
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+else:
+    print("No checkpoint found. Starting from scratch.")
+
+def generate_report(image_bytes):
+    image = Image.open(io.BytesIO(image_bytes))
+    image_tensor = load_image(image)
+    
+    visual_features, tags_embedding = encoder(image_tensor)
+    
+    dec_input = torch.unsqueeze(
+        torch.tensor(tokenizer_wrapper.GPT2_encode('startseq', pad=False)), 0
+    )
+
+    generation_config = {
+        "visual_features": visual_features,
+        "tags_embedding": tags_embedding,
+        "num_beams": 1,
+        "max_length": FLAGS.max_sequence_length,
+        "min_length": 3,
+        "eos_token_ids": tokenizer_wrapper.GPT2_eos_token_id(),
+        "pad_token_id": tokenizer_wrapper.GPT2_pad_token_id(),
+        "do_sample": False,
+        "early_stopping": True,
+    }
+
+    tokens = decoder.generate(dec_input, **generation_config)
+    sentence = tokenizer_wrapper.GPT2_decode(tokens[0])
+    sentence = tokenizer_wrapper.filter_special_words(sentence)
+    print(sentence)
+    
+    return {"report": sentence}

generator.py

@@ -0,0 +1,93 @@
+import numpy as np
+import os
+import pandas as pd
+from torch.utils.data import Dataset, DataLoader
+from PIL import Image
+from skimage.transform import resize
+import torch
+from torchvision import transforms
+
+
+class AugmentedImageSequence(Dataset):
+    """
+    Thread-safe image generator with imgaug support in PyTorch
+    """
+
+    def __init__(self, dataset_csv_file, class_names, source_image_dir, tokenizer_wrapper, batch_size=16,
+                 target_size=(224, 224), augmenter=None, verbose=0, steps=None,
+                 shuffle_on_epoch_end=True, random_state=1):
+        """
+        :param dataset_csv_file: str, path of dataset csv file
+        :param class_names: list of str
+        :param batch_size: int
+        :param target_size: tuple(int, int)
+        :param augmenter: imgaug object. Do not specify resize in augmenter.
+                          It will be done automatically according to input_shape of the model.
+        :param verbose: int
+        """
+        self.dataset_df = pd.read_csv(dataset_csv_file)
+        self.source_image_dir = source_image_dir
+        self.batch_size = batch_size
+        self.target_size = target_size
+        self.augmenter = augmenter
+        self.tokenizer_wrapper = tokenizer_wrapper
+        self.verbose = verbose
+        self.shuffle = shuffle_on_epoch_end
+        self.random_state = random_state
+        self.class_names = class_names
+        self.prepare_dataset()
+        if steps is None:
+            self.steps = int(np.ceil(len(self.x_path) / float(self.batch_size)))
+        else:
+            self.steps = int(steps)
+        
+        self.transform = transforms.Compose([
+            # transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+
+    def __len__(self):
+        return self.steps
+
+    def __getitem__(self, idx):
+        batch_x_path = self.x_path[idx * self.batch_size:(idx + 1) * self.batch_size]
+        batch_x = torch.stack([self.load_image(x_path) for x_path in batch_x_path])
+        batch_x = self.transform_batch_images(batch_x)
+        batch_y = torch.tensor(self.y[idx * self.batch_size:(idx + 1) * self.batch_size])
+        return batch_x, batch_y,  batch_x_path.tolist()
+
+    def load_image(self, image_file):
+        image_path = os.path.join(self.source_image_dir, image_file)
+        image = Image.open(image_path).convert("RGB")
+        image_array = np.asarray(image) / 255.
+        image_array = resize(image_array, self.target_size)
+        image_tensor = torch.tensor(image_array, dtype=torch.float32).permute(2, 0, 1)  # Convert to CxHxW
+        return image_tensor
+
+    def transform_batch_images(self, batch_x):
+        if self.augmenter is not None:
+            batch_x = torch.stack([torch.tensor(self.augmenter.augment_image(img.permute(1, 2, 0).numpy())) for img in batch_x])
+        batch_x = self.transform(batch_x)
+        return batch_x
+
+    def get_y_true(self):
+        """
+        Use this function to get y_true for DataLoader.
+        Ensure shuffle_on_epoch_end is False before using.
+        """
+        if self.shuffle:
+            raise ValueError("get_y_true() can only be used when shuffle_on_epoch_end is False.")
+        return torch.tensor(self.y[:self.steps * self.batch_size], dtype=torch.float32)
+
+    def prepare_dataset(self):
+        df = self.dataset_df.sample(frac=1., random_state=self.random_state)
+        ## @TODO: tokenize the targets
+        self.x_path, self.y = df["Image Index"].values,  self.tokenizer_wrapper.GPT2_encode(
+                df[self.class_names].values)
+        
+
+    def on_epoch_end(self):
+        if self.shuffle:
+            self.random_state += 1
+            self.prepare_dataset()
+

main.py

@@ -0,0 +1,37 @@
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from generate_report import generate_report
+from utils import convert_to_png
+import io
+
+
+app = FastAPI(title="FastAPI Example App", version="0.1.0")
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@app.get("/")
+async def read_root():
+    return {"message": "Hello oopa"}
+
+@app.post("/upload-image/")
+async def upload_image(file: UploadFile = File(...)):
+    try:
+        image = await convert_to_png(file)
+        image_io = io.BytesIO()
+        image.save(image_io, format="PNG")
+        image_data = image_io.getvalue()
+        report = generate_report(image_data)
+        return {"report": report}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+    
+
+
+
+    

pretrained_visual_model/best_weights.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c299f950c6a69bb3c82f20afbcf312c0d53e800b27c2420eb4f1c1e7edf77b8
+size 83811496

pretrained_visual_model/fine_tuned_chexnet.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6596d037bb3cc1c1959b7379f9fe8a4bf238963a09d11cd1bbe5d1688192c2c9
+size 29517896