', add_special_tokens=False).input_ids[0]
+ self.config.model['seg_token_idx'] = self.tokenizer('[SEG]', add_special_tokens=False).input_ids[0]
+ self.config.model['bop_token_idx'] = self.tokenizer('', add_special_tokens=False).input_ids[0]
+ self.config.model['eop_token_idx'] = self.tokenizer('
', add_special_tokens=False).input_ids[0]
+
+ model_args = {
+ 'seg_token_idx': self.config.model['seg_token_idx'],
+ 'bbox_token_idx': self.config.model['bbox_token_idx'],
+ 'eop_token_idx': self.config.model['eop_token_idx'],
+ 'bop_token_idx': self.config.model['bop_token_idx'],
+ }
+
+ self.model = GLaMMForCausalLM.from_pretrained(
+ self.config.model_path,
+ torch_dtype=torch.bfloat16,
+ low_cpu_mem_usage=True,
+ **model_args
+ )
+ self.model = prepare_model_for_inference(self.model, self.config.model)
+
+ def _init_processor(self) -> None:
+ """Set the self.processor to follow the example given.
+
+ This should follow the processor setting and tokenizers under:
+ https://github.com/mbzuai-oryx/groundingLMM/blob/main/app.py
+ """
+ processor = {
+ 'global_enc_processor': CLIPImageProcessor.from_pretrained(self.config.model['vision_tower']),
+ 'grounding_transform': ResizeLongestSide(self.config.model['image_size'])
+ }
+ self.processor = processor
+
+ def _generate_prompt(self, prompt: str) -> str:
+ """Generates the GLaMM model prompt which will not use the chat template.
+
+ Args:
+ prompt (str): The input prompt string.
+
+ Returns:
+ str: The prompt to return, set by the config.
+ """
+ prompt = f'The {DEFAULT_IMAGE_TOKEN} provides an overview of the picture.\n{prompt}'
+ if self.config.model['use_mm_start_end']:
+ replace_token = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN
+ prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token)
+ return prompt
+
+ def _generate_processor_output(self, prompt: str, img_path: str) -> dict:
+ """Generate the processor argument to be input into the processor.
+
+ Args:
+ prompt (str): The generated prompt string with the input text and the image labels.
+ img_path (str): The specified image path.
+
+ Returns:
+ dict: The corresponding processor arguments per image and prompt.
+
+ Raises:
+ ValueError: If the image path is not defined.
+ """
+ if img_path is None:
+ raise ValueError('GLAMM cannot have text-only generation.')
+
+ image_np = cv2.imread(img_path)
+ image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
+ orig_h, orig_w = image_np.shape[:2]
+ original_size_list = [(orig_h, orig_w)]
+
+ # Global encoder
+ global_enc_image = self.processor['global_enc_processor'].preprocess(
+ image_np, return_tensors='pt')['pixel_values'][0].unsqueeze(0).cuda().bfloat16()
+
+ # Grounding encoder
+ grounding_input = self.processor['grounding_transform'].apply_image(image_np)
+ resize_list = [grounding_input.shape[:2]]
+ grounding_enc_image = grounding_enc_processor(
+ torch.from_numpy(grounding_input).permute(2, 0, 1).contiguous()
+ ).unsqueeze(0).cuda().bfloat16()
+
+ input_ids = tokenizer_image_token(prompt, self.tokenizer, return_tensors='pt').unsqueeze(0).cuda()
+
+ return {
+ 'input_ids': input_ids,
+ 'global_enc_image': global_enc_image,
+ 'grounding_enc_image': grounding_enc_image,
+ 'resize_list': resize_list,
+ 'original_size_list': original_size_list,
+ 'bboxes': None
+ }
+
+ def _forward(self, data: dict) -> 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 (dict): The given data tensor.
+ """
+ with torch.no_grad():
+ output_ids, _ = self.model.evaluate(
+ data['global_enc_image'],
+ data['grounding_enc_image'],
+ data['input_ids'],
+ data['resize_list'],
+ data['original_size_list'],
+ max_tokens_new=self.config.forward['max_new_tokens'],
+ bboxes=data['bboxes']
+ )
+ logging.debug('Completed forward pass')
diff --git a/src/models/internlm_xc/__init__.py b/src/models/internlm_xc/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..b77d034fe4831a0af67433cd64582b8ca220a89f
--- /dev/null
+++ b/src/models/internlm_xc/__init__.py
@@ -0,0 +1,9 @@
+"""__init__.py.
+
+InternLM package, includes which of this module to export.
+"""
+from .internlm_xc import InternLMXComposerModel
+
+__all__ = [
+ 'InternLMXComposerModel'
+]
diff --git a/src/models/internlm_xc/__pycache__/__init__.cpython-310.pyc b/src/models/internlm_xc/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b74103631a62bcc69f57007019e178550988866e
Binary files /dev/null and b/src/models/internlm_xc/__pycache__/__init__.cpython-310.pyc differ
diff --git a/src/models/internlm_xc/__pycache__/internlm_xc.cpython-310.pyc b/src/models/internlm_xc/__pycache__/internlm_xc.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7e7de4f488012356be5435d0c29d6f73975db649
Binary files /dev/null and b/src/models/internlm_xc/__pycache__/internlm_xc.cpython-310.pyc differ
diff --git a/src/models/internlm_xc/internlm_xc.py b/src/models/internlm_xc/internlm_xc.py
new file mode 100644
index 0000000000000000000000000000000000000000..5df015aa9c03f7b74b217bd2363d914a0bf00d9c
--- /dev/null
+++ b/src/models/internlm_xc/internlm_xc.py
@@ -0,0 +1,98 @@
+"""internlm.py.
+
+File for providing the InternLM-XComposer model implementation.
+"""
+import logging
+
+import torch
+from transformers import AutoModel, AutoProcessor
+
+from src.models.base import ModelBase
+from src.models.config import Config
+
+
+class InternLMXComposerModel(ModelBase):
+ """InternLM model implementation."""
+
+ def __init__(self, config: Config) -> None:
+ """Initialization of the InternLM model.
+
+ Args:
+ config (Config): Parsed config
+ """
+ # initialize the parent class
+ super().__init__(config)
+
+ def _load_specific_model(self) -> None:
+ """Overridden function to populate self.model."""
+ self.model = AutoModel.from_pretrained(
+ self.model_path,
+ trust_remote_code=True,
+ **self.config.model
+ ) if hasattr(self.config, 'model') else (
+ AutoModel.from_pretrained(
+ self.model_path,
+ trust_remote_code=True
+ )
+ )
+
+ def _init_processor(self) -> None:
+ """Overridden function to instantiate the model's processor."""
+ self.processor = AutoProcessor.from_pretrained(
+ self.model_path, trust_remote_code=True)
+ self.model.tokenizer = self.processor
+
+ def _generate_prompt(self, prompt: str) -> str:
+ """Overridden function to generate the prompt for the model.
+
+ Args:
+ prompt (str): The input prompt to be processed.
+
+ Returns:
+ str: The formatted prompt ready for model input.
+ """
+ return prompt
+
+ def _generate_processor_output(self, prompt: str, img_path: str) -> dict:
+ """Overridden function to generate the format the prompt for the processor.
+
+ Args:
+ prompt (str): The input prompt to be processed.
+ img_path (str): The path to the image to be processed.
+
+ Returns:
+ dict: The formatted inputs for the processor.
+
+ Raises:
+ ValueError: If no prompt is provided when required.
+ """
+ logging.debug('Loading data...')
+
+ # Manually format input as we do not need a processor
+ inputs = {}
+
+ # Text prompts are required for this model
+ if not prompt:
+ raise ValueError(
+ 'No input prompt was provided for the InternLM-XC model')
+
+ # If there are images, load them and add image token to prompt
+ if self.config.has_images():
+ inputs['query'] = f'; {prompt}'
+ inputs['image'] = [img_path]
+ else:
+ inputs['query'] = prompt
+
+ return inputs
+
+ def _forward(self, data: dict) -> None:
+ """Overridden function to run the model forward pass.
+
+ Args:
+ data (dict): The input data for the model.
+ """
+ device_type = str(self.config.device)
+ logging.debug(f'DATA: {data}')
+ with torch.autocast(device_type=device_type):
+ _, _ = self.model.chat(
+ self.processor, **data, **self.config.forward)
diff --git a/src/models/internvl/__init__.py b/src/models/internvl/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f392ac3ef4bcbc8aaed84ec2d341d9e33e142356
--- /dev/null
+++ b/src/models/internvl/__init__.py
@@ -0,0 +1,9 @@
+"""__init__.py.
+
+Clip package, includes which of this module to export.
+"""
+from .internvl import InternVLModel
+
+__all__ = [
+ 'InternVLModel'
+]
diff --git a/src/models/internvl/__pycache__/__init__.cpython-310.pyc b/src/models/internvl/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..db804025796f5dffdb90b3ddf472d694eff8f18e
Binary files /dev/null and b/src/models/internvl/__pycache__/__init__.cpython-310.pyc differ
diff --git a/src/models/internvl/__pycache__/internvl.cpython-310.pyc b/src/models/internvl/__pycache__/internvl.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..43a9f7d8806bd5d39e0badebb592e73457acebcd
Binary files /dev/null and b/src/models/internvl/__pycache__/internvl.cpython-310.pyc differ
diff --git a/src/models/internvl/__pycache__/utils.cpython-310.pyc b/src/models/internvl/__pycache__/utils.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..083673a177e765ee3b535fe1749c128f0d33a81e
Binary files /dev/null and b/src/models/internvl/__pycache__/utils.cpython-310.pyc differ
diff --git a/src/models/internvl/internvl.py b/src/models/internvl/internvl.py
new file mode 100644
index 0000000000000000000000000000000000000000..a45027dc641bbdf7891e589e8a084133b223c067
--- /dev/null
+++ b/src/models/internvl/internvl.py
@@ -0,0 +1,82 @@
+"""internvl.py.
+
+File for providing the Intern-VL model implementation.
+"""
+
+import logging
+
+import torch
+from transformers import AutoModel, AutoTokenizer
+from transformers.feature_extraction_utils import BatchFeature
+
+from src.models.base import ModelBase
+from src.models.config import Config
+from src.models.internvl.utils import load_image
+
+
+class InternVLModel(ModelBase):
+ """InternVL model implementation."""
+
+ def __init__(self, config: Config) -> None:
+ """Initialization of the InternVL model.
+
+ Args:
+ config (Config): Parsed config
+ """
+ # initialize the parent class
+ super().__init__(config)
+
+ def _load_specific_model(self) -> None:
+ """Overridden function to populate self.model."""
+ self.model = AutoModel.from_pretrained(
+ self.model_path, **self.config.model
+ ) if hasattr(self.config, 'model') else (
+ AutoModel.from_pretrained(
+ self.model_path
+ )
+ )
+
+ def _generate_prompt(self, prompt: str) -> str:
+ """Generates the InternVL model prompt which will not use the chat template.
+
+ Args:
+ prompt (str): The prompt to return, set by the config.
+
+ Returns:
+ str: The prompt to return, set by the config.
+ """
+ return prompt
+
+ def _init_processor(self) -> None:
+ """Initialize the InternVL processor which need to be done manually."""
+ self.processor = None # no intended processor here
+ self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True, use_fast=False)
+ self.img_processor = load_image
+
+ def _generate_processor_output(self, prompt: str, img_path: str) -> 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): The specified image path.
+
+ Returns:
+ dict: The corresponding processor output per image and prompt.
+ """
+ return {'prompt': prompt,
+ 'pixel_values': None if img_path is None else self.img_processor(img_path, max_num=12).to(dtype=torch.bfloat16).to(self.config.device)}
+
+ 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.
+ """
+ generation_config = self.config.forward
+ with torch.no_grad():
+ _ = self.model.chat(self.tokenizer, data['pixel_values'], data['prompt'], generation_config)
+ logging.debug('Completed forward pass...')
diff --git a/src/models/internvl/utils.py b/src/models/internvl/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..0978a3201b10cadfc56b2e17f5ff631321a5df7c
--- /dev/null
+++ b/src/models/internvl/utils.py
@@ -0,0 +1,135 @@
+"""Helper functions from official huggingface library of InternVL."""
+
+from typing import List, Optional, Tuple
+
+import torch
+import torchvision.transforms as T
+from PIL import Image
+from torchvision.transforms.functional import InterpolationMode
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+
+def build_transform(input_size: Optional[int] = 448) -> T.Compose:
+ """Helper function that transform image.
+
+ Args:
+ input_size (int, optional): The input size. Defaults to 448.
+
+ Returns:
+ T.Compose: The composed transform.
+ """
+ MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
+ return T.Compose([
+ T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
+ T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
+ T.ToTensor(),
+ T.Normalize(mean=MEAN, std=STD)
+ ])
+
+
+def find_closest_aspect_ratio(
+ aspect_ratio: float, target_ratios: List[Tuple[float, float]],
+ width: int, height: int, image_size: int) -> Tuple[int, int]:
+ """Helper function that find closest aspect ratio.
+
+ Args:
+ aspect_ratio (float): The existing image aspect ratio.
+ target_ratios (list): The target aspect ratios.
+ width (int): The original image width.
+ height (int): The original image height.
+ image_size (int): The target image size.
+
+ Returns:
+ tuple: The closest aspect ratio.
+ """
+ best_ratio_diff = float('inf')
+ best_ratio = (1, 1)
+ area = width * height
+ for ratio in target_ratios:
+ target_aspect_ratio = ratio[0] / ratio[1]
+ ratio_diff = abs(aspect_ratio - target_aspect_ratio)
+ if ratio_diff < best_ratio_diff:
+ best_ratio_diff = ratio_diff
+ best_ratio = ratio
+ elif ratio_diff == best_ratio_diff:
+ if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
+ best_ratio = ratio
+ return best_ratio
+
+
+def dynamic_preprocess(
+ image: Image, min_num: Optional[int] = 1, max_num: Optional[int] = 12,
+ image_size: Optional[int] = 448, use_thumbnail: Optional[bool] = False) -> List[Image]:
+ """Helper function.
+
+ Args:
+ image (Image): The input image.
+ min_num (int, optional): The minimum number of image patches. Defaults to 1.
+ max_num (int, optional): The maximum number of image patches. Defaults to 12.
+ image_size (int, optional): The target image size. Defaults to 448.
+ use_thumbnail (bool, optional): Whether to use thumbnail. Defaults to False.
+
+ Returns:
+ list: The processed images.
+ """
+ orig_width, orig_height = image.size
+ aspect_ratio = orig_width / orig_height
+
+ # calculate the existing image aspect ratio
+ target_ratios = {
+ (i, j)
+ for n in range(min_num, max_num + 1)
+ for i in range(1, n + 1)
+ for j in range(1, n + 1)
+ if i * j <= max_num and i * j >= min_num
+ }
+ target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+ # find the closest aspect ratio to the target
+ target_aspect_ratio = find_closest_aspect_ratio(
+ aspect_ratio, target_ratios, orig_width, orig_height, image_size)
+
+ # calculate the target width and height
+ target_width = image_size * target_aspect_ratio[0]
+ target_height = image_size * target_aspect_ratio[1]
+ blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
+
+ # resize the image
+ resized_img = image.resize((target_width, target_height))
+ processed_images = []
+ for i in range(blocks):
+ box = (
+ (i % (target_width // image_size)) * image_size,
+ (i // (target_width // image_size)) * image_size,
+ ((i % (target_width // image_size)) + 1) * image_size,
+ ((i // (target_width // image_size)) + 1) * image_size
+ )
+ # split the image
+ split_img = resized_img.crop(box)
+ processed_images.append(split_img)
+ assert len(processed_images) == blocks
+ if use_thumbnail and len(processed_images) != 1:
+ thumbnail_img = image.resize((image_size, image_size))
+ processed_images.append(thumbnail_img)
+ return processed_images
+
+
+def load_image(image_file: str, input_size: Optional[int] = 448, max_num: Optional[int] = 12) -> torch.Tensor:
+ """Load image to pixel values.
+
+ Args:
+ image_file (str): The image file path.
+ input_size (int, optional): The input size. Defaults to 448.
+ max_num (int, optional): The max number of image patches. Defaults to 12.
+
+ Returns:
+ torch.Tensor: The corresponding pixel values.
+ """
+ image = Image.open(image_file).convert('RGB')
+ transform = build_transform(input_size=input_size)
+ images = dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
+ pixel_values = [transform(image) for image in images]
+ pixel_values = torch.stack(pixel_values)
+ return pixel_values
diff --git a/src/models/janus/__init__.py b/src/models/janus/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..cb320a1c9f10c4aea57a1e596b9f259545fe3782
--- /dev/null
+++ b/src/models/janus/__init__.py
@@ -0,0 +1,9 @@
+"""__init__.py.
+
+Janus package, includes which of this module to export.
+"""
+from .janus import JanusModel
+
+__all__ = [
+ 'JanusModel'
+]
diff --git a/src/models/janus/__pycache__/__init__.cpython-310.pyc b/src/models/janus/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..53515ac06ead83dc2e79d631b3ceec704d1dfe89
Binary files /dev/null and b/src/models/janus/__pycache__/__init__.cpython-310.pyc differ
diff --git a/src/models/janus/__pycache__/janus.cpython-310.pyc b/src/models/janus/__pycache__/janus.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..57d0b08d9992d21ab6e727ff45cfa0ee8ea39190
Binary files /dev/null and b/src/models/janus/__pycache__/janus.cpython-310.pyc differ
diff --git a/src/models/janus/janus.py b/src/models/janus/janus.py
new file mode 100644
index 0000000000000000000000000000000000000000..f2e2a7172fe3a3dc2109ddf4cdae136063916719
--- /dev/null
+++ b/src/models/janus/janus.py
@@ -0,0 +1,93 @@
+"""janus.py.
+
+File for providing the Janus model implementation.
+"""
+import torch
+from transformers import JanusForConditionalGeneration, JanusProcessor
+
+from src.models.base import ModelBase
+from src.models.config import Config
+
+
+class JanusModel(ModelBase):
+ """Janus model implementation."""
+
+ def __init__(self, config: Config) -> None:
+ """Initialize the Janus model.
+
+ Args:
+ config (Config): Parsed config.
+ """
+ super().__init__(config)
+
+ def _load_specific_model(self) -> None:
+ """Populate self.model with the specified Janus model."""
+ # require this import to force the models script to load
+ self.model = (
+ JanusForConditionalGeneration.from_pretrained(
+ self.model_path,
+ **self.config.model
+ ) if hasattr(self.config, 'model') else
+ JanusForConditionalGeneration.from_pretrained(
+ self.model_path,
+ )
+ )
+ self.model.to(torch.bfloat16)
+
+ def _init_processor(self) -> None:
+ """Initialize the Janus processor."""
+ self.processor = JanusProcessor.from_pretrained(self.model_path)
+
+ def _generate_prompt(self, prompt: str) -> str:
+ """Generates the prompt string with the input messages.
+
+ Args:
+ prompt (str): prompt content.
+
+ Returns:
+ str: Returns the prompt content as is.
+ """
+ return prompt
+
+ def _generate_processor_output(self, prompt: str, img_path: str) -> dict:
+ """Override the base function to produce processor arguments for Janus.
+
+ Args:
+ prompt (str): The input prompt to be processed.
+ img_path (str): The path to the image to be processed.
+
+ Returns:
+ dict: The formatted inputs for the processor.
+ """
+ # Do the _generate_prompt first
+ messages = [
+ {
+ 'role': 'user',
+ 'content': [
+ {'type': 'image', 'image': img_path},
+ {'type': 'text', 'text': prompt}
+ ]
+ }
+ ]
+
+ inputs = self.processor.apply_chat_template(
+ messages,
+ add_generation_prompt=True,
+ generation_mode='text',
+ tokenize=True,
+ return_dict=True,
+ return_tensors='pt',
+ ).to(self.config.device, dtype=torch.bfloat16)
+
+ return inputs
+
+ def _forward(self, data: dict) -> 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 (dict): The given data tensor.
+ """
+ _ = self.model.generate(**data, **self.config.forward)
diff --git a/src/models/llava/__init__.py b/src/models/llava/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..dbbb07fb4e8fe3fa4429c5315103cbf551220e13
--- /dev/null
+++ b/src/models/llava/__init__.py
@@ -0,0 +1,9 @@
+"""__init__.py.
+
+Llava package, includes which of this module to export.
+"""
+from .llava import LlavaModel
+
+__all__ = [
+ 'LlavaModel'
+]
diff --git a/src/models/llava/__pycache__/__init__.cpython-310.pyc b/src/models/llava/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b17ea03ab5a6bee3ca11d2fe6bc6e70e13deae6c
Binary files /dev/null and b/src/models/llava/__pycache__/__init__.cpython-310.pyc differ
diff --git a/src/models/llava/__pycache__/llava.cpython-310.pyc b/src/models/llava/__pycache__/llava.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..86e63076a3b4b435bda8c95551faafcba70ab50c
Binary files /dev/null and b/src/models/llava/__pycache__/llava.cpython-310.pyc differ
diff --git a/src/models/llava/llava.py b/src/models/llava/llava.py
new file mode 100644
index 0000000000000000000000000000000000000000..effda4dafb07c7dcd36f49af5198c2e3d79bfcc5
--- /dev/null
+++ b/src/models/llava/llava.py
@@ -0,0 +1,31 @@
+"""llava.py.
+
+File for providing the Llava model implementation.
+"""
+from transformers import LlavaForConditionalGeneration
+
+from src.models.base import ModelBase
+from src.models.config import Config
+
+
+class LlavaModel(ModelBase):
+ """Llava model implementation."""
+
+ def __init__(self, config: Config) -> None:
+ """Initialization of the llava model.
+
+ Args:
+ config (Config): Parsed config
+ """
+ # initialize the parent class
+ super().__init__(config)
+
+ def _load_specific_model(self) -> None:
+ """Overridden function to populate self.model."""
+ self.model = LlavaForConditionalGeneration.from_pretrained(
+ self.model_path, **self.config.model
+ ) if hasattr(self.config, 'model') else (
+ LlavaForConditionalGeneration.from_pretrained(
+ self.model_path
+ )
+ )
diff --git a/src/models/minicpm/__init__.py b/src/models/minicpm/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..bde64cfe6a8e1e841a67a6b97f61d5c5b7e05394
--- /dev/null
+++ b/src/models/minicpm/__init__.py
@@ -0,0 +1,9 @@
+"""__init__.py.
+
+MiniCPM package, includes which of this module to export.
+"""
+from .minicpm import MiniCPMModel
+
+__all__ = [
+ 'MiniCPMModel'
+]
diff --git a/src/models/minicpm/__pycache__/__init__.cpython-310.pyc b/src/models/minicpm/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..72d34dd3cfb3e80f0ad6b4e159a5ee7b697b8272
Binary files /dev/null and b/src/models/minicpm/__pycache__/__init__.cpython-310.pyc differ
diff --git a/src/models/minicpm/__pycache__/minicpm.cpython-310.pyc b/src/models/minicpm/__pycache__/minicpm.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7b7612896110f7f019736e310ca86080ede2d984
Binary files /dev/null and b/src/models/minicpm/__pycache__/minicpm.cpython-310.pyc differ
diff --git a/src/models/minicpm/minicpm.py b/src/models/minicpm/minicpm.py
new file mode 100644
index 0000000000000000000000000000000000000000..41aea8d53d3315c2de693335d4db160895954ca6
--- /dev/null
+++ b/src/models/minicpm/minicpm.py
@@ -0,0 +1,77 @@
+"""minicpm.py.
+
+File for providing the MiniCPM model implementation.
+"""
+
+import logging
+
+import torch
+from PIL import Image
+from transformers import AutoModel, AutoTokenizer
+from transformers.feature_extraction_utils import BatchFeature
+
+from src.models.base import ModelBase
+from src.models.config import Config
+
+
+class MiniCPMModel(ModelBase):
+ """MiniCPM model implementation."""
+
+ def __init__(self, config: Config) -> None:
+ """Initialization of the MiniCPM model.
+
+ Args:
+ config (Config): Parsed config
+ """
+ # initialize the parent class
+ super().__init__(config)
+
+ def _load_specific_model(self) -> None:
+ """Overridden function to populate self.model."""
+ self.model = AutoModel.from_pretrained(
+ self.model_path, **getattr(self.config, 'model', {})
+ )
+
+ def _generate_prompt(self, prompt: str) -> str:
+ """Generates the MiniCPM model prompt which will not use the chat template.
+
+ Args:
+ prompt (str): The prompt content.
+
+ Returns:
+ str: The prompt to return, set by the config.
+ """
+ return prompt
+
+ def _init_processor(self) -> None:
+ """Initialize the MiniCPM tokenizer."""
+ self.processor = None # no intended processor here
+ self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)
+
+ def _generate_processor_output(self, prompt: str, img_path: str) -> 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): The specified image path.
+
+ Returns:
+ dict: The corresponding processor output per image and prompt.
+ """
+ msgs = [{'role': 'user', 'content': prompt}]
+ image = Image.open(img_path).convert('RGB')
+ return {'msgs': msgs, 'image': image}
+
+ 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.
+ """
+ with torch.no_grad():
+ _ = self.model.chat(**data, context=None, tokenizer=self.tokenizer, **self.config.forward)
+ logging.debug('Completed forward pass...')
diff --git a/src/models/molmo/__init__.py b/src/models/molmo/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..a7f568316132ef42981332c299b374145dc791ee
--- /dev/null
+++ b/src/models/molmo/__init__.py
@@ -0,0 +1,9 @@
+"""__init__.py.
+
+Molmo package, includes which of this module to export.
+"""
+from .molmo import MolmoModel
+
+__all__ = [
+ 'MolmoModel'
+]
diff --git a/src/models/molmo/__pycache__/__init__.cpython-310.pyc b/src/models/molmo/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0cb4cda185aa0603cb293a3610ba999d4dd894c3
Binary files /dev/null and b/src/models/molmo/__pycache__/__init__.cpython-310.pyc differ
diff --git a/src/models/molmo/__pycache__/molmo.cpython-310.pyc b/src/models/molmo/__pycache__/molmo.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4fef7096313167fed99868e1211d39d48e072a4b
Binary files /dev/null and b/src/models/molmo/__pycache__/molmo.cpython-310.pyc differ
diff --git a/src/models/molmo/molmo.py b/src/models/molmo/molmo.py
new file mode 100644
index 0000000000000000000000000000000000000000..ae450c240f81109d8dee4407b81ae4ac77a557dc
--- /dev/null
+++ b/src/models/molmo/molmo.py
@@ -0,0 +1,101 @@
+"""molmo.py.
+
+File for providing the Molmo model implementation.
+"""
+import logging
+
+import torch
+from PIL import Image
+from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
+
+from src.models.base import ModelBase
+from src.models.config import Config
+
+
+class MolmoModel(ModelBase):
+ """Molmo model implementation."""
+
+ def __init__(self, config: Config) -> None:
+ """Initialization of the molmo model.
+
+ Args:
+ config (Config): Parsed config
+ """
+ # initialize the parent class
+ super().__init__(config)
+
+ def _load_specific_model(self) -> None:
+ """Overridden function to populate self.model."""
+ self.model = AutoModelForCausalLM.from_pretrained(
+ self.model_path, **getattr(self.config, 'model', {}), trust_remote_code=True
+ )
+
+ def _init_processor(self) -> None:
+ """Initializes the processor."""
+ self.processor = AutoProcessor.from_pretrained(
+ self.config.model_path, **getattr(self.config, 'model', {}), trust_remote_code=True
+ )
+
+ def _generate_prompt(self, prompt: str, add_generation_prompt: bool = True, has_images: bool = False) -> str:
+ """Generates the Molmo model prompt which will not use the chat template.
+
+ [Note from Martin] I'd hack these parameters a bit for gradio, follow Base.
+
+ Args:
+ prompt (str): The prompt to return, set by the config.
+ 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 prompt to return, set by the config.
+ """
+ return prompt
+
+ def _generate_processor_output(self, prompt: str, img_path: str) -> dict:
+ """Generate the processor argument to be input into the processor.
+
+ Args:
+ prompt (str): The generated prompt string with the input text and
+ the image labels.
+ img_path (str): The specified image path.
+
+ Returns:
+ dict: The corresponding processor arguments per image and prompt.
+
+ Raises:
+ ValueError: If no prompt is provided when required.
+ """
+ if img_path is None:
+ raise ValueError('Molmo cannot have text-only generation.')
+
+ # prepare the data inputs according to
+ # https://huggingface.co/allenai/Molmo-7B-D-0924
+ data_inputs = self.processor.process(
+ images=[Image.open(img_path)],
+ text=prompt
+ )
+
+ # move inputs to the correct device and make a batch of size 1
+ return {
+ k: v.to(self.config.device).unsqueeze(0)
+ for k, v in data_inputs.items()
+ }
+
+ def _forward(self, data: dict) -> 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 (dict): The given data tensor.
+ """
+ generation_config = self.config.forward
+ with torch.no_grad():
+ _ = self.model.generate_from_batch(
+ data,
+ GenerationConfig(**generation_config),
+ tokenizer=self.processor.tokenizer
+ )
+ logging.debug('Completed forward pass...')
diff --git a/src/models/paligemma/__init__.py b/src/models/paligemma/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..eb24b8656af6ac0e9bbde3cd1d0eb4f9c6b95419
--- /dev/null
+++ b/src/models/paligemma/__init__.py
@@ -0,0 +1,9 @@
+"""__init__.py.
+
+Paligemma package, includes which of this module to export.
+"""
+from .paligemma import PaligemmaModel
+
+__all__ = [
+ 'PaligemmaModel'
+]
diff --git a/src/models/paligemma/__pycache__/__init__.cpython-310.pyc b/src/models/paligemma/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..73790d75a7f4fd1750f35e13875785777741f15c
Binary files /dev/null and b/src/models/paligemma/__pycache__/__init__.cpython-310.pyc differ
diff --git a/src/models/paligemma/__pycache__/paligemma.cpython-310.pyc b/src/models/paligemma/__pycache__/paligemma.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e09d0edbc20b43e8afe290163474a60ba0e34f28
Binary files /dev/null and b/src/models/paligemma/__pycache__/paligemma.cpython-310.pyc differ
diff --git a/src/models/paligemma/paligemma.py b/src/models/paligemma/paligemma.py
new file mode 100644
index 0000000000000000000000000000000000000000..15cce0f9f8664ebc92549011349cf1abfda700ca
--- /dev/null
+++ b/src/models/paligemma/paligemma.py
@@ -0,0 +1,56 @@
+"""paligemma.py.
+
+File for providing the Paligemma model implementation.
+"""
+from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
+
+from src.models.base import ModelBase
+from src.models.config import Config
+
+
+class PaligemmaModel(ModelBase):
+ """PaligemmaModel model implementation."""
+
+ def __init__(self, config: Config) -> None:
+ """Initialization of the paligemma model.
+
+ Args:
+ config (Config): Parsed config
+ """
+ # initialize the parent class
+ super().__init__(config)
+
+ def _load_specific_model(self) -> None:
+ """Overridden function to populate Paligemma model.
+
+ Huggingface token is required to get access to the model.
+ Replace in configs/paligemma-3b.yaml file with you own hugging face security token.
+ Note: 'token' is a general Hugging Face Hub access token, not specific to PaliGemma.
+ It enables loading private models or authenticated access.
+ See: https://huggingface.co/docs/hub/en/security-tokens
+ """
+ self.model = PaliGemmaForConditionalGeneration.from_pretrained(
+ self.model_path, **self.config.model
+ )
+
+ def _init_processor(self) -> None:
+ """Initialize the Paligemma processor.
+
+ Huggingface token is required.
+ Replace in configs/paligemma-3b.yaml file with you own hugging face security token.
+ Note: 'token' is a general Hugging Face Hub access token, not specific to PaliGemma.
+ It enables loading private models or authenticated access.
+ See: https://huggingface.co/docs/hub/en/security-tokens
+ """
+ self.processor = AutoProcessor.from_pretrained(self.model_path, token=self.config.model['token'])
+
+ def _generate_prompt(self, prompt: str) -> str:
+ """Generates the Paligemma model prompt which will not use the chat template.
+
+ Args:
+ prompt (str): The input prompt for the model.
+
+ Returns:
+ str: The prompt to return, set by the config.
+ """
+ return prompt
diff --git a/src/models/pixtral/__init__.py b/src/models/pixtral/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c9c3d663b263d35cfc5eb3c7f29116c9015f0dcf
--- /dev/null
+++ b/src/models/pixtral/__init__.py
@@ -0,0 +1,9 @@
+"""__init__.py.
+
+Pixtral package, includes which of this module to export.
+"""
+from .pixtral import PixtralModel
+
+__all__ = [
+ 'PixtralModel'
+]
diff --git a/src/models/pixtral/__pycache__/__init__.cpython-310.pyc b/src/models/pixtral/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7c08c4b988f05b0887f5754086223643e2e599de
Binary files /dev/null and b/src/models/pixtral/__pycache__/__init__.cpython-310.pyc differ
diff --git a/src/models/pixtral/__pycache__/pixtral.cpython-310.pyc b/src/models/pixtral/__pycache__/pixtral.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..395dcd02837cfd9a79ab6e9a9d13687dc4ae432a
Binary files /dev/null and b/src/models/pixtral/__pycache__/pixtral.cpython-310.pyc differ
diff --git a/src/models/pixtral/pixtral.py b/src/models/pixtral/pixtral.py
new file mode 100644
index 0000000000000000000000000000000000000000..681e3ebd3cd8a81a073b5c4117bf6ecc37752cbf
--- /dev/null
+++ b/src/models/pixtral/pixtral.py
@@ -0,0 +1,103 @@
+"""pixtral.py.
+
+File for providing the Pixtral model implementation.
+"""
+import logging
+
+import torch
+from huggingface_hub import snapshot_download
+from mistral_common.protocol.instruct.messages import (ImageChunk, TextChunk,
+ UserMessage)
+from mistral_common.protocol.instruct.request import ChatCompletionRequest
+from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+from mistral_inference.transformer import Transformer
+from PIL import Image
+
+from src.models.base import ModelBase
+from src.models.config import Config
+
+
+class PixtralModel(ModelBase):
+ """Pixtral model implementation."""
+
+ def __init__(self, config: Config) -> None:
+ """Initialization of the Pixtral model.
+
+ Args:
+ config (Config): Parsed config
+ """
+ # initialize the parent class
+ super().__init__(config)
+
+ def _load_specific_model(self) -> None:
+ """Overridden function to populate self.model."""
+ snapshot_download(
+ repo_id=self.model_path,
+ allow_patterns=['params.json', 'consolidated.safetensors', 'tekken.json'],
+ local_dir=self.config.download_path,
+ )
+
+ self.model = Transformer.from_folder(self.config.download_path, **getattr(self.config, 'model', {}))
+
+ def _generate_prompt(self, prompt: str) -> str:
+ """Generates the Pixtral model prompt which will not use the chat template.
+
+ Args:
+ prompt (str): The input prompt for the model.
+
+ Returns:
+ str: The prompt to return, set by the config.
+ """
+ return prompt
+
+ def _init_processor(self) -> None:
+ """Initialize the Pixtral Tokenizer."""
+ self.processor = None # no intended processor here
+ self.tokenizer = MistralTokenizer.from_file(f'{self.config.download_path}/tekken.json')
+
+ def _generate_processor_output(self, prompt: str, img_path: str | None) -> dict:
+ """Generate the processor outputs from the prompt and image path.
+
+ Pixtral uses a specific chat template format with special image tokens.
+
+ Args:
+ prompt (str): The generated prompt string with the input text and
+ the image labels.
+ img_path (str or None): The specified image path, or None for text-only.
+
+ Returns:
+ dict: The corresponding processor output per image and prompt.
+ """
+ user_content = [TextChunk(text=prompt)]
+ if img_path is not None:
+ image = Image.open(img_path)
+ user_content = [ImageChunk(image=image)] + user_content
+
+ completion_request = ChatCompletionRequest(messages=[UserMessage(content=user_content)])
+ encoded = self.tokenizer.encode_chat_completion(completion_request)
+
+ res = {
+ 'input_ids': torch.tensor(encoded.tokens, dtype=torch.long, device=self.model.device),
+ 'seqlens': [len(encoded.tokens)],
+ }
+
+ if img_path is not None:
+ res['images'] = [
+ torch.tensor(img, device=self.model.device, dtype=self.model.dtype)
+ for img in encoded.images
+ ]
+
+ return res
+
+ def _forward(self, data: dict) -> 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 (dict): The given data tensor.
+ """
+ with torch.no_grad():
+ _ = self.model.forward(**data)
+ logging.debug('Completed forward pass...')
diff --git a/src/models/plm/__init__.py b/src/models/plm/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..89a2c7992adb0bfdecf304cd00b6f0e4e1168159
--- /dev/null
+++ b/src/models/plm/__init__.py
@@ -0,0 +1,9 @@
+"""__init__.py.
+
+Perception Language Model(Plm) package, includes which of this module to export.
+"""
+from .plm import PlmModel
+
+__all__ = [
+ 'PlmModel'
+]
diff --git a/src/models/plm/__pycache__/__init__.cpython-310.pyc b/src/models/plm/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..64f764c808697c778275d3f07b528097d69f7db6
Binary files /dev/null and b/src/models/plm/__pycache__/__init__.cpython-310.pyc differ
diff --git a/src/models/plm/__pycache__/plm.cpython-310.pyc b/src/models/plm/__pycache__/plm.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..dbb27b035d0454f4b49c2405afb0a501235d76ff
Binary files /dev/null and b/src/models/plm/__pycache__/plm.cpython-310.pyc differ
diff --git a/src/models/plm/plm.py b/src/models/plm/plm.py
new file mode 100644
index 0000000000000000000000000000000000000000..7ee55f644608ce30bef0a185dfa9c10e349391b8
--- /dev/null
+++ b/src/models/plm/plm.py
@@ -0,0 +1,55 @@
+"""plm.py.
+
+File for providing the Plm model implementation.
+"""
+
+import logging
+
+import torch
+from transformers import AutoModelForImageTextToText, AutoProcessor
+from transformers.feature_extraction_utils import BatchFeature
+
+from src.models.base import ModelBase
+from src.models.config import Config
+
+
+class PlmModel(ModelBase):
+ """PLM model implementation."""
+
+ def __init__(self, config: Config) -> None:
+ """Initialization of the PLM model.
+
+ Args:
+ config (Config): Parsed config
+ """
+ # initialize the parent class
+ super().__init__(config)
+
+ def _load_specific_model(self) -> None:
+ """Overridden function to populate self.model."""
+ self.model = AutoModelForImageTextToText.from_pretrained(
+ self.model_path, **self.config.model
+ ) if hasattr(self.config, 'model') else (
+ AutoModelForImageTextToText.from_pretrained(
+ self.model_path
+ )
+ )
+ self.model.to(self.config.device)
+
+ def _init_processor(self) -> None:
+ """Initialize the self.processor by loading from the path."""
+ self.processor = AutoProcessor.from_pretrained(self.model_path, use_fast=True)
+
+ 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.generate(**data, **self.config.forward)
+ logging.debug('Completed forward pass...')
diff --git a/src/models/qwen/__init__.py b/src/models/qwen/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..3a858b45d7427abe9198bea7ed6ffc24543252eb
--- /dev/null
+++ b/src/models/qwen/__init__.py
@@ -0,0 +1,9 @@
+"""__init__.py.
+
+Qwen package, includes which of this module to export.
+"""
+from .qwen import QwenModel
+
+__all__ = [
+ 'QwenModel'
+]
diff --git a/src/models/qwen/__pycache__/__init__.cpython-310.pyc b/src/models/qwen/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9e6de312f596394e271a61ae86be82b9e4273ffb
Binary files /dev/null and b/src/models/qwen/__pycache__/__init__.cpython-310.pyc differ
diff --git a/src/models/qwen/__pycache__/qwen.cpython-310.pyc b/src/models/qwen/__pycache__/qwen.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..025fe53bf75cd191f9a9a8d31a487299a901966e
Binary files /dev/null and b/src/models/qwen/__pycache__/qwen.cpython-310.pyc differ
diff --git a/src/models/qwen/qwen.py b/src/models/qwen/qwen.py
new file mode 100644
index 0000000000000000000000000000000000000000..11625b4561497a9eda4c6573ea4a6a67f1a8a132
--- /dev/null
+++ b/src/models/qwen/qwen.py
@@ -0,0 +1,31 @@
+"""qwen.py.
+
+File for providing the Qwen model implementation.
+"""
+from transformers import Qwen2VLForConditionalGeneration
+
+from src.models.base import ModelBase
+from src.models.config import Config
+
+
+class QwenModel(ModelBase):
+ """Qwen model implementation."""
+
+ def __init__(self, config: Config) -> None:
+ """Initialization of the qwen model.
+
+ Args:
+ config (Config): Parsed config
+ """
+ # initialize the parent class
+ super().__init__(config)
+
+ def _load_specific_model(self) -> None:
+ """Overridden function to populate self.model."""
+ self.model = Qwen2VLForConditionalGeneration.from_pretrained(
+ self.model_path, **self.config.model
+ ) if hasattr(self.config, 'model') else (
+ Qwen2VLForConditionalGeneration.from_pretrained(
+ self.model_path
+ )
+ )
diff --git a/src/probe/README.md b/src/probe/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..c2d729fccef69f963891083861b40450ddbba735
--- /dev/null
+++ b/src/probe/README.md
@@ -0,0 +1,77 @@
+# Probe Implementation
+Using the output database of extracted features (generated using `src/main.py`), you can easily intialize probes using the script in `probe/main.py`.
+
+To do this, you need a probe config `yaml` file specifying: the probe `model` attributes, `training` and `test` configuration and the input `data`. The following is an example of a probe configuration file for features extracted using LLAVA-1.5-7b:
+```yaml
+model:
+ - activation: ReLU # a valid activation function from torch.nn
+ - hidden_size: 512 # the input and output size of the intermediate layers
+ - num_layers: 2 # the number of layers of the probe model
+ - save_dir: /path/to/save_dir # the location to save the probe results
+
+training:
+ - batch_size: [64, 128, 1024]
+ - num_epochs: [50, 100, 200]
+ - learning_rate: [0.001, 0.0005, 0.0001]
+ - optimizer: AdamW # a valid optimizer from torch.nn
+ - loss: CrossEntropyLoss # a valid loss metric from torch.nn
+
+test:
+ - batch_size: 32
+ - loss: CrossEntropyLoss # a valid loss metric from torch.nn
+
+data:
+ - input_db: /path/to/input_db
+ - db_name: tensors # the name of the database in input_db, the default is `tensors` if unspecified
+ - input_layer: language_model.model.layers.16.post_attention_layernorm # the layer value in input_db to filter by
+
+```
+
+> For `data`, the `input_layer` value must be specified as the probe model is only trained on data extracted from a single layer to avoid dimension mismatches or ambiguous results.
+
+## Training
+Note that the training procedure conducts an iterative (naive) hyperparameter search on all configurations of the `training` hyperparameters: `batch_size`, `num_epochs` and `learning_rate`. It uses $k$-fold cross validation (with a default of $k=5$) to store the lowest validation loss from each configuration and returns the configuration with the minimum validation loss.
+
+```yaml
+training:
+ - batch_size: [64, 128, 1024]
+ - num_epochs: [50, 100, 200]
+ - learning_rate: [0.001, 0.0005, 0.0001]
+ ...
+```
+
+## Testing
+After retrieving the best training configuration, the current script trains and tests a probe model on two versions of the data: a **Main** condition (where the data is untouched) and a **Shuffled** condition (where the labels are shuffled).
+
+This is based on work from [Hewitt and Liang (2019)](https://aclanthology.org/D19-1275/), where they propose the inclusion of "control tasks" to ensure that the probe is selective in its learning. The goal is to achieve high task-specific accuracy and low control task accuracy. The output attributes described in [Results](#results) follow the former experimental design.
+
+## Results
+Finally, the program saves both the probe model and result values in `save_dir`. The current script saves the attributes to a `{save_dir}/probe_data.txt`. The following is an example output:
+```python
+{
+ # The final train configuration with the lowest validation loss
+ "train_config": {
+ "batch_size": 64,
+ "num_epochs": 200,
+ "learning_rate": 0.0001,
+ "optimizer": "AdamW",
+ "loss": "CrossEntropyLoss"
+ },
+
+ # Probe accuracy on the shuffled data
+ "shuffle_accuracy": 0.3295668661594391,
+ "shuffle_loss": 2.1123009968163153,
+ "shuffle_preds": [1, 2, 0, ... ],
+ "shuffle_labels": [0, 0, 0, ... ],
+
+ # Probe accuracy on the original data
+ "test_accuracy": 0.6436911225318909,
+ "test_loss": 1.7786695135290789,
+ "test_preds": [2, 0, 2, ...],
+ "test_labels": [0, 0, 2, ...],
+
+ # The statistical significance of the difference between the shuffled and unshuffled data using a z-test
+ "pvalue": 4.8257723322914694e-116
+}
+```
+Details on the reason behind shuffling and experimental design can be found in the [Testing](#testing) section.
diff --git a/src/probe/main.py b/src/probe/main.py
new file mode 100644
index 0000000000000000000000000000000000000000..ca735c73ee70d26b8637bc3cb2b57337da28721e
--- /dev/null
+++ b/src/probe/main.py
@@ -0,0 +1,520 @@
+"""Probe classes for information analysis in models.
+
+Example command: python -m src.probe.probe -c configs/probe/qwen/clevr-boolean-l13-example.yaml
+"""
+
+import argparse
+import io
+import itertools
+import json
+import logging
+import os
+import random
+import sqlite3
+from typing import Any, Dict, Optional
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import yaml
+from sklearn.model_selection import KFold, train_test_split
+from statsmodels.stats.proportion import proportions_ztest
+from torch.utils.data import DataLoader, Dataset, Subset, TensorDataset
+
+
+class ProbeConfig:
+ """Configuration class for the probe."""
+
+ def __init__(self) -> None:
+ """Initialize the configuration.
+
+ Raises:
+ ValueError: If the configuration file is not found.
+ """
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ '-c', '--config', type=str, help='Path to the probe configuration file'
+ )
+
+ parser.add_argument(
+ '--debug',
+ default=False,
+ action='store_true',
+ help='Flag to print out debug statements',
+ )
+
+ parser.add_argument(
+ '-d',
+ '--device',
+ type=str,
+ default='cuda' if torch.cuda.is_available() else 'cpu',
+ help='The device to send the model and tensors to',
+ )
+
+ args = parser.parse_args()
+
+ assert args.config is not None, 'Config file must be provided.'
+ with open(args.config, 'r') as file:
+ data = yaml.safe_load(file)
+ for key in data.keys():
+ setattr(self, key, data[key])
+
+ # Set debug mode based on config
+ logging.getLogger().setLevel(logging.DEBUG if args.debug else logging.INFO)
+
+ # Load model device
+ if 'cuda' in args.device and not torch.cuda.is_available():
+ raise ValueError('No GPU found on this machine')
+
+ self.device = args.device
+ logging.debug(self.device)
+
+ # Load data mapping
+ assert (
+ hasattr(self, 'data')
+ ), 'The `data` field must be specified in the config, with an input database path.'
+
+ data_mapping = {}
+ for mapping in self.data:
+ data_mapping = {**data_mapping, **mapping}
+
+ # Check if specific layer in specified for the database
+ data_mapping.setdefault('input_layer', None)
+
+ # Set default database name if not specified
+ if 'db_name' not in data_mapping:
+ logging.debug(
+ 'Input database name attribute `db_name` not specified, setting to default `tensors`.')
+ data_mapping.setdefault('db_name', 'tensors')
+ self.data = data_mapping
+
+ # Load model mapping
+ model_mapping = {}
+ if hasattr(self, 'model'):
+ for mapping in self.model:
+ model_mapping = {**model_mapping, **mapping}
+
+ # Set default model config if not provided
+ # input_size and output_size will be set when the data is loaded
+ model_mapping.update({k: v for k, v in {
+ 'activation': 'ReLU',
+ 'hidden_size': 256,
+ 'num_layers': 2,
+ }.items() if k not in model_mapping})
+ logging.debug(model_mapping)
+ self.model = model_mapping
+
+ # Load training mapping
+ train_mapping = {}
+ if hasattr(self, 'training'):
+ for mapping in self.training:
+ train_mapping = {**train_mapping, **mapping}
+
+ logging.debug(train_mapping)
+ # Set default training config if not provided
+ train_mapping.update({k: v for k, v in {
+ 'optimizer': 'AdamW',
+ 'learning_rate': 1e-3,
+ 'loss': 'CrossEntropyLoss',
+ 'num_epochs': 10,
+ 'batch_size': 32
+ }.items() if k not in train_mapping})
+
+ self.training = train_mapping
+
+ # Load test mapping
+ test_mapping = {}
+ if hasattr(self, 'test'):
+ for mapping in self.test:
+ test_mapping = {**test_mapping, **mapping}
+
+ # Set default test config if not provided
+ test_mapping.update({k: v for k, v in {
+ 'optimizer': 'AdamW',
+ 'learning_rate': 1e-3,
+ 'loss': 'CrossEntropyLoss',
+ 'num_epochs': 10,
+ 'batch_size': 32
+ }.items() if k not in test_mapping})
+
+ self.test = test_mapping
+
+
+class Probe(nn.Module):
+ """Probe class for extracting information from models."""
+
+ def __init__(self, config: Dict[str, Any]) -> None:
+ """Intialize the probe with the given configuration.
+
+ Args:
+ config (Dict[str, Any]): Configuration dictionary for the probe.
+ """
+ super(Probe, self).__init__()
+ self.config = config
+
+ # Load input data to parse model input_size and output_size
+ self.data = self.load_data()
+
+ # Intialize the model
+ self.build_model()
+
+ def build_model(self) -> None:
+ """Builds the probe model from scratch."""
+ # Intialize probe model
+ layers = list()
+ layers.append(
+ nn.Linear(self.config.model['input_size'],
+ self.config.model['hidden_size'])
+ )
+ layers.append(getattr(nn, self.config.model['activation'])())
+
+ # Intialize intermediate layers based on config
+ for _ in range(self.config.model['num_layers'] - 2):
+ layers.append(
+ nn.Linear(self.config.model['hidden_size'],
+ self.config.model['hidden_size'])
+ )
+ layers.append(getattr(nn, self.config.model['activation'])())
+
+ # Final layer to output the desired size
+ layers.append(
+ nn.Linear(self.config.model['hidden_size'],
+ self.config.model['output_size'])
+ )
+
+ # Combine all layers to construct the model
+ self.model = nn.Sequential(*layers)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """Forward pass of the probe model.
+
+ Args:
+ x (torch.Tensor): Input tensor.
+
+ Returns:
+ torch.Tensor: Output tensor.
+ """
+ logging.debug('Forward pass with input: %s', x.shape)
+ return self.model(x)
+
+ def load_data(self, shuffle: bool = False) -> TensorDataset:
+ """Load tensors from the database.
+
+ Args:
+ shuffle (bool): Whether to shuffle the data.
+
+ Returns:
+ TensorDataset: A dataset containing the loaded tensors.
+ """
+ logging.debug('Loading tensors from the database...')
+ # Connect to database
+ connection = sqlite3.connect(self.config.data['input_db'])
+ cursor = connection.cursor()
+
+ # Build query and fetch results
+ cursor.execute(
+ f"SELECT layer, tensor, label FROM {self.config.data['db_name']}"
+ )
+ results = cursor.fetchall()
+
+ # Close the connection
+ connection.close()
+
+ # Gather unique class labels
+ all_labels = set([result[2] for result in results])
+ self.config.model.setdefault('output_size', len(all_labels))
+ assert (
+ 'output_size' in self.config.model and len(
+ all_labels) == self.config.model['output_size']
+ ), 'Input attribute `output_size` does not match number of classes in dataset. Leave blank to assign automatically.'
+
+ # Label to index mapping
+ label_to_idx = {label: i for i, label in enumerate(all_labels)}
+
+ features, targets = [], []
+ probe_layer = self.config.data.get('input_layer', None)
+ if not probe_layer:
+ logging.debug(
+ 'No `input_layer` attribute provided for database loading, extracting all tensors...')
+
+ input_size = self.config.data.get('input_size', None)
+ for layer, tensor_bytes, label in results:
+ if (probe_layer and layer == probe_layer) or (not probe_layer):
+ tensor = torch.load(io.BytesIO(tensor_bytes),
+ map_location=self.config.device)
+ if tensor.ndim > 2:
+ # Apply mean pooling if tensor is not already pooled
+ tensor = tensor.mean(dim=1)
+ # Squeeze to shape (hidden_dim)
+ tensor = tensor.squeeze()
+
+ if not input_size:
+ # Set model config input_size once
+ input_size = tensor.shape[0] # pooled tensor
+ self.config.model.setdefault('input_size', input_size)
+ assert (
+ 'input_size' in self.config.model and input_size == self.config.model[
+ 'input_size']
+ ), 'Input attribute `input_size` does not match input tensor dimension. Leave blank to assign automatically.'
+
+ features.append(tensor)
+ targets.append(label_to_idx[label])
+
+ if shuffle:
+ random.shuffle(targets)
+
+ # Stack lists into batched tensors
+ X = torch.stack(features)
+ Y = torch.tensor(targets)
+ logging.debug(f'Features shape {X.shape}, Targets shape {Y.shape}')
+
+ # Move tensors to same device as model
+ X, Y = X.to(self.config.device), Y.to(self.config.device)
+
+ return TensorDataset(X, Y)
+
+ def cross_validate(self, config: dict, data: Dataset, nfolds: Optional[int] = 5) -> float:
+ """Trains the model using the config hyperparameters across k folds.
+
+ Args:
+ config (dict): The configuration dictionary.
+ data (Dataset): The dataset to train on.
+ nfolds (Optional[int]): The number of folds for cross-validation.
+
+ Returns:
+ float: The average validation loss across all folds.
+ """
+ kf = KFold(n_splits=nfolds, shuffle=True, random_state=42)
+ val_losses = []
+ for fold, (train_idx, val_idx) in enumerate(kf.split(range(len(data)))):
+ logging.debug(f'===Starting fold {fold}/{nfolds}===')
+ train_set, val_set = Subset(data, train_idx), Subset(data, val_idx)
+
+ # Reinitialize model after each fold to prevent contamination
+ self.build_model()
+
+ result = self.train(config, train_set, val_set)
+ val_losses.append(result['val_loss'] * len(val_set))
+
+ # Return the mean validation loss across all folds
+ return sum(val_losses) / len(data)
+
+ def train(self, train_config: dict, train_set: Dataset, val_set: Optional[Dataset] = None) -> dict:
+ """Train the probe model.
+
+ Args:
+ train_config (dict): The training configuration.
+ train_set (Dataset): The training dataset.
+ val_set (Dataset, optional): The validation dataset.
+
+ Returns:
+ dict: The training results, including validation loss and accuracy.
+ """
+ logging.debug(
+ f'Training the probe model with config {train_config}...')
+
+ # Set the device
+ device = torch.device(self.config.device)
+ self.model.to(device)
+
+ # Initialize the optimizer
+ optimizer_class = getattr(optim, train_config['optimizer'])
+ optimizer = optimizer_class(
+ self.parameters(), lr=train_config['learning_rate'])
+
+ # Intialize the loss function
+ loss_fn = getattr(nn, train_config['loss'])()
+ train_loader = DataLoader(
+ train_set, batch_size=train_config['batch_size'], shuffle=True)
+
+ for epoch in range(train_config['num_epochs']):
+ # Set the model to training mode
+ self.model.train()
+ total_loss = 0
+ for X, Y in train_loader:
+ optimizer.zero_grad()
+
+ outputs = self.model(X.float())
+ loss = loss_fn(outputs, Y)
+
+ loss.backward()
+ optimizer.step()
+
+ total_loss += loss.item() * X.size(0)
+
+ mean_train_loss = total_loss / len(train_set)
+ logging.debug(
+ f"--Epoch {epoch + 1}/{train_config['num_epochs']}: Train loss: {mean_train_loss:.4f}")
+
+ if val_set:
+ val_loader = DataLoader(
+ val_set, batch_size=train_config['batch_size'])
+ # Set model to eval mode and calculate validation loss
+ self.model.eval()
+ val_loss = 0
+ preds, labels = [], []
+ with torch.no_grad():
+ for X_val, Y_val in val_loader:
+ outputs = self.model(X_val.float())
+ loss = loss_fn(outputs, Y_val)
+ val_loss += loss.item() * X_val.size(0)
+
+ preds.append(outputs)
+ labels.append(Y_val)
+
+ preds = torch.cat(preds, dim=0)
+ labels = torch.cat(labels, dim=0)
+
+ val_loss = val_loss / len(val_set)
+ val_acc = (preds.argmax(dim=1) == labels).float().mean().item()
+ logging.debug(
+ f'Validation accuracy: {val_acc}, Validation mean loss: {val_loss}')
+
+ return {'preds': preds, 'labels': labels, 'val_loss': val_loss, 'val_acc': val_acc}
+
+ # TODO: Return train details here
+ return {}
+
+ def evaluate(self, test_set: Dataset) -> dict:
+ """Evaluate the probe model on the input test set.
+
+ Args:
+ test_set (Dataset): The test dataset.
+
+ Returns:
+ dict: The evaluation results, including loss and accuracy.
+ """
+ self.model.eval()
+
+ device = torch.device(self.config.device)
+ self.model.to(device)
+
+ test_config = self.config.test
+ test_loader = DataLoader(
+ test_set, batch_size=test_config['batch_size'])
+
+ loss_fn = getattr(nn, test_config['loss'])()
+ total_loss = 0.0
+ num_correct, num_samples = 0, 0
+ all_preds, all_labels = [], []
+ with torch.no_grad():
+ for X, Y in test_loader:
+ outputs = self.model(X.float())
+ loss = loss_fn(outputs, Y)
+ total_loss += loss.item() * X.size(0) # to account for incomplete batches
+
+ preds = outputs.argmax(dim=1)
+ num_correct += (preds == Y).sum()
+ num_samples += Y.size(0)
+
+ all_preds.append(preds)
+ all_labels.append(Y)
+
+ mean_loss = float(total_loss / len(test_set))
+ accuracy = float(num_correct / num_samples)
+
+ all_preds = torch.cat(all_preds, dim=0).cpu().numpy()
+ all_labels = torch.cat(all_labels, dim=0).cpu().numpy()
+ logging.debug(
+ f'Test accuracy: {accuracy}, Test mean loss: {mean_loss}')
+ return {'accuracy': accuracy,
+ 'loss': mean_loss,
+ 'labels': all_labels,
+ 'preds': all_preds}
+
+ def save_model(self, metadata: Optional[dict] = None) -> None:
+ """Saves the trained model to a user-specified path.
+
+ Args:
+ metadata (Optional[dict]): Metadata to save alongside the model.
+ """
+ save_dir = self.config.model.get('save_dir') or 'probe_output'
+ os.makedirs(save_dir, exist_ok=True)
+
+ save_path = os.path.join(save_dir, 'probe.pth')
+ try:
+ torch.save(self.model.state_dict(), save_path)
+ logging.debug(f'Model saved to {save_path}')
+ except Exception as e:
+ logging.error(f'Failed to save probe model: {e}')
+
+ if metadata:
+ try:
+ data_path = os.path.join(save_dir, 'probe_data.json')
+ with open(data_path, 'w') as f:
+ f.write(json.dumps(metadata, indent=2))
+ logging.debug(f'Probe metadata saved to {data_path}')
+ except Exception as e:
+ logging.error(f'Failed to save metadata: {e}')
+
+
+def main() -> None:
+ """Main function to run the probe."""
+ config = ProbeConfig()
+ probe = Probe(config)
+
+ # Load data and split into train/val and test
+ data = probe.data
+ indices = list(range(len(data)))
+
+ train_idx, test_idx = train_test_split(
+ indices, test_size=0.2, random_state=42)
+ train_set, test_set = Subset(data, train_idx), Subset(data, test_idx)
+
+ # Load all combinations of hyperparameters
+ train_keys = list(config.training.keys())
+ train_configs = list(itertools.product(
+ *[[config.training[k]] if not isinstance(config.training[k], list) else config.training[k] for k in train_keys]))
+ logging.debug(
+ f'Hyperparamer tuning using {len(train_configs)} config combinations...')
+
+ # Train using k-fold cross validation on all configs and store the lowest validation losses
+ val_losses = []
+ for config in train_configs:
+ val_loss = probe.cross_validate(
+ dict(zip(train_keys, config)), train_set)
+ val_losses.append(val_loss)
+
+ # Finally, train the model on the whole train_set using best config
+ min_idx = val_losses.index(min(val_losses))
+ final_config = dict(zip(train_keys, train_configs[min_idx]))
+ logging.debug(
+ f'Model config results after hyperparameter tuning: {final_config}')
+
+ # Shuffle the data and train the model again to test generalization
+ shffl_data = probe.load_data(shuffle=True)
+ shuffl_train, shuffl_test = Subset(
+ shffl_data, train_idx), Subset(shffl_data, test_idx)
+
+ probe.build_model()
+ probe.train(final_config, shuffl_train)
+ shffl_results = probe.evaluate(shuffl_test)
+
+ # Reinitialize model to finally train with best config
+ probe.build_model()
+ probe.train(final_config, train_set)
+ test_results = probe.evaluate(test_set)
+
+ # Calculate p-value using proportions z-test
+ shffl_correct = (shffl_results['preds'] == shffl_results['labels']).sum()
+ test_correct = (test_results['preds'] == test_results['labels']).sum()
+ pvalue = proportions_ztest([test_correct, shffl_correct],
+ [len(test_results['preds']), len(shffl_results['preds'])])[1]
+
+ # Save results to file with non-shuffled model to file
+ probe.save_model({'train_config': final_config,
+ 'shuffle_accuracy': shffl_results['accuracy'],
+ 'shuffle_loss': shffl_results['loss'],
+ 'shuffle_preds': shffl_results['preds'].tolist(),
+ 'shuffle_labels': shffl_results['labels'].tolist(),
+ 'test_accuracy': test_results['accuracy'],
+ 'test_loss': test_results['loss'],
+ 'test_preds': test_results['preds'].tolist(),
+ 'test_labels': test_results['labels'].tolist(),
+ 'pvalue': pvalue})
+
+ # TODO: implement a demo
+
+
+if __name__ == '__main__':
+ main()
diff --git a/src/utils.py b/src/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..c173deb29795f605f72c2673c9687734dab11875
--- /dev/null
+++ b/src/utils.py
@@ -0,0 +1,43 @@
+"""Utility functions for interacting with the SQLite database."""
+import io
+import logging
+import sqlite3
+from typing import Any, List, Optional
+
+import torch
+
+
+def select_tensors(
+ db_path: str,
+ table_name: str,
+ keys: List[str] = ['layer', 'pooling_method', 'tensor_dim', 'tensor'],
+ sql_where: Optional[str] = None,
+ ) -> List[Any]:
+ """Select and return all tensors from the specified SQLite database and table.
+
+ Args:
+ db_path (str): Path to the SQLite database file.
+ table_name (str): Name of the table to query.
+ keys (List[str]): List of keys to select from the database.
+ sql_where (str): Optional SQL WHERE clause to filter results.
+
+ Returns:
+ List[Any]: A list of tensors retrieved from the database.
+ """
+ if 'tensor' not in keys:
+ logging.warning("'tensor' key should be included to retrieve tensors; automatically adding it.")
+ keys.append('tensor')
+ final_results = []
+ with sqlite3.connect(db_path) as connection:
+ cursor = connection.cursor()
+ query = f'SELECT {", ".join(keys)} FROM {table_name}'
+ if sql_where:
+ assert sql_where.strip().lower().startswith('where'), "sql_where should start with 'WHERE'"
+ query += f' {sql_where}'
+ cursor.execute(query)
+ results = cursor.fetchall()
+ for row in results:
+ result_item = {key: value for key, value in zip(keys, row)}
+ result_item['tensor'] = torch.load(io.BytesIO(result_item['tensor']), map_location='cpu')
+ final_results.append(result_item)
+ return final_results
+
VLM-Lens
+
+[](https://www.python.org/downloads/release/python-31012/)
+[](https://www.apache.org/licenses/LICENSE-2.0)
+[](https://compling-wat.github.io/vlm-lens/)
+[](tutorial-notebooks/guide.ipynb)
+[](https://colab.research.google.com/drive/13WC4HA6syXFotmn7S8WsVz4OmoHsfHV9?usp=sharing)
+
+