fusion-bench 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

fusion_bench/taskpool/clip_vision/taskpool.py

@@ -3,7 +3,17 @@ import json
 import logging
 import os
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast  # noqa: F401
+from typing import (  # noqa: F401
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    Union,
+    cast,
+)
 
 import torch
 from omegaconf import DictConfig
@@ -25,6 +35,10 @@ from fusion_bench.tasks.clip_classification import get_classnames_and_templates
 from fusion_bench.utils import instantiate
 from fusion_bench.utils.parameters import count_parameters
 
+if TYPE_CHECKING:
+    from fusion_bench.models.surgery.surgerymodelwrapper import SurgeryModelWrapper
+
+# disable tokenizers parallelism by default to avoid deadlocks
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 log = logging.getLogger(__name__)
@@ -198,14 +212,16 @@ class CLIPVisionModelTaskPool(
         classifier: HFCLIPClassifier,
         test_loader: DataLoader,
         num_classes: int,
+        task_name: str = None,
     ):
         """
-        Evaluate the classifier on the test dataset.
+        Evaluate the classifier on the test dataset (single-task evaluation).
 
         Args:
             classifier (HFCLIPClassifier): The classifier to evaluate.
             test_loader (DataLoader): The data loader for the test dataset.
             num_classes (int): The number of classes in the classification task.
+            task_name (str): The name of the task.
 
         Returns:
             Dict[str, float]: A dictionary containing the accuracy and loss of the classifier on the test dataset.
@@ -228,7 +244,12 @@ class CLIPVisionModelTaskPool(
             )
         ):
             inputs, targets = batch
-            outputs = classifier(inputs, return_image_embeds=True, return_dict=True)
+            outputs = classifier(
+                inputs,
+                return_image_embeds=True,
+                return_dict=True,
+                task_name=task_name,
+            )
             logits: Tensor = outputs["logits"]
 
             loss = F.cross_entropy(logits, targets)
@@ -246,12 +267,18 @@ class CLIPVisionModelTaskPool(
         results = {"accuracy": acc, "loss": loss}
         return results
 
-    def evaluate(self, model: Union[CLIPVisionModel, CLIPVisionTransformer], name=None):
+    def evaluate(
+        self,
+        model: Union[CLIPVisionModel, CLIPVisionTransformer],
+        name=None,
+        **kwargs,
+    ):
         """
         Evaluate the model on the image classification task.
 
         Args:
             model (Union[CLIPVisionModel, CLIPVisionTransformer]): The model to evaluate.
+            name (Optional[str]): The name of the model. This will be logged into the report if not None.
 
         Returns:
             Dict[str, Any]: A dictionary containing the evaluation results for each task.
@@ -261,8 +288,17 @@ class CLIPVisionModelTaskPool(
 
         report = {}
         # CLIPVisionModel works the same with CLIPVisonTransformer, so we can use it directly
-        self.clip_model.vision_model = model
-        classifier = HFCLIPClassifier(self.clip_model, processor=self.processor)
+        if hasattr(model, "is_surgery_model") and model.is_surgery_model:
+            log.info("running evaluation on a surgery model.")
+            model: "SurgeryModelWrapper" = model
+            self.clip_model.vision_model = model
+        else:
+            # replace the vision encoder with the model
+            self.clip_model.vision_model = model
+        classifier = HFCLIPClassifier(
+            self.clip_model,
+            processor=self.processor,
+        )
         classifier = cast(HFCLIPClassifier, self.fabric.to_device(classifier))
         # collect basic model information
         training_params, all_params = count_parameters(model)
@@ -285,6 +321,7 @@ class CLIPVisionModelTaskPool(
                 classifier,
                 test_dataloader,
                 num_classes=len(classnames),
+                task_name=task_name,
             )
             report[task_name] = result
             self.on_task_evaluation_end()

fusion_bench/taskpool/llama/reward_model.py

@@ -0,0 +1,157 @@
+"""
+The dataset contains the following fields:
+
+- chosen_input_ids: The input token ids for the winner.
+- chosen_attention_mask: The attention mask for the winner.
+- rejected_input_ids: The input token ids for the loser.
+- rejected_attention_mask: The attention mask for the loser.
+"""
+
+import functools
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast
+
+import lightning as L
+import torch
+from omegaconf import DictConfig
+from torch.utils.data import Subset
+import numpy as np
+from tqdm.auto import tqdm
+
+from fusion_bench.dataset.llama.collate import bradley_terry_rm_collate
+from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.taskpool import BaseTaskPool
+from fusion_bench.utils import instantiate
+
+if TYPE_CHECKING:
+    from transformers import LlamaForSequenceClassification
+
+
+def evaluate_batch(model: "LlamaForSequenceClassification", batch):
+    batch_size = batch["input_ids"].size(0)
+    assert batch_size % 2 == 0, "Batch size must be even."
+
+    outputs = model(
+        input_ids=batch["input_ids"],
+        attention_mask=batch["attention_mask"],
+    )
+
+    rewards = outputs[0]
+    chosen_reward = rewards[: batch_size // 2]
+    rejected_rewards = rewards[batch_size // 2 :]
+
+    loss = -torch.log(torch.sigmoid(chosen_reward - rejected_rewards)).mean()
+    correct = (chosen_reward > rejected_rewards).sum().item()
+    total = batch_size // 2
+
+    return {
+        "loss": loss.item(),
+        "correct": correct,
+        "total": total,
+    }
+
+
+def evaluate_dataloader(model: "LlamaForSequenceClassification", dataloader):
+    """
+    Compute the accuracy of the reward model on the given dataloader.
+
+    Args:
+        model: The reward model
+        dataloader: The dataloader for the dataset
+
+    Returns:
+        float: The accuracy of the reward model on the dataset
+    """
+    metrics = {
+        "loss": 0.0,
+        "correct": 0,
+        "total": 0,
+    }
+    with torch.no_grad():
+        for batch in (pbar := tqdm(dataloader)):
+            batch_result = evaluate_batch(model, batch)
+            new_total = metrics["total"] + batch_result["total"]
+            metrics["loss"] = (
+                metrics["loss"] * metrics["total"] / new_total
+                + batch_result["loss"] * batch_result["total"] / new_total
+            )
+            metrics["correct"] += batch_result["correct"]
+            metrics["total"] += batch_result["total"]
+            pbar.set_postfix(metrics)
+
+    metrics["accuracy"] = metrics["correct"] / metrics["total"]
+    return metrics
+
+
+class RewardModelEvaluationTaskPool(
+    BaseTaskPool,
+    LightningFabricMixin,
+):
+    def __init__(
+        self,
+        test_datasets: List[DictConfig],
+        dataloader_kwargs: DictConfig,
+        tokenizer: Optional[DictConfig],
+        max_num_samples: int = -1,
+        seed: int = 0,
+        **kwargs,
+    ):
+        self.seed = seed
+        L.seed_everything(seed)
+        self._test_datasets = test_datasets
+        self.dataloader_kwargs = dataloader_kwargs
+        self._tokenizer = tokenizer
+        self.max_num_samples = max_num_samples
+        super().__init__(**kwargs)
+
+    def setup(self):
+        if self._tokenizer is None:
+            # try to load the tokenizer from the model pool
+            tokenizer = self._program.modelpool.load_tokenizer()
+        else:
+            tokenizer = instantiate(self._tokenizer)
+        self.tokenizer = tokenizer
+
+        test_datasets = {
+            dataset_name: instantiate(self._test_datasets[dataset_name])
+            for dataset_name in self._test_datasets
+        }
+        if self.max_num_samples > 0:
+            test_datasets = {
+                dataset_name: Subset(
+                    test_dataset,
+                    np.random.permutation(len(test_dataset))[: self.max_num_samples],
+                )
+                for dataset_name, test_dataset in test_datasets.items()
+            }
+        test_dataloaders = {
+            dataset_name: torch.utils.data.DataLoader(
+                test_dataset,
+                collate_fn=functools.partial(
+                    bradley_terry_rm_collate,
+                    pad_token_id=tokenizer.pad_token_id,
+                ),
+                **self.dataloader_kwargs,
+            )
+            for dataset_name, test_dataset in test_datasets.items()
+        }
+
+        self.test_dataloaders = {
+            dataset_name: self.fabric.setup_dataloaders(test_dataloader)
+            for dataset_name, test_dataloader in test_dataloaders.items()
+        }
+
+    @torch.no_grad()
+    def evaluate(self, model: "LlamaForSequenceClassification"):
+        self.setup()
+
+        model = self.fabric.setup_module(model)
+        if model.config.pad_token_id is None:
+            model.config.pad_token_id = self.tokenizer.pad_token_id
+
+        model.eval()
+        report = {}
+        for dataset_name, test_dataloader in self.test_dataloaders.items():
+            report[dataset_name] = evaluate_dataloader(model, test_dataloader)
+
+        print(report)
+        return report

fusion_bench/taskpool/nyuv2_taskpool.py

@@ -60,4 +60,6 @@ class NYUv2TaskPool(TaskPool):
             num_workers=self.config.num_workers,
         )
         report = self.trainer.validate(model, val_loader)
+        if isinstance(report, list) and len(report) == 1:
+            report = report[0]
         return report

fusion_bench/tasks/flan_t5_text_generation/glue_load_dataset.py

@@ -3,9 +3,10 @@ import os
 from typing import Optional
 
 from datasets import load_dataset, load_from_disk
+from omegaconf import DictConfig
 
 from fusion_bench.utils import instantiate, timeit_context
-from omegaconf import DictConfig
+
 from .glue_preprocessors import glue_processors
 from .glue_prompt_templates import glue_prompt_templates
 

fusion_bench/utils/hydra_utils.py

@@ -4,3 +4,25 @@ import hydra.core.hydra_config
 def get_hydra_output_dir():
     hydra_cfg = hydra.core.hydra_config.HydraConfig.get()
     return hydra_cfg.runtime.output_dir
+
+
+def config_priority_get(priority_config, general_config, key, default):
+    """
+    Retrieve a configuration value with priority.
+
+    This function retrieves the value associated with `key` from `priority_config` if it exists.
+    If the key is not found in `priority_config`, it retrieves the value from `general_config`.
+    If the key is not found in either configuration, it returns the provided `default` value.
+
+    Args:
+        priority_config (dict): The configuration dictionary with higher priority.
+        general_config (dict): The general configuration dictionary.
+        key (str): The key to look up in the configuration dictionaries.
+        default: The default value to return if the key is not found in either configuration.
+
+    Returns:
+        The value associated with `key` from `priority_config` or `general_config`, or the `default` value if the key is not found.
+    """
+    if key in priority_config:
+        return priority_config[key]
+    return general_config.get(key, default)

fusion_bench/utils/plot/token.py

@@ -0,0 +1,52 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import seaborn as sns
+
+
+def visualize_model_inputs(input_ids, attention_mask, labels, tokenizer=None):
+    """
+    Visualize model inputs: attention mask, labels and input_ids
+
+    Parameters:
+    -----------
+    attention_mask: numpy array or tensor
+        The attention mask array
+    labels: numpy array or tensor
+        The labels array
+    input_ids: numpy array or tensor
+        The input ids array
+    tokenizer: optional
+        The tokenizer object to decode input_ids
+    """
+
+    # Convert inputs to numpy if they're tensors
+    attention_mask = np.array(attention_mask)
+    labels = np.array(labels)
+    input_ids = np.array(input_ids)
+
+    # Create figure with 3 subplots
+    fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(15, 10))
+
+    # Plot attention mask
+    sns.heatmap(attention_mask.reshape(1, -1), ax=ax1, cmap="Blues", cbar=True)
+    ax1.set_title("**Attention Mask**")
+    ax1.set_ylabel("Sequence")
+
+    # Plot labels
+    sns.heatmap(labels.reshape(1, -1), ax=ax2, cmap="Reds", cbar=True)
+    ax2.set_title("**Labels**")
+    ax2.set_ylabel("Sequence")
+
+    # Plot input_ids
+    sns.heatmap(input_ids.reshape(1, -1), ax=ax3, cmap="Greens", cbar=True)
+    ax3.set_title("**Input IDs**")
+    ax3.set_ylabel("Sequence")
+
+    # If tokenizer is provided, add decoded tokens as x-axis labels
+    if tokenizer:
+        decoded_tokens = [tokenizer.decode(token_id) for token_id in input_ids]
+        ax3.set_xticks(np.arange(len(decoded_tokens)) + 0.5)
+        ax3.set_xticklabels(decoded_tokens, rotation=45, ha="right")
+
+    plt.tight_layout()
+    return fig

fusion_bench/utils/plot/token_notebook.py

@@ -0,0 +1,127 @@
+import numpy as np
+from IPython.display import HTML, display
+
+
+def create_color_style():
+    return """
+    <style>
+        .token-container { font-family: monospace; white-space: pre; }
+        .attention { background-color: #90EE90; }  /* Light green */
+        .label { background-color: #FFB6C6; }      /* Light red */
+        .token { color: #0066cc; }                 /* Blue */
+        .stats { font-weight: bold; }
+    </style>
+    """
+
+
+def escape_special_chars(text):
+    """Convert special characters to their string representation"""
+    return (
+        text.replace("\n", "\\n")
+        .replace("\t", "\\t")
+        .replace("\r", "\\r")
+        .replace(" ", "␣")
+    )  # Optional: show spaces with visible character
+
+
+def visualize_tokens_html(input_ids, attention_mask, labels, tokenizer):
+    """
+    Visualize model inputs using HTML colored text representation for Jupyter Notebook
+    with special characters shown as strings
+    """
+    # Convert to numpy if tensors
+    attention_mask = np.array(attention_mask).flatten()
+    labels = np.array(labels).flatten()
+    input_ids = np.array(input_ids).flatten()
+
+    # Decode tokens and escape special characters
+    tokens = [escape_special_chars(tokenizer.decode(id_)) for id_ in input_ids]
+
+    # Create HTML output
+    html_output = [create_color_style()]
+
+    # Header
+    html_output.append("<h3>**Token Visualization**</h3>")
+
+    # Legend
+    html_output.append(
+        """
+    <div style='margin: 10px 0;'>
+        <strong>Legend:</strong><br>
+        <span class='attention'>&nbsp;&nbsp;&nbsp;&nbsp;</span> Active Attention<br>
+        <span class='label'>&nbsp;&nbsp;&nbsp;&nbsp;</span> Label Present<br>
+        <span class='token'>Text</span> Token Text<br>
+        Special Characters: \\n (newline), \\t (tab), ␣ (space)
+    </div>
+    """
+    )
+
+    # Token alignment
+    html_output.append("<strong>Token Alignment:</strong>")
+    html_output.append("<div class='token-container'>")
+
+    # Calculate maximum token length for better alignment
+    max_token_len = max(len(str(token)) for token in tokens)
+
+    for i, (input_id, token, mask, label) in enumerate(
+        zip(input_ids, tokens, attention_mask, labels)
+    ):
+        # Pad token for alignment
+        token_text = f"{token:{max_token_len}s}"
+
+        # Create classes for styling
+        classes = []
+        if mask == 1:
+            classes.append("attention")
+        if label != -100 and label != 0:
+            classes.append("label")
+
+        class_str = f"class='{' '.join(classes)}'" if classes else ""
+
+        # Create the line
+        line = f"Position {i:3d}: <span {class_str}><span class='token'>{token_text}</span></span> "
+        line += (
+            f"(Mask: {int(mask)}, Label: {int(label)}, Inpu_id: {int(input_id)})<br>"
+        )
+        html_output.append(line)
+
+    html_output.append("</div>")
+
+    # Statistics
+    html_output.append(
+        """
+    <div class='stats' style='margin-top: 10px;'>
+        Statistics:<br>
+        Total tokens: {}<br>
+        Active attention tokens: {}<br>
+        Labeled tokens: {}<br>
+    </div>
+    """.format(
+            len(tokens), attention_mask.sum(), sum(labels != -100)
+        )
+    )
+
+    # Display the HTML
+    display(HTML("".join(html_output)))
+
+
+# Example usage:
+"""
+from transformers import AutoTokenizer
+import torch
+
+# Initialize tokenizer
+tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
+
+# Sample input with special characters
+text = "Hello,\nhow are\tyou?"
+inputs = tokenizer(text, return_tensors='pt')
+labels = torch.zeros_like(inputs['input_ids'])  # dummy labels
+
+visualize_tokens_html(
+    inputs['attention_mask'][0],
+    labels[0],
+    inputs['input_ids'][0],
+    tokenizer
+)
+"""

fusion_bench/utils/type.py

@@ -6,18 +6,20 @@ from typing_extensions import TypeAlias
 
 try:
     import torch
-    from torch import Tensor
+    from torch import Tensor, nn
 
     StateDictType: TypeAlias = Dict[str, Tensor]
+    TorchModelType = TypeVar("TorchModelType", bound=nn.Module)
+
 except ImportError:
     pass
 
 
-ModuleType = type(sys)
+PyModuleType = type(sys)
 T = TypeVar("T")
 T1 = TypeVar("T1")
 T2 = TypeVar("T2")
 T3 = TypeVar("T3")
 T4 = TypeVar("T4")
 
-__all__ = ["StateDictType", "ModuleType", "T", "T1", "T2", "T3", "T4"]
+__all__ = ["StateDictType", "PyModuleType", "TorchModelType", "T", "T1", "T2", "T3", "T4"]

{fusion_bench-0.2.5.dist-info → fusion_bench-0.2.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fusion_bench
-Version: 0.2.5
+Version: 0.2.7
 Summary: A Comprehensive Benchmark of Deep Model Fusion
 Author-email: Anke Tang <tang.anke@foxmail.com>
 License: MIT License