class_docstr: PNPVQA model consists of three submodels for zero-shot VQA:
      1. Image-questioning matching model
      2. Image captioning model
      3. Question answering model
  
  Supported model types:
      - base: BLIPITM, BLIPCaption, PNPUnifiedQAv2FiD (t5-base)
      - large: BLIPITM, BLIPCaption, PNPUnifiedQAv2FiD (t5-large)
      - 3b: BLIPITM, BLIPCaption, PNPUnifiedQAv2FiD (t5-3b)
  
  Usage:
      >>> from lavis.models import load_model
      >>> model = load_model("pnp_vqa", "base", is_eval=True)
      >>> model = load_model("pnp_vqa", "large", is_eval=True)
      >>> model = load_model("pnp_vqa", "3b", is_eval=True)
class_docstr_tok: ["PNPVQA","model","consists","of","three","submodels","for","zero-shot","VQA",":","1",".","Image-questioning","matching","model","2",".","Image","captioning","model","3",".","Question","answering","model","Supported","model","types",":","-","base",":","BLIPITM",",","BLIPCaption",",","PNPUnifiedQAv2FiD","(","t5-base",")","-","large",":","BLIPITM",",","BLIPCaption",",","PNPUnifiedQAv2FiD","(","t5-large",")","-","3b",":","BLIPITM",",","BLIPCaption",",","PNPUnifiedQAv2FiD","(","t5-3b",")","Usage",":",">",">",">","from","lavis.models","import","load_model",">",">",">","model","=","load_model","(","``","pnp_vqa","''",",","``","base","''",",","is_eval=True",")",">",">",">","model","=","load_model","(","``","pnp_vqa","''",",","``","large","''",",","is_eval=True",")",">",">",">","model","=","load_model","(","``","pnp_vqa","''",",","``","3b","''",",","is_eval=True",")"]
docstr: Args:
      samples (dict): A dictionary containing the following keys:
          - image (torch.Tensor): A tensor of shape (batch_size, 3, H, W)
          - text_input (list): A list of strings of length batch_size
          - gradcams (torch.Tensor): A tensor of shape (batch_size, H*W)
      cap_max_length (int): The maximum length of the caption to be generated.
      cap_min_length (int): The minimum length of the caption to be generated.
      top_p (float): The cumulative probability for nucleus sampling.
      top_k (float): The number of the highest probability tokens for top-k sampling.
      repetition_penalty (float): The parameter for repetition penalty. 1.0 means no penalty.
      num_captions (int): Number of captions generated for each image.
      num_patches (int): Number of patches sampled for each image.
  
  Returns:
      samples (dict): A dictionary containing the following keys:
          - image (torch.Tensor): A tensor of shape (batch_size, 3, H, W)
          - text_input (list): A list of strings of length batch_size
          - gradcams (torch.Tensor): A tensor of shape (batch_size, H*W)
          - captions (nested list): A nested list of strings of total length batch_size * num_captions
docstr_tok: ["Args",":","samples","(","dict",")",":","A","dictionary","containing","the","following","keys",":","-","image","(","torch.Tensor",")",":","A","tensor","of","shape","(","batch_size",",","3",",","H",",","W",")","-","text_input","(","list",")",":","A","list","of","strings","of","length","batch_size","-","gradcams","(","torch.Tensor",")",":","A","tensor","of","shape","(","batch_size",",","H","*","W",")","cap_max_length","(","int",")",":","The","maximum","length","of","the","caption","to","be","generated",".","cap_min_length","(","int",")",":","The","minimum","length","of","the","caption","to","be","generated",".","top_p","(","float",")",":","The","cumulative","probability","for","nucleus","sampling",".","top_k","(","float",")",":","The","number","of","the","highest","probability","tokens","for","top-k","sampling",".","repetition_penalty","(","float",")",":","The","parameter","for","repetition","penalty",".","1.0","means","no","penalty",".","num_captions","(","int",")",":","Number","of","captions","generated","for","each","image",".","num_patches","(","int",")",":","Number","of","patches","sampled","for","each","image",".","Returns",":","samples","(","dict",")",":","A","dictionary","containing","the","following","keys",":","-","image","(","torch.Tensor",")",":","A","tensor","of","shape","(","batch_size",",","3",",","H",",","W",")","-","text_input","(","list",")",":","A","list","of","strings","of","length","batch_size","-","gradcams","(","torch.Tensor",")",":","A","tensor","of","shape","(","batch_size",",","H","*","W",")","-","captions","(","nested","list",")",":","A","nested","list","of","strings","of","total","length","batch_size","*","num_captions"]
code: def forward_cap(
      self,
      samples,
      cap_max_length=20,
      cap_min_length=0,
      top_p=1,
      top_k=50,
      repetition_penalty=1.0,
      num_captions=100,
      num_patches=20,
  ):
      """
      Args:
          samples (dict): A dictionary containing the following keys:
              - image (torch.Tensor): A tensor of shape (batch_size, 3, H, W)
              - text_input (list): A list of strings of length batch_size
              - gradcams (torch.Tensor): A tensor of shape (batch_size, H*W)
          cap_max_length (int): The maximum length of the caption to be generated.
          cap_min_length (int): The minimum length of the caption to be generated.
          top_p (float): The cumulative probability for nucleus sampling.
          top_k (float): The number of the highest probability tokens for top-k sampling.
          repetition_penalty (float): The parameter for repetition penalty. 1.0 means no penalty.
          num_captions (int): Number of captions generated for each image.
          num_patches (int): Number of patches sampled for each image.
  
      Returns:
          samples (dict): A dictionary containing the following keys:
              - image (torch.Tensor): A tensor of shape (batch_size, 3, H, W)
              - text_input (list): A list of strings of length batch_size
              - gradcams (torch.Tensor): A tensor of shape (batch_size, H*W)
              - captions (nested list): A nested list of strings of total length batch_size * num_captions
      """
      encoder_out = self.image_captioning_model.forward_encoder(samples)
      captions = [[] for _ in range(encoder_out.size(0))]
  
      min_num_captions = 0
  
      while min_num_captions < num_captions:
          encoder_out_samples = []
          for i in range(num_captions):
              patch_id = (
                  torch.multinomial(
                      samples["gradcams"].to(
                          self.image_captioning_model.device
                      ),
                      num_patches,
                  ).reshape(encoder_out.size(0), -1)
                  + 1
              )
              patch_id = (
                  patch_id.sort(dim=1)
                  .values.unsqueeze(-1)
                  .expand(-1, -1, encoder_out.size(2))
              )
              encoder_out_sample = torch.gather(
                  encoder_out, 1, patch_id
              )
              encoder_out_samples.append(encoder_out_sample)
  
          stacked = torch.stack(encoder_out_samples, dim=1)
          image_embeds = torch.flatten(
              stacked, start_dim=0, end_dim=1
          )  # (bsz*num_seq, num_patch, dim)
  
          image_atts = torch.ones(
              image_embeds.size()[:-1], dtype=torch.long
          ).to(self.image_captioning_model.device)
          model_kwargs = {
              "encoder_hidden_states": image_embeds,
              "encoder_attention_mask": image_atts,
          }
  
          prompt = [
              self.image_captioning_model.prompt
          ] * image_embeds.size(0)
          prompt = self.image_captioning_model.tokenizer(
              prompt, return_tensors="pt"
          ).to(self.image_captioning_model.device)
          prompt.input_ids[
              :, 0
          ] = self.image_captioning_model.tokenizer.bos_token_id
          prompt.input_ids = prompt.input_ids[:, :-1]
  
          decoder_out = self.image_captioning_model.text_decoder.generate(
              input_ids=prompt.input_ids,
              max_length=cap_max_length,
              min_length=cap_min_length,
              do_sample=True,
              top_p=top_p,
              top_k=top_k,
              num_return_sequences=1,
              eos_token_id=self.image_captioning_model.tokenizer.sep_token_id,
              pad_token_id=self.image_captioning_model.tokenizer.pad_token_id,
              repetition_penalty=repetition_penalty,
              **model_kwargs,
          )
  
          outputs = self.image_captioning_model.tokenizer.batch_decode(
              decoder_out, skip_special_tokens=True
          )
  
          for counter, output in enumerate(outputs):
              ind = counter // num_captions
              if len(captions[ind]) < num_captions:
                  caption = output[
                      len(self.image_captioning_model.prompt) :
                  ]
                  overlap_caption = [
                      1 for caps in captions[ind] if caption in caps
                  ]
                  if len(overlap_caption) == 0:
                      captions[ind].append(caption)
  
          min_num_captions = min([len(i) for i in captions])
  
      samples["captions"] = captions
  
      return samples
  
file_path: lavis/lavis/models/pnp_vqa_models/pnp_vqa.py
imports: import torch
  import torch.nn
  from itertools import chain
  from lavis.common.registry import registry
  from lavis.models.base_model import BaseModel
  from torch.nn import CrossEntropyLoss, MSELoss
  from transformers import T5ForConditionalGeneration
  from lavis.models.pnp_vqa_models import prepare_qa_input
  from lavis.models.blip_models.blip_image_text_matching import compute_gradcam
  from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
label_desc: Use image node_id 3 for calling the PNPVQA obj's underlying member method code with example usage: obj.forward_cap(samples, cap_max_length, cap_min_length, top_p, top_k, repetition_penalty, num_captions, num_patches) and returns: samples