lavis/models/img2prompt_models/img2prompt_vqa.py

"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import random

import spacy
import torch
import torch.nn.functional as F
from transformers import T5ForConditionalGeneration, T5Tokenizer

from lavis.common.dist_utils import download_cached_file
from lavis.common.registry import registry
from lavis.models.base_model import BaseModel
from lavis.models.blip_models.blip_image_text_matching import compute_gradcam

open_pos = ["NOUN", "VERB", "ADJ", "ADV", "NUM"]



@registry.register_model("img2prompt_vqa")
class Img2PromptVQA(BaseModel):
    """
    Img2Prompt_VQA model consists of three submodels for zero-shot VQA:
        1. Image-questioning matching model
        2. Image captioning model
        3. Large Language 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("img2prompt_vqa", "base", is_eval=True)
    """

    PRETRAINED_MODEL_CONFIG_DICT = {
        "base": "configs/models/img2prompt-vqa/img2prompt_vqa_base.yaml",
    }

    def __init__(
        self,
        image_question_matching_model,
        image_captioning_model,
        question_generation_model,
        question_generation_tokenizer,
        offload_model=False,
    ):
        super().__init__()

        self.image_question_matching_model = image_question_matching_model
        self.image_captioning_model = image_captioning_model
        self.question_generation_model = question_generation_model
        self.question_generation_tokenizer = question_generation_tokenizer
        self.offload_model = offload_model
        self.nlp = spacy.load("en_core_web_sm")

    def forward_itm(self, samples, block_num=7):
        """
        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
            block_num (int): The index of cross-attention block for gradcam computation.

        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)
        """
        image = samples["image"]
        question = [text.strip("?") for text in samples["text_input"]]
        tokenized_text = self.image_question_matching_model.tokenizer(
            question, padding="longest", truncation=True, return_tensors="pt"
        ).to(self.image_question_matching_model.device)
        with torch.set_grad_enabled(True):
            gradcams, _ = compute_gradcam(
                model=self.image_question_matching_model,
                visual_input=image,
                text_input=question,
                tokenized_text=tokenized_text,
                block_num=block_num,
            )

        gradcams = [gradcam_[1] for gradcam_ in gradcams]
        samples["gradcams"] = torch.stack(gradcams).reshape(
            samples["image"].size(0), -1
        )

        return samples

    def itm_rank(self, image_embeds, image_atts, encoder_input_ids, match_head="itm"):
        # breakpoint()
        encoder_input_ids = encoder_input_ids.clone()
        encoder_input_ids = encoder_input_ids[:, self.prompt_length - 1 :]
        text_attention_mask = (encoder_input_ids != self.tokenizer.pad_token_id).long()

        if match_head == "itm":
            # encoder_input_ids = encoder_input_ids.clone()
            encoder_input_ids[:, 0] = self.tokenizer.enc_token_id
            output = self.text_encoder(
                encoder_input_ids,
                attention_mask=text_attention_mask,
                encoder_hidden_states=image_embeds,
                encoder_attention_mask=image_atts,
                return_dict=True,
            )
            itm_output = self.itm_head(output.last_hidden_state[:, 0, :])
            return itm_output  # , mask, token_length

        elif match_head == "itc":
            encoder_input_ids[:, 0] = self.tokenizer.cls_token_id
            text_output = self.text_encoder(
                encoder_input_ids,
                attention_mask=text_attention_mask,
                return_dict=True,
                mode="text",
            )
            image_feat = F.normalize(self.vision_proj(image_embeds[:, 0, :]), dim=-1)
            text_feat = F.normalize(
                self.text_proj(text_output.last_hidden_state[:, 0, :]), dim=-1
            )

            sim = image_feat @ text_feat.t()
            return sim

    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
            )

            itm_outputs = self.image_question_matching_model.itm_rank(
                image_embeds, image_atts, encoder_input_ids=decoder_out
            )  # caption filter

            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]
                    # print(itm_outputs)
                    if (
                        len(overlap_caption) == 0 and itm_outputs[counter] >= 0.5
                    ):  # image filter
                        captions[ind].append(caption)

            min_num_captions = min([len(i) for i in captions])

        samples["captions"] = captions

        return samples

    def answer_extraction(self, caption, num_question_generation=30):
        cap_use = ""
        # print(caption)
        caption = caption
        ans_to_cap_dict = {}
        answers = []
        for cap_idx, cap in enumerate(caption):
            # print(cap)
            cap_use += cap
            cap = cap.strip().strip(".")
            # print(cap)
            cap = self.nlp(cap)
            for token in cap:  # Noun /Verb/Adj//NUM
                if token.pos_ in open_pos:
                    if token.text.lower() not in ans_to_cap_dict:
                        ans_to_cap_dict[token.text.lower()] = [cap_idx]
                    else:
                        if cap_idx not in ans_to_cap_dict[token.text.lower()]:
                            ans_to_cap_dict[token.text.lower()].append(cap_idx)
                    answers.append(token.text)
            for ent in cap.ents:

                if ent.text not in answers:
                    if ent.text.lower() not in ans_to_cap_dict:
                        ans_to_cap_dict[ent.text.lower()] = [cap_idx]
                    else:
                        if cap_idx not in ans_to_cap_dict[ent.text.lower()]:
                            ans_to_cap_dict[ent.text.lower()].append(cap_idx)
                    answers.append(ent.text)
            for chunk in cap.noun_chunks:
                if len(chunk.text.split()) < 4:
                    if chunk.text.lower() not in ans_to_cap_dict:
                        ans_to_cap_dict[chunk.text.lower()] = [cap_idx]
                    else:
                        if cap_idx not in ans_to_cap_dict[chunk.text.lower()]:
                            ans_to_cap_dict[chunk.text.lower()].append(cap_idx)
                    #                 print(chunk.text)
                    answers.append(chunk.text)
        answers = sorted(answers, key=answers.count, reverse=True)
        real_answers = []
        for i in answers:
            i = i + "."
            if i not in real_answers:
                real_answers.append(i)

        contexts_for_question_generation = []
        answers = []
        for ans in real_answers[
            :num_question_generation
        ]:  # Generate questions for 30 answers with max frequencies.
            contexts_for_question_generation.append(
                "answer: %s  context: %s." % (ans, cap_use)
            )
            answers.append(ans)
        contexts_for_question_generation.append(
            "answer: %s  context: %s." % ("yes.", cap_use)
        )
        answers.append("yes.")
        return contexts_for_question_generation, answers, ans_to_cap_dict

    def forward_qa_generation(self, samples):
        caption = samples["captions"][0]
        (
            contexts_for_question_generation,
            answers,
            ans_to_cap_dict,
        ) = self.answer_extraction(caption)
        inputs = self.question_generation_tokenizer(
            contexts_for_question_generation,
            padding="longest",
            truncation=True,
            max_length=2048,
            return_tensors="pt",
        ).to(self.device)
        question_size = inputs.input_ids.shape[0]
        cur_b = 0
        true_input_size = 10
        outputs_list = []
        while cur_b < question_size:
            outputs = self.question_generation_model.generate(
                input_ids=inputs.input_ids[cur_b : cur_b + true_input_size],
                attention_mask=inputs.attention_mask[cur_b : cur_b + true_input_size],
                num_beams=3,
                max_length=30,
            )
            questions = self.question_generation_tokenizer.batch_decode(
                outputs, skip_special_tokens=True
            )
            outputs_list += questions
            cur_b += true_input_size
        questions = outputs_list
        samples["questions"] = questions
        samples["answers"] = answers
        samples["ans_to_cap_dict"] = ans_to_cap_dict
        # results.append({"question_id": ques_id, "question":questions,"answer":answers})
        return samples

    def create_context_prompt(self, samples, num_caps_per_img=30):
        ans_dict_queid = samples["ans_to_cap_dict"]
        # print(ans_dict_queid)
        caption = samples["captions"][0]
        answers = samples["answers"]
        Context_Prompt = ""
        mycontexts_id = []
        for idx in range(num_caps_per_img):
            cap_id_list = ans_dict_queid.get(
                answers[(len(answers) - 1 - idx) % len(answers)][:-1].lower(), [0]
            )
            for cap_id in cap_id_list:
                if cap_id not in mycontexts_id:
                    Context_Prompt += caption[cap_id]
                    mycontexts_id.append(cap_id)
                    break  # We just take one cap for each answer
        samples["Context_Prompt"] = Context_Prompt
        return Context_Prompt

    def create_task_prompt(
        self, samples, question_type="neural", num_question_per_img=30
    ):
        syn_question_queid = samples["questions"]
        syn_ans_queid = samples["answers"]
        Task_Prompt = ""
        for idx in range(num_question_per_img):
            # if config['random_question']:
            #     qa_idx = random.randint(0, len(syn_question_queid) - 1)
            # else:
            qa_idx = idx
            if (
                question_type != "rule" and num_question_per_img > 0 and idx < 1
            ):  ## yes and no questions for vqav2
                # Task_Prompt += "Question:"
                # Task_Prompt += syn_question_queid_next[-1]
                # Task_Prompt += '\n'
                # Task_Prompt += "Answer:no\n"
                Task_Prompt += "Question:"
                Task_Prompt += syn_question_queid[-1]
                Task_Prompt += "\n"
                Task_Prompt += "Answer:"
                Task_Prompt += "yes\n"
                Task_Prompt += "Question:Is this a toilet?\n"
                Task_Prompt += "Answer:no\n"
            if "question_type" == "rule":  # Rule-Based Question Generation
                Noun_Questions = [
                    "What item is this in this picture?",
                    "What item is that in this picture?",
                ]

                Verb_Questions = [
                    "What action is being done in this picture?",
                    "Why is this item doing in this picture?",
                    "Which action is being taken in this picture?",
                    "What action is item doing in this picture?",
                    "What action is item performing in this picture?",
                ]

                Adj_Questions = [
                    "How to describe one item in this picture?",
                    "What is item's ADJ TYPE in this picture?",
                    "What is the ADJ TYPE in this picture?",
                ]

                Task_Prompt += "Question:"
                doc = self.nlp(syn_ans_queid[(qa_idx) % len(syn_ans_queid)][:-1].lower())
                if doc[-1].pos_ == "NOUN":
                    Task_Prompt += Noun_Questions[
                        random.randint(0, len(Noun_Questions) - 1)
                    ]
                elif doc[-1].pos_ == "VERB":
                    Task_Prompt += Verb_Questions[
                        random.randint(0, len(Verb_Questions) - 1)
                    ]
                elif doc[-1].pos_ == "ADJ":
                    Task_Prompt += Adj_Questions[
                        random.randint(0, len(Adj_Questions) - 1)
                    ]

                Task_Prompt += "\n"

                Task_Prompt += "Answer:"
                Task_Prompt += syn_ans_queid[(qa_idx) % len(syn_ans_queid)][:-1].lower()
                Task_Prompt += "\n"
        samples["Task_Prompt"] = Task_Prompt
        # print(Task_Prompt)
        return Task_Prompt

    def prompts_construction(
        self,
        samples,
        question_type="neural",
        num_caps_per_img=30,
        num_question_per_img=30,
    ):
        Prompt = "Please reason the answer of the questions according to the given contexts.\n"

        Context_Prompt = self.create_context_prompt(samples, num_caps_per_img)

        Task_Prompt = self.create_task_prompt(
            samples, question_type, num_question_per_img
        )

        Img2Prompt = (
            Prompt
            + "Contexts:"
            + Context_Prompt
            + "\n"
            + Task_Prompt
            + "Question:"
            + samples["text_input"][0]
            + "\nAnswer:"
        )
        return Img2Prompt

    def prepare_LLM_input(
        self,
        samples,
        num_beams=1,
        inference_method="generate",
        max_len=20,
        min_len=0,
        internal_bsz_fid=1,
        num_captions=50,
        num_captions_fid=1,
        cap_max_length=20,
        cap_min_length=10,
        top_k=50,
        top_p=1,
        repetition_penalty=1,
        num_patches=20,
        block_num=7,
    ):
        """
        Args:
            samples (dict): A dictionary containing the following keys:
                - image (torch.Tensor): A tensor of shape (batch_size, 3, H, W). Default H=480, W=480.
                - text_input (str or [str]): String or a list of strings, each string is a question.
                                             The number of questions must be equal to the batch size. If a single string, will be converted to a list of string, with length 1 first.
            num_beams (int): Number of beams for beam search. 1 means no beam search.
            inference_method (str): Inference method. Must be "generate". The model will generate answers.
            max_len (int): Maximum length of generated answers.
            min_len (int): Minimum length of generated answers.
            internal_bsz_fid (int): Internal batch size when using FiD decoding.
            num_captions (int): Number of captions generated for each image.
            num_captions_fid (int): Number of captions concatenated with a question during FiD decoding.
            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_k (float): The number of the highest probability tokens for top-k sampling.
            top_p (float): The cumulative probability for nucleus sampling.
            repetition_penalty (float): The parameter for repetition penalty. 1.0 means no penalty.
            num_patches (int): Number of patches sampled for each image.
            block_num (int): The index of cross-attention block for gradcam computation.

        Returns:
            List: A list of strings, each string is an answer.
            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
        """
        assert inference_method in [
            "generate",
        ], "Inference method must be 'generate', got {}.".format(inference_method)

        if isinstance(samples["text_input"], str):
            samples["text_input"] = [samples["text_input"]]

        assert len(samples["text_input"]) == samples["image"].size(
            0
        ), "The number of questions must be equal to the batch size."

        samples = self.forward_itm(samples, block_num=block_num)

        samples = self.forward_cap(
            samples,
            cap_max_length=cap_max_length,
            cap_min_length=cap_min_length,
            top_k=top_k,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            num_captions=num_captions,
            num_patches=num_patches,
        )

        if self.offload_model:
            samples["image"] = samples["image"].to("cpu")
            self.image_question_matching_model.to("cpu")
            self.image_captioning_model.to("cpu")
        torch.cuda.empty_cache()

        pred_answers = self.forward_qa(
            samples,
            num_beams=num_beams,
            max_len=max_len,
            min_len=min_len,
            internal_bsz_fid=internal_bsz_fid,
            num_captions=num_captions,
            num_captions_fid=num_captions_fid,
        )

        if self.offload_model:
            self.image_question_matching_model.to(self.question_answering_model.device)
            self.image_captioning_model.to(self.question_answering_model.device)

        return pred_answers, samples["captions"], samples["gradcams"]

    @classmethod
    def from_config(cls, model_config):
        itm_config = model_config.image_question_matching_model
        cap_config = model_config.image_captioning_model

        itm_cls = registry.get_model_class(itm_config.arch)
        cap_cls = registry.get_model_class(cap_config.arch)

        image_question_matching_model = itm_cls.from_config(itm_config)
        image_captioning_model = cap_cls.from_config(cap_config)

        question_generation_tokenizer = T5Tokenizer.from_pretrained(
            "google/t5-large-lm-adapt"
        )
        question_generation_model = T5ForConditionalGeneration.from_pretrained(
            "google/t5-large-lm-adapt"
        )
        cached_file = download_cached_file(
            "https://storage.googleapis.com/sfr-vision-language-research/LAVIS/projects/img2prompt/T5_large_QG.pth",
            check_hash=False,
            progress=True,
        )
        checkpoint = torch.load(cached_file, map_location="cpu")
        state_dict = checkpoint["model"]
        question_generation_model.load_state_dict(state_dict)
        model = cls(
            image_question_matching_model=image_question_matching_model,
            image_captioning_model=image_captioning_model,
            question_generation_model=question_generation_model,
            question_generation_tokenizer=question_generation_tokenizer,
            offload_model=False,
        )

        return model