src/models/base.py

"""base.py.

Provides the common classes used such as the ModelSelection enum as well as the
abstract base class for models.
"""
import io
import logging
import os
import sqlite3
from abc import ABC, abstractmethod
from collections.abc import Iterator
from typing import Callable, List, Optional, TypedDict

import torch
import tqdm
from PIL import Image
from transformers import AutoProcessor
from transformers.feature_extraction_utils import BatchFeature

from .config import Config


class ModelInput(TypedDict):
    """Definition for the general model input dictionary."""
    image: str | Image.Image
    prompt: str
    label: Optional[str]
    data: BatchFeature
    row_id: Optional[str]


class ModelBase(ABC):
    """Provides an abstract base class for everything to implement."""

    def __init__(self, config: Config) -> None:
        """Initialization of the model base class.

        Args:
            config (Config): Parsed config.
        """
        self.model_path = config.model_path
        self.config = config

        # log the modules -- note that this causes an exit
        if self.config.log_named_modules:
            self._log_named_modules()
            exit(0)

        # load the specific model
        logging.debug(
            f'Loading model {self.config.architecture.value}; {self.model_path}'
        )
        self._load_specific_model()

        # load the processor
        self._init_processor()

    def _log_named_modules(self) -> None:
        """Logs the named modules based on the loaded model."""
        file_path = 'logs/' + self.model_path + '.txt'
        directory_path = os.path.dirname(file_path)

        # if the path exists to the file, don't load the model again
        if os.path.isfile(file_path):
            logging.debug(f'Named modules are cached in {file_path}')
            return

        # in which case, we first load the model, then output its modules
        self._load_specific_model()

        # otherwise, we log the output to that file, and creating directories
        # as needed
        if not os.path.exists(directory_path):
            os.makedirs(directory_path)

        with open(file_path, 'w') as output_file:
            output_file.writelines(
                [f'{name}\n' for name, _ in self.model.named_modules()]
            )

    @abstractmethod
    def _load_specific_model(self) -> None:
        """Abstract method that loads the specific model."""
        pass

    def _init_processor(self) -> None:
        """Initialize the self.processor by loading from the path."""
        self.processor = AutoProcessor.from_pretrained(self.model_path)

    def _generate_state_hook(self,
                             name: str,
                             model_input: ModelInput
                             ) -> Callable[[torch.nn.Module, tuple, torch.Tensor], None]:
        """Generates the state hook depending on the embedding type.

        Args:
            name (str): The module name.
            model_input (ModelInput): The input dictionary
                containing the image path, prompt, label (if applicable) and
                the data itself.

        Returns:
            hook function: The hook function to return.
        """
        image_path, prompt = model_input['image'], model_input['prompt']
        label = model_input.get('label', None)
        row_id = model_input.get('row_id', None)

        # Modify image path to be an absolute path if necessary
        if isinstance(image_path, str) and image_path != self.config.NO_IMG_PROMPT:
            image_path = os.path.abspath(image_path)

            # this image path should already exist, error out if someone isn't
            # properly providing an image path
            assert os.path.exists(image_path)

        def generate_states_hook(module: torch.nn.Module, input: tuple, output: torch.Tensor) -> None:
            """Hook handle function that saves the embedding output to a tensor.

            This tensor will be saved within a SQL database, according to the
            connection that was initialized previously.

            Args:
                module (torch.nn.Module): The module that save its hook on.
                input (tuple): The input used.
                output (torch.Tensor): The embeddings to save.
            """
            if not isinstance(output, torch.Tensor):
                logging.warning(f'Output type of {str(type(module))} is not a tensor, skipped.')
                return

            cursor = self.connection.cursor()

            # Convert the tensor to a binary blob
            tensor_blob = io.BytesIO()

            # It currently averages the output across the sequence length dimension, i.e., mean pooling
            # WARNING: When contributing new models, ensure that dim 1 is always the sequence length dimension
            final_output = getattr(output, self.config.pooling_method)(dim=1) if hasattr(
                self.config, 'pooling_method') and hasattr(output, self.config.pooling_method) else output
            output_dim = final_output.shape[-1]
            torch.save(final_output, tensor_blob)

            # Insert the tensor into the table
            cursor.execute(f"""
                INSERT INTO {self.config.DB_TABLE_NAME}
                (name, architecture, image_path, image_id, prompt, label, layer, pooling_method, tensor_dim, tensor)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
            """, (
                self.model_path,
                self.config.architecture.value,
                image_path if isinstance(image_path, str) else None,
                row_id,
                prompt,
                label,
                name,
                self.config.pooling_method if hasattr(self.config, 'pooling_method') else None,
                output_dim,
                tensor_blob.getvalue())
            )

            self.connection.commit()

            logging.debug(
                f'Ran hook and saved tensor for {image_path} using prompt '
                f'{prompt} on layer {name}.'
            )

        return generate_states_hook

    def _register_module_hooks(self,
                               model_input: ModelInput
                               ) -> List[torch.utils.hooks.RemovableHandle]:
        """Register the generated hook function to the modules in the config.

        At the same time, we need to add in the image path itself and the prompt
        which will be used for the database input.

        Args:
            model_input (ModelInput): The input dictionary
                containing the image path, prompt, label (if applicable) and
                the data itself.

        Raises:
            RuntimeError: Calls a runtime error if no hooks were registered

        Returns:
            List[torch.utils.hooks.RemovableHandle]: A list of handles that one
                can remove after the forward pass.
        """
        logging.debug(
            f'Registering module hook for {model_input["image"]} using prompt "{model_input["prompt"]}"'
        )

        # a list of hooks to remove after the forward pass
        hooks = []

        # for each module, register the state hook and save the output to database
        for name, module in self.model.named_modules():
            if self.config.matches_module(name):
                hooks.append(module.register_forward_hook(
                    self._generate_state_hook(name, model_input)
                ))
                logging.debug(f'Registered hook to {name}')

        if len(hooks) == 0:
            raise RuntimeError(
                'No hooks were registered. Double-check the configured modules.'
            )

        return hooks

    def _forward(self, data: BatchFeature) -> None:
        """Given some input data, performs a single forward pass.

        This function itself can be overriden, while _hook_and_eval
        should be left in tact.

        Args:
            data (BatchFeature): The given data tensor.
        """
        data.to(self.config.device)
        with torch.no_grad():
            _ = self.model(**data)
        logging.debug('Completed forward pass...')

    def _hook_and_eval(self, model_input: ModelInput) -> None:
        """Given some input, performs a single forward pass.

        Args:
            model_input (ModelInput): The given input dictionary.
        """
        logging.debug('Starting forward pass')
        self.model.eval()

        # now set up the modules to register the hook to
        hooks = self._register_module_hooks(model_input)

        # then ensure that the data is correct
        self._forward(model_input['data'])

        for hook in hooks:
            hook.remove()
        logging.debug('Unregistered all hooks..')

    def _initialize_db(self) -> None:
        """Initializes a database based on config."""
        # Connect to the database, creating it if it doesn't exist
        self.connection = sqlite3.connect(self.config.output_db)
        logging.debug(f'Database created at {self.config.output_db}')

        cursor = self.connection.cursor()

        # Create a table
        cursor.execute(
            f"""
                CREATE TABLE IF NOT EXISTS {self.config.DB_TABLE_NAME} (
                    id INTEGER PRIMARY KEY,
                    name TEXT NOT NULL,
                    architecture TEXT NOT NULL,
                    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                    image_path TEXT NULL,
                    image_id INTEGER NULL,
                    prompt TEXT NOT NULL,
                    label TEXT NULL,
                    layer TEXT NOT NULL,
                    pooling_method TEXT NULL,
                    tensor_dim INTEGER NOT NULL,
                    tensor BLOB NOT NULL
                );
            """
        )

    def _cleanup(self) -> None:
        """Cleanups the database by closing the connection."""
        self.connection.close()

    def _generate_processor_output(self, prompt: str, img_path: str | Image.Image) -> dict:
        """Generate the processor outputs from the prompt and image path.

        Args:
            prompt (str): The generated prompt string with the input text and
                the image labels.
            img_path (str | Image.Image): The specified input image path or image object.

        Returns:
            dict: The corresponding processor output per image and prompt.
        """
        data = {
            'text': prompt,
            'return_tensors': 'pt'
        }

        if img_path:
            img = Image.open(img_path) if isinstance(img_path, str) else img_path
            data['images'] = [img.convert('RGB')]

        return self.processor(**data)

    def _generate_prompt(self, prompt: str, add_generation_prompt: bool = True, has_images: bool = False) -> str:
        """Generates the prompt string with the input messages.

        TODO: move `add_generation_prompt` to the config.
        [Note from Martin] I'd argue that we should keep it as a parameter here
        since in gradio we want to hack these parameters a bit.

        Args:
            prompt (str): The input prompt string.
            add_generation_prompt (bool): Whether to add a start token of a bot
                response.
            has_images (bool): Whether the model has images or not.

        Returns:
            str: The generated prompt with the input text and the image labels.
        """
        logging.debug('Loading data...')
        # build the input dict for the chat template
        input_msgs_formatted = [{
            'role': 'user',
            'content': []
        }]

        # add the image if it exists
        if self.config.has_images() or has_images:
            input_msgs_formatted[0]['content'].append({
                'type': 'image'
            })

        # add the prompt if it exists
        if prompt:
            input_msgs_formatted[0]['content'].append({
                'type': 'text',
                'text': prompt
            })

        # apply the chat template to get the prompt
        return self.processor.apply_chat_template(
            input_msgs_formatted,
            add_generation_prompt=add_generation_prompt
        )

    def _load_input_data(self) -> Iterator[ModelInput]:
        """From a configuration, loads the input image and text data.

        For each prompt and input image, create a separate batch feature that
        will be ran separately and saved separately within the database.

        Yields:
            List[ModelInput]: List of input data, this input data is made of
                a tuple of strings (first an image path, then a prompt) and
                a batch feature which is either a torch.Tensor or a dictionary.
        """
        # by default use the processor, which may not exist for each model
        logging.debug('Generating embeddings through its processor...')
        if self.config.dataset:
            # Use the dataset to load input data, which includes (id, prompt, image_path)
            for row in self.config.dataset:
                prompt = self._generate_prompt(row['prompt'])
                data = self._generate_processor_output(
                    prompt=prompt,
                    img_path=row['image']
                )

                yield {
                    'image': row['image'],
                    'prompt': row['prompt'],
                    'label': row['label'] if 'label' in self.config.dataset.column_names else None,
                    'data': data,
                    'row_id': row['id'],
                }

        else:
            if not self.config.has_images():
                yield {
                    'image': self.config.NO_IMG_PROMPT,  # TODO: Check this?
                    'prompt': self.config.prompt,
                    'data': self._generate_processor_output(
                        prompt=self._generate_prompt(),
                        img_path=None
                    )
                }
            else:
                prompt = self._generate_prompt(self.config.prompt)
                for img_path in self.config.image_paths:
                    data = self._generate_processor_output(
                        prompt=prompt,
                        img_path=img_path
                    )
                    yield {
                        'image': img_path,
                        'prompt': self.config.prompt,
                        'data': data
                    }

    @property
    def _data_size(self) -> int:
        """Returns the total number of data points.

        Returns:
            int: The total number of data points.
        """
        if self.config.dataset:
            return len(self.config.dataset)
        else:
            if not self.config.has_images():
                return 1
            else:
                return len(self.config.image_paths)

    def run(self) -> None:
        """Get the hidden states from the model and saving them."""
        # let's first initialize a database connection
        self._initialize_db()

        # then convert to gpu
        self.model.to(self.config.device)

        # then reset the starting point in tracking maximum GPU memory, if using cuda
        if self.config.device.type == 'cuda':
            torch.cuda.reset_peak_memory_stats(self.config.device)

        # then run everything else
        for item in tqdm.tqdm(self._load_input_data(), desc='Running forward hooks on data', total=self._data_size):
            self._hook_and_eval(item)

        # then output peak memory usage, if using cuda
        if self.config.device.type == 'cuda':
            logging.debug(f'Peak GPU memory allocated: {torch.cuda.max_memory_allocated(self.config.device) / 1e6:.2f} MB')

        # finally clean up, closing database connection, etc.
        self._cleanup()