fusion-bench 0.2.9__py3-none-any.whl

fusion_bench/taskpool/gpt2_text_classification.py

@@ -0,0 +1,149 @@
+import functools
+import itertools
+import logging
+from copy import deepcopy
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from omegaconf import DictConfig
+from torch.utils.data import DataLoader
+from torchmetrics import Accuracy, MeanMetric
+from tqdm.autonotebook import tqdm
+from transformers import (
+    GPT2ForSequenceClassification,
+    GPT2Model,
+    GPT2Tokenizer,
+    default_data_collator,
+)
+from typing_extensions import override
+
+from fusion_bench.dataset.gpt2_glue import TokenizedGLUE
+from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.taskpool import BaseTaskPool
+from fusion_bench.utils import instantiate
+
+log = logging.getLogger(__name__)
+
+tokenizer: GPT2Tokenizer = None
+
+
+@functools.cache
+def load_gpt2_dataset(name: str, split: Optional[str] = None):
+    global tokenizer
+    dataset = TokenizedGLUE(tokenizer=tokenizer).load_dataset(name)
+    if split is not None:
+        dataset = dataset[split]
+    return dataset
+
+
+class GPT2TextClassificationTaskPool(BaseTaskPool, LightningFabricMixin):
+    """
+    A task pool for GPT2 text classification tasks.
+    This class manages the tasks and provides methods for loading test dataset and evaluation.
+    """
+
+    _config_mapping = BaseTaskPool._config_mapping | {
+        "_test_datasets": "test_datasets",
+        "_tokenizer": "tokenizer",
+        "dataloader_kwargs": "dataloader_kwargs",
+        "fast_dev_run": "fast_dev_run",
+    }
+
+    def __init__(
+        self,
+        test_datasets: DictConfig,
+        tokenizer: DictConfig,
+        dataloader_kwargs: DictConfig,
+        fast_dev_run: bool,
+        **kwargs,
+    ):
+        self._test_datasets = test_datasets
+        self._tokenizer = tokenizer
+        self.dataloader_kwargs = dataloader_kwargs
+        self.fast_dev_run = fast_dev_run
+        super().__init__(**kwargs)
+
+        self.setup()
+
+    def setup(self):
+        global tokenizer
+        self.tokenizer = tokenizer = instantiate(self._tokenizer)
+
+    def get_classifier(
+        self, task_name: str, model: GPT2Model
+    ) -> GPT2ForSequenceClassification:
+        modelpool = self._program.modelpool
+        classifier = modelpool.load_classifier(task_name)
+        classifier.transformer = deepcopy(model)
+        return classifier
+
+    @torch.no_grad()
+    def evaluate_single_task(
+        self,
+        task_name: str,
+        model: GPT2Model,
+        test_loader: DataLoader,
+    ):
+        loss_metric = MeanMetric()
+        # load classifier and replace the backbone with the passed model
+        model: GPT2ForSequenceClassification = self.get_classifier(task_name, model)
+        accuracy = Accuracy("multiclass", num_classes=model.num_labels)
+        model = self.fabric.setup(model)
+
+        if self.config.get("fast_dev_run", False):
+            log.info("Running under fast_dev_run mode, evaluating on a single batch.")
+            test_loader = itertools.islice(test_loader, 1)
+        else:
+            test_loader = test_loader
+
+        for batch in (
+            pbar := tqdm(
+                test_loader, desc="Evaluating", leave=False, dynamic_ncols=True
+            )
+        ):
+            input_ids = batch["input_ids"]
+            attention_mask = batch["attention_mask"]
+            labels = batch["labels"]
+
+            outputs = model(input_ids, attention_mask=attention_mask)
+            logits = outputs.logits
+            loss = F.cross_entropy(logits, labels)
+
+            accuracy(logits.detach().cpu(), labels.detach().cpu())
+            loss_metric.update(loss.detach().cpu())
+            pbar.set_postfix(
+                {
+                    "accuracy": accuracy.compute().item(),
+                    "loss": loss_metric.compute().item(),
+                }
+            )
+
+        acc = accuracy.compute().item()
+        loss = loss_metric.compute().item()
+        results = {"accuracy": acc, "loss": loss}
+        log.info(f"Results for task {task_name}: {results}")
+        return results
+
+    def get_test_dataloader(self, task_name: str):
+        dataset = instantiate(self._test_datasets[task_name])
+        dataloader_kwargs = {
+            "shuffle": False,
+        }
+        dataloader_kwargs.update(self.dataloader_kwargs)
+        dataloader = DataLoader(
+            dataset, collate_fn=default_data_collator, **dataloader_kwargs
+        )
+        if self.fabric is not None:
+            dataloader = self.fabric.setup_dataloaders(dataloader)
+        return dataloader
+
+    @override
+    def evaluate(self, model: GPT2Model):
+        report = {}
+        for task_name in (pbar := tqdm(self._test_datasets, desc="Evaluating tasks")):
+            pbar.set_description(f"Evaluating task {task_name}")
+            dataloader = self.get_test_dataloader(task_name)
+            result = self.evaluate_single_task(task_name, model, dataloader)
+            report[task_name] = result
+        return report

fusion_bench/taskpool/llama/__init__.py

@@ -0,0 +1 @@
+from .test_generation import LlamaTestGenerationTaskPool

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 numpy as np
+import torch
+from omegaconf import DictConfig
+from torch.utils.data import Subset
+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/llama/test_generation.py

@@ -0,0 +1,185 @@
+import itertools
+import logging
+import time
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast
+
+import torch
+
+from fusion_bench import BaseTaskPool
+from fusion_bench.taskpool.dummy import get_model_summary
+from fusion_bench.utils.devices import get_device
+from fusion_bench.utils.rich_utils import print_bordered
+
+if TYPE_CHECKING:
+    from transformers import LlamaForCausalLM, PreTrainedTokenizer
+
+    from fusion_bench.modelpool import CausalLMPool
+log = logging.getLogger(__name__)
+
+
+def generate_text(
+    model: "LlamaForCausalLM",
+    tokenizer: "PreTrainedTokenizer",
+    prompt: str,
+    max_length: int = 1024,
+    temperature: float = 0.01,
+    top_p=0.9,
+    device: torch.device = None,
+):
+    """
+    Generate text using the loaded model.
+
+    Args:
+        model: The loaded language model
+        tokenizer: The loaded tokenizer
+        prompt (str): Input prompt text
+        max_length (int): Maximum length of generated sequence
+        temperature (float): Controls randomness (higher = more random)
+        top_p (float): Nucleus sampling parameter
+
+    Returns:
+        str: Generated text
+    """
+    if device is None:
+        device = get_device(model)
+
+    # Encode the prompt
+    inputs = tokenizer(prompt, return_tensors="pt")
+
+    # Move to GPU if available
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+
+    # Generate
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_length=max_length,
+            temperature=temperature,
+            top_p=top_p,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+            do_sample=True,
+        )
+
+    # Decode and return the generated text
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=False)
+    response = generated_text[len(prompt) :]
+    return {
+        "generated_text": generated_text,
+        "response": response,
+        "num_tokens": len(outputs[0]) - len(inputs["input_ids"][0]),
+    }
+
+
+class LlamaTestGenerationTaskPool(BaseTaskPool):
+    """
+    This task pool is used to evaluate a language model on a set of prompts.
+    For the purpose of debugging, it can also be used in an interactive mode.
+    """
+
+    def __init__(
+        self,
+        test_prompts: List[str],
+        max_length: int = 1024,
+        temperature: float = 0.01,
+        top_p: float = 0.9,
+        iterative_mode: bool = False,
+        **kwargs,
+    ):
+        """
+        Args:
+            test_prompts (List[str]): A list of prompts to be used for testing the model.
+            max_length (int, optional): The maximum length of the generated text. Defaults to 1024.
+            temperature (float, optional): The sampling temperature for text generation. Defaults to 0.01.
+            top_p (float, optional): The cumulative probability for nucleus sampling. Defaults to 0.9.
+            iterative_mode (bool, optional): If True, enables interactive mode for debugging. Defaults to False.
+        """
+        self.test_prompts = test_prompts
+        self.max_length = max_length
+        self.temperature = temperature
+        self.top_p = top_p
+        self.iterative_mode = iterative_mode
+        super().__init__(**kwargs)
+
+    def evaluate(
+        self,
+        model: Union["LlamaForCausalLM", Any],
+        tokenizer: Optional["PreTrainedTokenizer"] = None,
+    ):
+        if tokenizer is None:
+            if self._program is None:
+                log.error(
+                    "`_program` is not set. This is probably happening when you are not runing the program via `fusion_bench` CLI."
+                    "Please pass `tokenizer` to this function."
+                )
+            modelpool: "CausalLMPool" = self._program.modelpool
+            tokenizer = modelpool.load_tokenizer()
+
+        report = get_model_summary(model)
+        if self.test_prompts is not None:
+            for prompt_idx, prompt in enumerate(self.test_prompts):
+                print(f"=== Generating text {prompt_idx+1}/{len(self.test_prompts)}")
+                report[f"conversation_{prompt_idx+1}"] = self._generate_text(
+                    model, tokenizer, prompt
+                )
+
+        if self.iterative_mode:
+            for prompt_idx in itertools.count():
+                # Prompt for input
+                # print usage instructions
+                print("Enter a prompt to generate text. Type 'exit' to exit the loop.")
+                prompt = input(
+                    f"Enter a prompt, or type 'exit' to quit ({prompt_idx+1}): "
+                )
+                if prompt == "exit":
+                    break
+                report[f"iterative_conversation_{prompt_idx+1}"] = self._generate_text(
+                    model, tokenizer, prompt
+                )
+
+        return report
+
+    def _generate_text(
+        self, model: "LlamaForCausalLM", tokenizer: "PreTrainedTokenizer", prompt: str
+    ) -> dict:
+        """
+        Generate text using the provided model and tokenizer for a given prompt.
+
+        This method generates text based on the given prompt using the specified model and tokenizer.
+        It prints the prompt and the generated response, and returns a dictionary containing the prompt,
+        response, wall time, number of characters, and number of tokens.
+
+        Args:
+            model: The language model to be used for text generation.
+            tokenizer: The tokenizer to be used for encoding and decoding text.
+            prompt (str): The input prompt for text generation.
+
+        Returns:
+            dict: A dictionary containing the following keys:
+                - "prompt" (str): The input prompt.
+                - "response" (str): The generated response.
+                - "wall_time" (float): The time taken to generate the response.
+                - "num_chars" (int): The number of characters in the generated response.
+                - "num_tokens" (int): The number of tokens in the generated response.
+        """
+        print(prompt)
+        start_time = time.time()
+        outputs = generate_text(
+            model,
+            tokenizer=tokenizer,
+            prompt=prompt,
+            max_length=self.max_length,
+            temperature=self.temperature,
+            top_p=self.top_p,
+        )
+        print_bordered(
+            outputs["response"], title="Generated Text", code_style="markdown"
+        )
+        print("\n")
+        return {
+            "prompt": prompt,
+            "response": outputs["response"],
+            "wall_time": time.time() - start_time,
+            "num_chars": len(outputs["response"]),
+            "num_tokens": outputs["num_tokens"],
+        }

fusion_bench/taskpool/nyuv2_taskpool.py

@@ -0,0 +1,65 @@
+import logging
+from pathlib import Path
+
+import lightning as L
+from omegaconf import DictConfig
+from torch import nn
+from torch.utils.data import DataLoader
+
+from fusion_bench.compat.taskpool.base_pool import TaskPool
+from fusion_bench.dataset.nyuv2 import NYUv2
+from fusion_bench.models.nyuv2.lightning_module import NYUv2MTLModule
+from fusion_bench.models.nyuv2.resnet_dilated import ResnetDilated
+
+log = logging.getLogger(__name__)
+
+
+class NYUv2TaskPool(TaskPool):
+    _trainer: L.Trainer = None
+
+    def __init__(self, taskpool_config: DictConfig):
+        self.config = taskpool_config
+
+    def load_datasets(self):
+        log.info("Loading NYUv2 dataset")
+        data_path = str(Path(self.config.data_dir) / "nyuv2")
+
+        train_dataset = NYUv2(root=data_path, train=True)
+        val_dataset = NYUv2(root=data_path, train=False)
+        return train_dataset, val_dataset
+
+    @property
+    def trainer(self):
+        if self._trainer is None:
+            self._trainer = L.Trainer(devices=1)
+        return self._trainer
+
+    def get_decoders(self):
+        from fusion_bench.modelpool.nyuv2_modelpool import NYUv2ModelPool
+
+        modelpool: NYUv2ModelPool = self._program.modelpool
+        decoders = nn.ModuleDict()
+        for task in self.config.tasks:
+            decoders[task] = modelpool.load_model(task, encoder_only=False).decoders[
+                task
+            ]
+        return decoders
+
+    def evaluate(self, encoder: ResnetDilated):
+        model = NYUv2MTLModule(
+            encoder,
+            self.get_decoders(),
+            tasks=self.config.tasks,
+            task_weights=[1] * len(self.config.tasks),
+        )
+        _, val_dataset = self.load_datasets()
+        val_loader = DataLoader(
+            val_dataset,
+            batch_size=self.config.batch_size,
+            shuffle=False,
+            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/__init__.py

@@ -0,0 +1,2 @@
+# flake8: noqa F401
+from .base_task import BaseTask

fusion_bench/tasks/base_task.py

@@ -0,0 +1,18 @@
+from abc import ABC, abstractmethod
+
+from omegaconf import DictConfig
+
+
+class BaseTask(ABC):
+    _taskpool = None
+
+    def __init__(self, task_config: DictConfig):
+        self.config = task_config
+
+    @abstractmethod
+    def evaluate(self, model):
+        """
+        Evaluate the model on the task.
+        Returns a dictionary containing the evaluation metrics.
+        """
+        raise NotImplementedError

fusion_bench/tasks/classification.py

@@ -0,0 +1,75 @@
+import functools
+import itertools
+import logging
+from abc import abstractmethod
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+from torchmetrics import Accuracy, MeanMetric
+from torchmetrics.classification.accuracy import MulticlassAccuracy
+from tqdm.autonotebook import tqdm
+
+from .base_task import BaseTask
+
+log = logging.getLogger(__name__)
+
+
+class ClassificationTask(BaseTask):
+    def __init__(self, task_config):
+        super().__init__(task_config)
+
+    @property
+    @abstractmethod
+    def num_classes(self):
+        """
+        Returns the number of classes in the dataset.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def test_loader(self):
+        """
+        Returns a test data loader.
+        """
+        pass
+
+    @torch.no_grad()
+    def evaluate(self, classifier: nn.Module, device=None):
+        accuracy: MulticlassAccuracy = Accuracy(
+            task="multiclass", num_classes=self.num_classes
+        )
+        classifier.eval()
+        loss_metric = MeanMetric()
+        # if fast_dev_run is set, we only evaluate on a batch of the data
+        if self.config.get("fast_dev_run", False):
+            log.info("Running under fast_dev_run mode, evaluating on a single batch.")
+            test_loader = itertools.islice(self.test_loader, 1)
+        else:
+            test_loader = self.test_loader
+
+        for batch in (
+            pbar := tqdm(
+                test_loader, desc="Evaluating", leave=False, dynamic_ncols=True
+            )
+        ):
+            inputs, targets = batch
+            if device is not None:
+                inputs, targets = inputs.to(device), targets.to(device)
+            logits: Tensor = classifier(inputs)
+
+            loss = F.cross_entropy(logits, targets)
+            loss_metric.update(loss.detach().cpu())
+            acc = accuracy(logits.detach().cpu(), targets.detach().cpu())
+            pbar.set_postfix(
+                {
+                    "accuracy": accuracy.compute().item(),
+                    "loss": loss_metric.compute().item(),
+                }
+            )
+
+        acc = accuracy.compute().item()
+        loss = loss_metric.compute().item()
+        results = {"accuracy": acc, "loss": loss}
+        return results