added necesary files-1

configs/MetaLlama3.json

@@ -0,0 +1,31 @@
+{
+    "model_name_or_path": "meta-llama/Meta-Llama-3-8B-Instruct",
+    "pooling_mode": "weighted_mean",
+    "dataset_name": "E5",
+    "dataset_file_path": "cache/echo-data",
+    "remove_unused_columns": false,
+    "learning_rate": 0.0002,
+    "num_train_epochs": 3,
+    "warmup_steps": 300,
+    "per_device_train_batch_size": 64,
+    "per_device_eval_batch_size": 64,
+    "gradient_accumulation_steps": 1,
+    "do_train": true,
+    "disable_tqdm": false,
+    "max_seq_length": 512,
+    "overwrite_output_dir": true,
+    "output_dir": "output/meta-llama/Meta-Llama-3-8B-Instruct",
+    "use_adapter": true,
+    "percent_prune": [25],
+    "autoprune": "small+large",
+    "logging_steps": 50,
+    "save_steps": 200,
+    "save_total_limit": 1,
+    "save_only_model": true,
+    "stop_after_n_steps": 1000,
+    "lora_r": 16,
+    "gradient_checkpointing": true,
+    "torch_dtype": "bfloat16",
+    "attn_implementation": "flash_attention_2",
+    "seed": 42
+}

configs/Mistral.json

@@ -0,0 +1,31 @@
+{
+    "model_name_or_path": "mistralai/Mistral-7B-Instruct-v0.2",
+    "pooling_mode": "weighted_mean",
+    "dataset_name": "E5",
+    "dataset_file_path": "cache/echo-data",
+    "remove_unused_columns": false,
+    "learning_rate": 0.0002,
+    "num_train_epochs": 3,
+    "warmup_steps": 300,
+    "per_device_train_batch_size": 64,
+    "per_device_eval_batch_size": 64,
+    "gradient_accumulation_steps": 1,
+    "do_train": true,
+    "disable_tqdm": false,
+    "max_seq_length": 512,
+    "overwrite_output_dir": true,
+    "output_dir": "output/mistralai/Mistral-7B-Instruct-v0.2",
+    "use_adapter": true,
+    "percent_prune": [22],
+    "autoprune": "small+large",
+    "logging_steps": 50,
+    "save_steps": 200,
+    "save_total_limit": 1,
+    "save_only_model": true,
+    "stop_after_n_steps": 1000,
+    "lora_r": 16,
+    "gradient_checkpointing": true,
+    "torch_dtype": "bfloat16",
+    "attn_implementation": "flash_attention_2",
+    "seed": 42
+}

configs/Phi.json

@@ -0,0 +1,31 @@
+{
+    "model_name_or_path": "microsoft/Phi-3-mini-4k-instruct",
+    "pooling_mode": "weighted_mean",
+    "dataset_name": "E5",
+    "dataset_file_path": "cache/echo-data",
+    "remove_unused_columns": false,
+    "learning_rate": 0.0002,
+    "num_train_epochs": 3,
+    "warmup_steps": 300,
+    "per_device_train_batch_size": 64,
+    "per_device_eval_batch_size": 64,
+    "gradient_accumulation_steps": 1,
+    "do_train": true,
+    "disable_tqdm": false,
+    "max_seq_length": 512,
+    "overwrite_output_dir": true,
+    "output_dir": "output/microsoft/Phi-3-mini-4k-instruct",
+    "use_adapter": true,
+    "percent_prune": [25],
+    "autoprune": "small+large",
+    "logging_steps": 50,
+    "save_steps": 200,
+    "save_total_limit": 1,
+    "save_only_model": true,
+    "stop_after_n_steps": 1000,
+    "lora_r": 16,
+    "gradient_checkpointing": true,
+    "torch_dtype": "bfloat16",
+    "attn_implementation": "flash_attention_2",
+    "seed": 42
+}

configs/Qwen.json

@@ -0,0 +1,31 @@
+{
+    "model_name_or_path": "Qwen/Qwen2-7B-Instruct",
+    "pooling_mode": "weighted_mean",
+    "dataset_name": "E5",
+    "dataset_file_path": "cache/echo-data",
+    "remove_unused_columns": false,
+    "learning_rate": 0.0002,
+    "num_train_epochs": 3,
+    "warmup_steps": 300,
+    "per_device_train_batch_size": 64,
+    "per_device_eval_batch_size": 64,
+    "gradient_accumulation_steps": 1,
+    "do_train": true,
+    "disable_tqdm": false,
+    "max_seq_length": 512,
+    "overwrite_output_dir": true,
+    "output_dir": "output/Qwen/Qwen2-7B-Instruct",
+    "use_adapter": true,
+    "percent_prune": [25],
+    "autoprune": "small+large",
+    "logging_steps": 50,
+    "save_steps": 200,
+    "save_total_limit": 1,
+    "save_only_model": true,
+    "stop_after_n_steps": 200,
+    "lora_r": 16,
+    "gradient_checkpointing": true,
+    "torch_dtype": "bfloat16",
+    "attn_implementation": "flash_attention_2",
+    "seed": 42
+}

l3prune/__init__.py

@@ -0,0 +1,2 @@
+from .llmencoder import LLMEncoder
+from .l3prune import l3prune

l3prune/dataset/E5Data.py

@@ -0,0 +1,155 @@
+import json
+import random
+import os
+
+from .dataset import DataSample, TrainSample, Dataset
+from accelerate.logging import get_logger
+
+logger = get_logger(__name__, log_level="INFO")
+
+E5_EMBEDDING_PROMPTS = {
+    "allnli": [
+        "Given a premise, retrieve a hypothesis that is entailed by the premise",
+        "Retrieve semantically similar text",
+    ],
+    "dureader": "Given a Chinese search query, retrieve web passages that answer the question",
+    "eli5_question_answer": "Provided a user question, retrieve the highest voted answers on Reddit ELI5 forum",
+    "fever": "Given a claim, retrieve documents that support or refute the claim",
+    "hotpot_qa": "Given a multi-hop question, retrieve documents that can help answer the question",
+    "miracl": "Given a question, retrieve Wikipedia passages that answer the question",
+    "mrtydi": "Given a question, retrieve Wikipedia passages that answer the question",
+    "msmarco_passage": "Given a web search query, retrieve relevant passages that answer the query",
+    "msmarco_document": "Given a web search query, retrieve relevant documents that answer the query",
+    "nq": "Given a question, retrieve Wikipedia passages that answer the question",
+    "quora_duplicates": [
+        "Given a question, retrieve questions that are semantically equivalent to the given question",
+        "Find questions that have the same meaning as the input question",
+    ],
+    "squad": "Retrieve Wikipedia passages that answer the question",
+    "t2ranking": "Given a Chinese search query, retrieve web passages that answer the question",
+    "trivia_qa": "Retrieve Wikipedia passages that answer the question",
+}
+
+
+class E5Data(Dataset):
+    def __init__(
+        self,
+        dataset_name: str = "E5",
+        split: str = "validation",
+        file_path: str = "cache/echo-data",
+        effective_batch_size: int = 32,
+        shuffle_individual_datasets: bool = True,
+        separator: str = "!@#$%^&*()",
+    ):
+        self.dataset_name = dataset_name
+        self.split = split
+        self.effective_batch_size = effective_batch_size
+        self.shuffle_individual_datasets = shuffle_individual_datasets
+        self.separator = separator
+
+        self.data = []
+        self.load_data(file_path)
+
+    def __len__(self):
+        return len(self.data)
+
+    def load_data(self, file_path: str = None):
+        logger.info(f"Loading E5 data from {file_path}...")
+        # file path is actually a directory
+
+        data_map = {}
+        all_samples = []
+        id_ = 0
+        for dataset in E5_EMBEDDING_PROMPTS:
+            logger.info(f"Loading dataset {dataset}...")
+            if dataset not in data_map:
+                data_map[dataset] = []
+            with open(os.path.join(file_path, f"{dataset}.jsonl"), "r") as f:
+                dataset_samples = f.readlines()
+
+            dataset_samples = [json.loads(d) for d in dataset_samples]
+
+            for i, sample in enumerate(dataset_samples):
+                instruction = (
+                    E5_EMBEDDING_PROMPTS[dataset]
+                    if isinstance(E5_EMBEDDING_PROMPTS[dataset], str)
+                    else E5_EMBEDDING_PROMPTS[dataset][i % 2]
+                )
+                query = f"{instruction}; " + self.separator + sample["query"]
+                if dataset in [
+                    "allnli_split2",
+                    "quora_duplicates_split1",
+                    "quora_duplicates_split2",
+                ]:
+                    pos = (
+                        f"{E5_EMBEDDING_PROMPTS[dataset]}; "
+                        + self.separator
+                        + sample["positive"]
+                    )
+                    neg = (
+                        f"{E5_EMBEDDING_PROMPTS[dataset]}; "
+                        + self.separator
+                        + sample["negative"]
+                    )
+                else:
+                    pos = self.separator + sample["positive"]
+                    neg = self.separator + sample["negative"]
+
+                data_map[dataset].append(id_)
+
+                all_samples.append(
+                    DataSample(
+                        id_=id_,
+                        query=query,
+                        positive=pos,
+                        negative=neg,
+                        task_name=dataset,
+                    )
+                )
+                id_ += 1
+
+        # combine split1 and split2
+        new_data_map = {}
+        for dataset in data_map:
+            new_dataset = dataset.replace("_split1", "").replace("_split2", "")
+            if new_dataset not in new_data_map:
+                new_data_map[new_dataset] = []
+            new_data_map[new_dataset] += data_map[dataset]
+        data_map = new_data_map
+
+        if self.shuffle_individual_datasets:
+            for task, samples in data_map.items():
+                random.shuffle(samples)
+
+        datasets = list(data_map.keys())
+
+        logger.info(
+            f"Batching Echo data properly for effective batch size of {self.effective_batch_size}..."
+        )
+        all_batches = []
+        for dataset in datasets:
+            dataset_samples = data_map[dataset]
+            for i in range(0, len(dataset_samples), self.effective_batch_size):
+                batch = dataset_samples[i : i + self.effective_batch_size]
+                if len(batch) == self.effective_batch_size:
+                    all_batches.append(batch)
+                else:
+                    logger.info(f"Skip 1 batch for dataset {dataset}.")
+        random.shuffle(all_batches)
+
+        final_idx_order = []
+        for batch in all_batches:
+            for idx in batch:
+                final_idx_order.append(idx)
+
+        self.data = [all_samples[idx] for idx in final_idx_order]
+        logger.info(f"Loaded {len(self.data)} samples.")
+
+    def __getitem__(self, index):
+        sample = self.data[index]
+        if self.split == "train":
+            return TrainSample(
+                texts=[sample.query, sample.positive, sample.negative], label=1.0
+            )
+        elif self.split == "validation":
+            assert False, "E5Data does not have a validation split."

l3prune/dataset/Wiki1M.py

@@ -0,0 +1,43 @@
+from .dataset import DataSample, TrainSample, Dataset
+from accelerate.logging import get_logger
+
+logger = get_logger(__name__, log_level="INFO")
+
+
+class Wiki1M(Dataset):
+    def __init__(
+        self,
+        dataset_name: str = "Wiki1M",
+        split: str = "validation",
+        file_path: str = "cache/wiki1m_for_simcse.txt",
+    ):
+        self.dataset_name = dataset_name
+        self.split = split
+        self.data = []
+        self.load_data(file_path)
+
+    def __len__(self):
+        return len(self.data)
+
+    def load_data(self, file_path: str = None):
+        logger.info(f"Loading Wiki1M data from {file_path}...")
+        id_ = 0
+        with open(file_path, "r") as f:
+            for line in f:
+                line = line.strip()
+                self.data.append(
+                    DataSample(
+                        id_=id_,
+                        query=line,
+                        positive=line,
+                    )
+                )
+                id_ += 1
+        logger.info(f"Loaded {len(self.data)} samples.")
+
+    def __getitem__(self, index):
+        sample = self.data[index]
+        if self.split == "train":
+            return TrainSample(texts=[sample.query, sample.positive], label=1.0)
+        elif self.split == "validation":
+            assert False, "Wiki1M does not have a validation split."

l3prune/dataset/__init__.py

@@ -0,0 +1,2 @@
+from .E5Data import E5Data
+from .Wiki1M import Wiki1M

l3prune/dataset/dataset.py

@@ -0,0 +1,53 @@
+from dataclasses import dataclass
+from typing import Union, List
+
+import torch
+
+
+@dataclass
+class DataSample:
+    id_: int
+    query: str
+    positive: str
+    negative: str = None
+    task_name: str = None
+
+
+class TrainSample:
+    """
+    Structure for one input example with texts, the label and a unique id
+    """
+
+    def __init__(
+        self, guid: str = "", texts: List[str] = None, label: Union[int, float] = 0
+    ):
+        """
+        Creates one TrainSample with the given texts, guid and label
+
+
+        :param guid
+            id for the example
+        :param texts
+            the texts for the example.
+        :param label
+            the label for the example
+        """
+        self.guid = guid
+        self.texts = texts
+        self.label = label
+
+    def __str__(self):
+        return "<TrainSample> label: {}, texts: {}".format(
+            str(self.label), "; ".join(self.texts)
+        )
+
+
+class Dataset(torch.utils.data.Dataset):
+    def load_data(self, file_path: str = None):
+        raise NotImplementedError()
+
+    def __getitem__(self, index):
+        raise NotImplementedError()
+
+    def __len__(self):
+        raise NotImplementedError()

l3prune/dataset/utils.py

@@ -0,0 +1,25 @@
+from ..dataset import E5Data
+from ..dataset import Wiki1M
+
+
+def load_dataset(dataset_name, split="validation", file_path=None, **kwargs):
+    """
+    Loads a dataset by name.
+
+    Args:
+        dataset_name (str): Name of the dataset to load.
+        split (str): Split of the dataset to load.
+        file_path (str): Path to the dataset file.
+    """
+    dataset_mapping = {
+        "E5": E5Data,
+        "Wiki1M": Wiki1M,
+    }
+
+    if dataset_name not in dataset_mapping:
+        raise NotImplementedError(f"Dataset name {dataset_name} not supported.")
+
+    if split not in ["train", "validation", "test"]:
+        raise NotImplementedError(f"Split {split} not supported.")
+
+    return dataset_mapping[dataset_name](split=split, file_path=file_path, **kwargs)

l3prune/l3prune.py

@@ -0,0 +1,62 @@
+import random
+import torch
+import numpy as np
+from tqdm import tqdm
+from functools import partial
+from .model_overrides import get_forward
+
+# A custom encode function to override the forward of the model
+def encode_custom(forward, encoder, sentence_feature):
+    embed_mask = None
+    if "embed_mask" in sentence_feature:
+        embed_mask = sentence_feature.pop("embed_mask")
+    out, reps = forward(encoder.model, **sentence_feature)
+    sentence_feature["embed_mask"] = embed_mask
+
+    return [encoder.get_pooling(sentence_feature, emb) for emb in reps]
+
+def l3prune(encoder, dataset, loss_fn, batch_size=64, num_batches=100):
+    dataset = [t for t in dataset]
+    subset = random.sample(dataset, batch_size*num_batches)
+    subset = [[encoder.prepare_for_tokenization(t) for t in s.texts] for s in subset]
+    subset = [subset[i:i + batch_size] for i in range(0, len(subset), batch_size)]
+
+    num_layers = encoder.model.config.num_hidden_layers
+    loss = {i: [] for i in range(1, num_layers+1)}
+    forward = get_forward(encoder.model)
+
+    with torch.no_grad():
+        # Override the forward of the model to get the intermediate representations in only one pass
+        if forward:
+            encode = partial(encode_custom, forward)
+            for batch in tqdm(subset):
+                features = []
+                for j in range(3):
+                    embs = [t[j] for t in batch]
+                    embs = encoder.tokenize(embs).to(encoder.model.device)
+                    embs = encode(encoder, embs)
+                    features += [embs]
+                q, d, d_neg = features
+                for i in range(num_layers):
+                    loss[i+1] += [loss_fn(q[i], d[i], d_neg[i])]
+        else:
+            # Without the override, we have to rerun the forward pass with each layer pruned
+            for l in range(num_layers, 0, -1):
+                encoder.prune(layer_prune=l)
+                for batch in tqdm(subset):
+                    features = []
+                    for j in range(3):
+                        embs = [t[j] for t in batch]
+                        embs = encoder.tokenize(embs).to(encoder.model.device)
+                        embs = encoder.forward(embs)
+                        features += [embs]
+                    q, d, d_neg = features
+                    loss[l] += [loss_fn(q, d, d_neg)]
+
+        loss = [torch.tensor(loss[i]).mean().float().detach() for i in range(1, num_layers+1)]
+    
+    # minima before and after midpoint
+    midpoint = num_layers // 2
+    small_p = np.argmin(loss[:midpoint]) + 1
+    large_p = np.argmin(loss[midpoint:]) + midpoint + 1
+    return small_p, large_p

l3prune/llmencoder.py

@@ -0,0 +1,470 @@
+import json
+import logging
+import os
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import torch
+import torch.multiprocessing as mp
+from peft import PeftModel
+from torch import Tensor, device, nn
+from tqdm.autonotebook import tqdm, trange
+from transformers import (
+    AutoModel,
+    AutoConfig,
+    PretrainedConfig,
+    AutoTokenizer,
+    LlamaConfig,
+    MistralConfig,
+    GemmaConfig,
+    Qwen2Config,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def batch_to_device(batch, target_device: device):
+    """
+    send a pytorch batch to a device (CPU/GPU)
+    """
+    for key in batch:
+        if isinstance(batch[key], Tensor):
+            batch[key] = batch[key].to(target_device)
+    return batch
+
+
+class LLMEncoder(nn.Module):
+    def __init__(
+        self,
+        model: AutoModel,
+        tokenizer: AutoTokenizer,
+        pooling_mode: str = "weighted_mean",
+        max_length: int = 512,
+        doc_max_length: int = 400,
+        skip_instruction: bool = True,
+    ):
+        super().__init__()
+        self.model = model
+        self.tokenizer = tokenizer
+        self.pooling_mode = pooling_mode
+        self.skip_instruction = skip_instruction
+        self.max_length = max_length
+        self.doc_max_length = doc_max_length
+        self.config = model.config
+
+    @classmethod
+    def from_pretrained(
+        self,
+        base_model_name_or_path,
+        peft_model_name_or_path=None,
+        cache_dir=None,
+        **kwargs,
+    ):
+        """
+        Load a pretrained model from a model identifier or path.
+        Args:
+            base_model_name_or_path: Model identifier or path to pretrained model.
+            peft_model_name_or_path: Path to any PEFT models to apply.
+        Returns: L3Prune model.
+        """
+
+        # pop out encoder args
+        keys = ["pooling_mode", "max_length", "doc_max_length", "skip_instruction"]
+        encoder_args = {
+            key: kwargs.pop(key, None) for key in keys if kwargs.get(key) is not None
+        }
+
+        tokenizer = AutoTokenizer.from_pretrained(base_model_name_or_path, cache_dir=cache_dir)
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.padding_side = "left"
+
+        config = AutoConfig.from_pretrained(base_model_name_or_path)
+        model = AutoModel.from_pretrained(base_model_name_or_path, cache_dir=cache_dir, **kwargs)
+
+        if os.path.isdir(base_model_name_or_path) and os.path.exists(
+            f"{base_model_name_or_path}/config.json"
+        ):
+            with open(f"{base_model_name_or_path}/config.json", "r") as fIn:
+                config_dict = json.load(fIn)
+            config = PretrainedConfig.from_dict(config_dict)
+            model.config._name_or_path = config._name_or_path
+
+        if peft_model_name_or_path is not None:
+            model = PeftModel.from_pretrained(
+                model,
+                peft_model_name_or_path,
+            )
+            model = model.merge_and_unload()
+
+        config = {}
+        if os.path.exists(f"{base_model_name_or_path}/l3prune_config.json"):
+            with open(f"{base_model_name_or_path}/l3prune_config.json", "r") as fIn:
+                l3prune_config = json.load(fIn)
+            config.update(l3prune_config)
+
+        for key, value in encoder_args.items():
+            config[key] = value
+
+        return self(model=model, tokenizer=tokenizer, **config)
+
+    def prune(self, percent_prune=0):
+        """
+        Prune a model to a percentage of layers of the base model. If percent_prune is equal to or greater than 1, 
+        it is taken as the specific layer number to prune to. For example, if percent_prune=0.3, 30% of the layers will be pruned. If
+        percent_prune=3, the model will be pruned to 3 layers.
+        """
+        # take it as the specific layer number to prune to
+        if percent_prune >= 1:
+            new_num_layers = int(percent_prune)
+        else:
+            new_num_layers = int(self.model.config.num_hidden_layers * (1 - percent_prune))
+        print(f"Pruning to {new_num_layers} layer.")
+        self.model.layers = self.model.layers[:new_num_layers]
+        self.model.config.num_hidden_layers = new_num_layers
+    
+    def prepare_for_tokenization(self, text):
+        if self.model.config._name_or_path == "meta-llama/Meta-Llama-3-8B-Instruct":
+            text = (
+                "<|start_header_id|>user<|end_header_id|>\n\n"
+                + text.strip()
+                + "<|eot_id|>"
+            )
+            return text
+        if self.model.config._name_or_path in [
+            "mistralai/Mistral-7B-Instruct-v0.2",
+            "meta-llama/Llama-2-7b-chat-hf",
+        ]:
+            text = "[INST] " + text.strip() + " [/INST]"
+        if self.model.config._name_or_path in [
+            "google/gemma-2-9b-it",
+        ]:
+            text = "<bos><start_of_turn>user\n" + text.strip() + "<end_of_turn>"
+        if self.model.config._name_or_path in [
+            "Qwen/Qwen2-1.5B-Instruct",
+            "Qwen/Qwen2-7B-Instruct",
+        ]:
+            text = "<|im_start|>user\n" + text.strip() + "<|im_end|>"
+        if self.pooling_mode == "eos_token":
+            if self.model.config._name_or_path == "meta-llama/Meta-Llama-3-8B":
+                text = text.strip() + "<|end_of_text|>"
+            elif isinstance(self.model.config, LlamaConfig) or isinstance(
+                self.model.config, MistralConfig
+            ):
+                text = text.strip() + " </s>"
+            elif isinstance(self.model.config, GemmaConfig):
+                text = text.strip() + "<eos>"
+            elif isinstance(self.model.config, Qwen2Config):
+                text = text.strip() + "<|endoftext|>"
+        return text
+
+    def tokenize(self, texts):
+        texts_2 = []
+        original_texts = []
+        for text in texts:
+            t = text.split("!@#$%^&*()")
+            texts_2.append(t[1] if len(t) > 1 else "")
+            original_texts.append("".join(t))
+
+        original = self.tokenizer(
+            original_texts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=self.max_length,
+        )
+        embed_mask = None
+        for t_i, t in enumerate(texts_2):
+            ids = self.tokenizer(
+                [t],
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=self.max_length,
+                add_special_tokens=False,
+            )
+            if embed_mask is None:
+                e_m = torch.zeros_like(original["attention_mask"][t_i])
+                if len(ids["input_ids"][0]) > 0:
+                    e_m[-len(ids["input_ids"][0]) :] = torch.ones(
+                        len(ids["input_ids"][0])
+                    )
+                embed_mask = e_m.unsqueeze(0)
+            else:
+                e_m = torch.zeros_like(original["attention_mask"][t_i])
+                if len(ids["input_ids"][0]) > 0:
+                    e_m[-len(ids["input_ids"][0]) :] = torch.ones(
+                        len(ids["input_ids"][0])
+                    )
+                embed_mask = torch.cat((embed_mask, e_m.unsqueeze(0)), dim=0)
+
+        original["embed_mask"] = embed_mask
+        return original
+
+    def _skip_instruction(self, sentence_feature):
+        assert (
+            sentence_feature["attention_mask"].shape
+            == sentence_feature["embed_mask"].shape
+        )
+        sentence_feature["attention_mask"] = sentence_feature["embed_mask"]
+
+    def forward(self, sentence_feature: Dict[str, Tensor]):
+        embed_mask = None
+        if "embed_mask" in sentence_feature:
+            embed_mask = sentence_feature.pop("embed_mask")
+        reps = self.model(**sentence_feature)
+        sentence_feature["embed_mask"] = embed_mask
+
+        return self.get_pooling(sentence_feature, reps.last_hidden_state)
+
+    def get_pooling(self, features, last_hidden_states):  # All models padded from left
+        assert (
+            self.tokenizer.padding_side == "left"
+        ), "Pooling modes are implemented for padding from left."
+        if self.skip_instruction:
+            self._skip_instruction(features)
+        seq_lengths = features["attention_mask"].sum(dim=-1)
+        if self.pooling_mode == "mean":
+            return torch.stack(
+                [
+                    last_hidden_states[i, -length:, :].mean(dim=0)
+                    for i, length in enumerate(seq_lengths)
+                ],
+                dim=0,
+            )
+        elif self.pooling_mode == "weighted_mean":
+            bs, l, _ = last_hidden_states.shape
+            complete_weights = torch.zeros(bs, l, device=last_hidden_states.device)
+            for i, seq_l in enumerate(seq_lengths):
+                if seq_l > 0:
+                    complete_weights[i, -seq_l:] = torch.arange(seq_l) + 1
+                    complete_weights[i] /= torch.clamp(
+                        complete_weights[i].sum(), min=1e-9
+                    )
+            return torch.sum(last_hidden_states * complete_weights.unsqueeze(-1), dim=1)
+        elif self.pooling_mode == "eos_token" or self.pooling_mode == "last_token":
+            return last_hidden_states[:, -1]
+        elif self.pooling_mode == "bos_token":
+            return last_hidden_states[
+                features["input_ids"] == self.tokenizer.bos_token_id
+            ]
+        else:
+            raise ValueError(f"{self.pooling_mode} is not implemented yet.")
+
+    def _convert_to_str(self, instruction, text):
+        tokenized_q = self.tokenizer(
+            text,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=self.max_length,
+            add_special_tokens=False,
+        )
+        tokenized_q_length = len(tokenized_q["input_ids"][0])
+
+        while tokenized_q_length > self.doc_max_length:
+            reduction_ratio = self.doc_max_length / tokenized_q_length
+            reduced_length = int(len(text.split()) * reduction_ratio)
+            text = " ".join(text.split()[:reduced_length])
+            tokenized_q = self.tokenizer(
+                text,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=self.max_length,
+                add_special_tokens=False,
+            )
+            tokenized_q_length = len(tokenized_q["input_ids"][0])
+
+        return (
+            f"{instruction.strip()} !@#$%^&*(){text}"
+            if instruction
+            else f"!@#$%^&*(){text}"
+        )
+
+    def encode(
+        self,
+        sentences: Union[str, List[str]],
+        batch_size: int = 32,
+        show_progress_bar: bool = True,
+        convert_to_numpy: bool = False,
+        convert_to_tensor: bool = False,
+        device: Optional[str] = None,
+    ):
+        """
+        Encode a list of sentences to their respective embeddings. The sentences can be a list of strings or a string.
+        Args:
+            sentences: sentence or sentences to encode.
+            batch_size: batch size for turning sentence tokens into embeddings.
+            show_progress_bar: whether to show progress bars during encoding steps.
+            convert_to_numpy: If true, return numpy arrays instead of torch tensors.
+            convert_to_tensor: If true, return torch tensors (default).
+            device: torch backend device identifier (e.g., 'cuda', 'cpu','mps' etc.). If not specified,
+            the default is to use cuda when available, otherwise cpu. Note that only the choice of 'cuda' supports
+            multiprocessing as currently implemented.
+
+        Returns: embeddings of the sentences. Embeddings are detached and always on the CPU (see _encode implementation).
+
+        """
+        if isinstance(sentences[0], str) and isinstance(sentences[-1], int):
+            sentences = [sentences]
+        # required for MEDI version of MTEB
+        if isinstance(sentences[0], str):
+            sentences = [[""] + [sentence] for sentence in sentences]
+
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        concatenated_input_texts = []
+        for sentence in sentences:
+            assert isinstance(sentence[0], str)
+            assert isinstance(sentence[1], str)
+            concatenated_input_texts.append(
+                self._convert_to_str(sentence[0], sentence[1])
+            )
+        sentences = concatenated_input_texts
+
+        self.eval()
+
+        if convert_to_tensor:
+            convert_to_numpy = False
+
+        length_sorted_idx = np.argsort([-self._text_length(sen) for sen in sentences])
+        sentences_sorted = [sentences[idx] for idx in length_sorted_idx]
+        all_embeddings = []
+
+        if torch.cuda.device_count() <= 1:
+            # This branch also support mps devices
+            self.to(device)
+            for start_index in trange(
+                0,
+                len(sentences),
+                batch_size,
+                desc="Batches",
+                disable=not show_progress_bar,
+            ):
+                sentences_batch = sentences_sorted[
+                    start_index : start_index + batch_size
+                ]
+                embeddings = self._encode(
+                    sentences_batch, device=device, convert_to_numpy=convert_to_numpy
+                )
+                all_embeddings.append(embeddings)
+        else:
+
+            num_proc = torch.cuda.device_count()
+            cuda_compatible_multiprocess = mp.get_context("spawn")
+            with cuda_compatible_multiprocess.Pool(num_proc) as p:
+                sentences_batches = [
+                    sentences_sorted[start_index : start_index + batch_size]
+                    for start_index in range(0, len(sentences), batch_size)
+                ]
+
+                progress_bar = tqdm(
+                    total=len(sentences_batches),
+                    desc="Batches",
+                    disable=not show_progress_bar,
+                )
+                results = []
+
+                def update(*args):
+                    progress_bar.update()
+
+                for batch in sentences_batches:
+                    results.append(
+                        p.apply_async(
+                            self._encode,
+                            args=(batch, None, convert_to_numpy, True),
+                            callback=update,
+                        )
+                    )
+
+                all_embeddings = [result.get() for result in results]
+                progress_bar.close()
+
+        all_embeddings = torch.cat(all_embeddings, dim=0)
+        all_embeddings = all_embeddings[np.argsort(length_sorted_idx)]
+        all_embeddings = all_embeddings.to(torch.float32)
+        if convert_to_numpy:
+            all_embeddings = np.asarray([emb.numpy() for emb in all_embeddings])
+        return all_embeddings
+
+    def save(self, output_path, merge_before_save=False, save_config=True):
+        if merge_before_save and isinstance(self.model, PeftModel):
+            self.model = self.model.merge_and_unload()
+            if hasattr(self.model, "_hf_peft_config_loaded"):
+                self.model._hf_peft_config_loaded = False
+
+        self.model.save_pretrained(output_path)
+        self.tokenizer.save_pretrained(output_path)
+
+        l3prune_config = {
+            "pooling_mode": self.pooling_mode,
+            "max_length": self.max_length,
+            "doc_max_length": self.doc_max_length,
+            "skip_instruction": self.skip_instruction,
+        }
+
+        if save_config:
+            os.makedirs(output_path, exist_ok=True)
+            with open(f"{output_path}/l3prune_config.json", "w") as fOut:
+                json.dump(l3prune_config, fOut, indent=4)
+
+    def _encode(
+        self,
+        sentences_batch,
+        device: Optional[str] = None,
+        convert_to_numpy: bool = False,
+        multiprocessing=False,
+    ):
+        if multiprocessing:
+            # multiprocessing only supports CUDA devices at this time, so we ignore the value of device
+            # and use cuda:rank for the device
+            rank = mp.current_process()._identity[0]
+            if device is None and torch.cuda.is_available():
+                device = f"cuda:{rank % torch.cuda.device_count()}"
+
+        self.to(device)
+        features = self.tokenize(
+            [self.prepare_for_tokenization(sentence) for sentence in sentences_batch]
+        )
+        features = batch_to_device(features, device)
+
+        with torch.no_grad():
+            embeddings = self.forward(features)
+            embeddings = embeddings.detach()
+            embeddings = embeddings.cpu()
+
+        return embeddings
+
+    def _text_length(self, text: Union[List[int], List[List[int]]]):
+        """
+        Help function to get the length for the input text. Text can be either a string (which means a single text)
+        a list of ints (which means a single tokenized text), or a tuple of list of ints
+        (representing several text inputs to the model).
+        """
+        if (
+            isinstance(text, str)
+            or (isinstance(text, list) and isinstance(text[0], int))
+            or len(text) == 0
+        ):  # Single text, list of ints, or empty
+            return len(text)
+        if isinstance(text, dict):  # {key: value} case
+            return len(next(iter(text.values())))
+        elif not hasattr(text, "__len__"):  # Object has no len() method
+            return 1
+        else:
+            return sum([len(t) for t in text])
+
+    def resize_token_embeddings(
+        self,
+        new_num_tokens: Optional[int] = None,
+        pad_to_multiple_of: Optional[int] = None,
+    ) -> nn.Embedding:
+        return self.model.resize_token_embeddings(
+            new_num_tokens=new_num_tokens, pad_to_multiple_of=pad_to_multiple_of
+        )
+
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        self.model.gradient_checkpointing_enable(
+            gradient_checkpointing_kwargs=gradient_checkpointing_kwargs
+        )

l3prune/loss/HardNegativeNLLLoss.py

@@ -0,0 +1,46 @@
+import torch
+from torch import nn, Tensor
+from .loss_utils import cos_sim, mismatched_sizes_all_gather
+
+
+class HardNegativeNLLLoss:
+    def __init__(
+        self,
+        scale: float = 20.0,
+        similarity_fct=cos_sim,
+    ):
+        self.scale = scale
+        self.similarity_fct = similarity_fct
+        self.cross_entropy_loss = nn.CrossEntropyLoss()
+
+    def __call__(
+        self,
+        q_reps: Tensor,
+        d_reps_pos: Tensor,
+        d_reps_neg: Tensor = None,
+    ):
+        if d_reps_neg is None:
+            d_reps_neg = d_reps_pos[:0, :]
+
+        if torch.distributed.is_initialized():
+            full_d_reps_pos = mismatched_sizes_all_gather(d_reps_pos)
+            full_d_reps_pos = torch.cat(full_d_reps_pos)
+
+            full_q_reps = mismatched_sizes_all_gather(q_reps)
+            full_q_reps = torch.cat(full_q_reps)
+
+            full_d_reps_neg = mismatched_sizes_all_gather(d_reps_neg)
+            full_d_reps_neg = torch.cat(full_d_reps_neg)
+        else:
+            full_d_reps_pos = d_reps_pos
+            full_q_reps = q_reps
+            full_d_reps_neg = d_reps_neg
+
+        d_reps = torch.cat([full_d_reps_pos, full_d_reps_neg], dim=0)
+        scores = self.similarity_fct(full_q_reps, d_reps) * self.scale
+        labels = torch.tensor(
+            range(len(scores)), dtype=torch.long, device=scores.device
+        )
+
+        loss = self.cross_entropy_loss(scores, labels)
+        return loss

l3prune/loss/__init__.py

@@ -0,0 +1 @@
+from .HardNegativeNLLLoss import HardNegativeNLLLoss

l3prune/loss/loss_utils.py

@@ -0,0 +1,107 @@
+import torch
+from torch import Tensor
+
+
+class AllGather(torch.autograd.Function):
+    """
+    all_gather with gradient back-propagation
+    """
+
+    @staticmethod
+    def forward(ctx, tensor_list, tensor, group, async_op):
+        torch.distributed.all_gather(
+            tensor_list, tensor, group=group, async_op=async_op
+        )
+        return tuple(tensor_list)
+
+    @staticmethod
+    def backward(ctx, *grad_list):
+        grad_list = list(grad_list)
+        rank = torch.distributed.get_rank()
+
+        dist_ops = [
+            torch.distributed.reduce(grad_list[i], i, async_op=True)
+            for i in range(torch.distributed.get_world_size())
+        ]
+
+        for op in dist_ops:
+            op.wait()
+
+        return None, grad_list[rank], None, None
+
+
+all_gather_with_grad = AllGather.apply
+
+
+def cos_sim(a: Tensor, b: Tensor):
+    """
+    Computes the cosine similarity cos_sim(a[i], b[j]) for all i and j.
+    :return: Matrix with res[i][j]  = cos_sim(a[i], b[j])
+    """
+    if not isinstance(a, torch.Tensor):
+        a = torch.tensor(a)
+
+    if not isinstance(b, torch.Tensor):
+        b = torch.tensor(b)
+
+    if len(a.shape) == 1:
+        a = a.unsqueeze(0)
+
+    if len(b.shape) == 1:
+        b = b.unsqueeze(0)
+
+    a_norm = torch.nn.functional.normalize(a, p=2, dim=1)
+    b_norm = torch.nn.functional.normalize(b, p=2, dim=1)
+    return torch.mm(a_norm, b_norm.transpose(0, 1))
+
+
+def mismatched_sizes_all_gather(
+    tensor: Tensor, group=None, async_op=False, mismatched_axis=0
+):
+    # all_gather doesn't support tensor lists where the first dimension is mismatched. This does.
+    assert torch.distributed.is_initialized(), "torch.distributed not initialized"
+    world_size = torch.distributed.get_world_size()
+    # let's get the sizes for everyone
+    mismatched_sizes = torch.tensor(
+        [tensor.shape[mismatched_axis]], dtype=torch.int64, device="cuda"
+    )
+    sizes = [torch.zeros_like(mismatched_sizes) for _ in range(world_size)]
+    torch.distributed.all_gather(
+        sizes, mismatched_sizes, group=group, async_op=async_op
+    )
+    sizes = torch.cat(sizes).cpu().tolist()
+    # now pad to the max dim-0 size
+    max_size = max(sizes)
+    padded = torch.zeros(
+        (
+            *tensor.shape[:mismatched_axis],
+            max_size,
+            *tensor.shape[mismatched_axis + 1 :],
+        ),
+        device=tensor.device,
+        dtype=tensor.dtype,
+    )
+    # selects the place where we're adding information
+    padded_to_fill = padded.narrow(mismatched_axis, 0, tensor.shape[mismatched_axis])
+    padded_to_fill[...] = tensor
+    # gather the padded tensors
+    tensor_list = [
+        torch.zeros(padded.shape, device=padded.device, dtype=padded.dtype)
+        for _ in range(world_size)
+    ]
+    all_gather_with_grad(tensor_list, padded, group, async_op)
+    # trim off the padding
+    for rank in range(world_size):
+        # checks that the rest is 0
+        assert (
+            not tensor_list[rank]
+            .narrow(
+                mismatched_axis,
+                sizes[rank],
+                padded.shape[mismatched_axis] - sizes[rank],
+            )
+            .count_nonzero()
+            .is_nonzero()
+        ), "This would remove non-padding information"
+        tensor_list[rank] = tensor_list[rank].narrow(mismatched_axis, 0, sizes[rank])
+    return tensor_list

l3prune/loss/utils.py

@@ -0,0 +1,9 @@
+from .HardNegativeNLLLoss import HardNegativeNLLLoss
+
+
+def load_loss(loss_class, *args, **kwargs):
+    if loss_class == "HardNegativeNLLLoss":
+        loss_cls = HardNegativeNLLLoss
+    else:
+        raise ValueError(f"Unknown loss class {loss_class}")
+    return loss_cls(*args, **kwargs)

l3prune/model_overrides.py

@@ -0,0 +1,458 @@
+import torch
+from typing import List, Optional, Tuple, Union
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers import LlamaConfig, Qwen2Config, Phi3Config, MistralConfig
+
+def qwen2_forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
+            )
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                use_cache = False
+
+        use_legacy_cache = False
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = self._update_causal_mask(
+            attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+        )
+
+        hidden_states = inputs_embeds
+
+        # decoder layers
+        layerwise_reps = ()
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = None
+
+        for decoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    causal_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    use_cache,
+                    cache_position,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=causal_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                    cache_position=cache_position,
+                )
+
+            hidden_states = layer_outputs[0]
+            layerwise_reps += (hidden_states,)
+
+            if use_cache:
+                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+        layerwise_reps = [self.norm(rep) for rep in layerwise_reps]
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = None
+        if use_cache:
+            next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return (BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        ), layerwise_reps)
+
+def phi3_forward(
+    self,
+    input_ids: torch.LongTensor = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[List[torch.FloatTensor]] = None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+    use_cache: Optional[bool] = None,
+    output_attentions: Optional[bool] = None,
+    output_hidden_states: Optional[bool] = None,
+    return_dict: Optional[bool] = None,
+    cache_position: Optional[torch.LongTensor] = None,
+) -> Union[Tuple, BaseModelOutputWithPast]:
+    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+    output_hidden_states = (
+        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+    )
+    use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+    if (input_ids is None) ^ (inputs_embeds is not None):
+        raise ValueError(
+            "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
+        )
+
+    if self.gradient_checkpointing and self.training:
+        if use_cache:
+            use_cache = False
+
+    use_legacy_cache = False
+
+    if inputs_embeds is None:
+        inputs_embeds = self.embed_tokens(input_ids)
+
+    if cache_position is None:
+        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+        cache_position = torch.arange(
+            past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+        )
+    if position_ids is None:
+        position_ids = cache_position.unsqueeze(0)
+
+    causal_mask = self._update_causal_mask(
+        attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+    )
+
+    inputs_embeds = self.embed_dropout(inputs_embeds)
+    hidden_states = inputs_embeds
+
+    # decoder layers
+    layerwise_reps = ()
+    all_hidden_states = () if output_hidden_states else None
+    all_self_attns = () if output_attentions else None
+    next_decoder_cache = None
+
+    for decoder_layer in self.layers:
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        if self.gradient_checkpointing and self.training:
+            layer_outputs = self._gradient_checkpointing_func(
+                decoder_layer.__call__,
+                hidden_states,
+                causal_mask,
+                position_ids,
+                output_attentions,
+                use_cache,
+                past_key_values,
+                cache_position,
+            )
+        else:
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+            )
+
+        hidden_states = layer_outputs[0]
+        layerwise_reps += (hidden_states,)
+
+        if use_cache:
+            next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+        if output_attentions:
+            all_self_attns += (layer_outputs[1],)
+
+    hidden_states = self.norm(hidden_states)
+    layerwise_reps = [self.norm(rep) for rep in layerwise_reps]
+
+    # add hidden states from the last decoder layer
+    if output_hidden_states:
+        all_hidden_states += (hidden_states,)
+
+    next_cache = None
+    if use_cache:
+        next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
+    if not return_dict:
+        return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+    return (BaseModelOutputWithPast(
+        last_hidden_state=hidden_states,
+        past_key_values=next_cache,
+        hidden_states=all_hidden_states,
+        attentions=all_self_attns,
+    ), layerwise_reps)
+
+def mistral_forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
+            )
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            use_cache = False
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        return_legacy_cache = False
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = self._update_causal_mask(
+            attention_mask, inputs_embeds, cache_position, past_key_values, use_cache, output_attentions
+        )
+
+        hidden_states = inputs_embeds
+
+        # decoder layers
+        layerwise_reps = ()
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = None
+
+        for decoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    causal_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    use_cache,
+                    cache_position,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=causal_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                    cache_position=cache_position,
+                )
+
+            hidden_states = layer_outputs[0]
+            layerwise_reps += (hidden_states,)
+
+            if use_cache:
+                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+        layerwise_reps = [self.norm(rep) for rep in layerwise_reps]
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = next_decoder_cache if use_cache else None
+        if return_legacy_cache:
+            next_cache = next_cache.to_legacy_cache()
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return (BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        ), layerwise_reps)
+
+def llama3_forward(
+    self,
+    input_ids: torch.LongTensor = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values = None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+    use_cache: Optional[bool] = None,
+    output_attentions: Optional[bool] = None,
+    output_hidden_states: Optional[bool] = None,
+    return_dict: Optional[bool] = None,
+    cache_position: Optional[torch.LongTensor] = None,
+):
+    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+    output_hidden_states = (
+        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+    )
+    use_cache = use_cache if use_cache is not None else self.config.use_cache
+    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+    if (input_ids is None) ^ (inputs_embeds is not None):
+        raise ValueError(
+            "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
+        )
+
+    if self.gradient_checkpointing and self.training and use_cache:
+        use_cache = False
+
+    if inputs_embeds is None:
+        inputs_embeds = self.embed_tokens(input_ids)
+
+    return_legacy_cache = False
+
+    if cache_position is None:
+        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+        cache_position = torch.arange(
+            past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+        )
+    if position_ids is None:
+        position_ids = cache_position.unsqueeze(0)
+
+    causal_mask = self._update_causal_mask(
+        attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+    )
+    hidden_states = inputs_embeds
+
+    # create position embeddings to be shared across the decoder layers
+    position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+    # decoder layers
+    layerwise_reps = ()
+    all_hidden_states = () if output_hidden_states else None
+    all_self_attns = () if output_attentions else None
+    next_decoder_cache = None
+
+    for decoder_layer in self.layers:
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        if self.gradient_checkpointing and self.training:
+            layer_outputs = self._gradient_checkpointing_func(
+                decoder_layer.__call__,
+                hidden_states,
+                causal_mask,
+                position_ids,
+                past_key_values,
+                output_attentions,
+                use_cache,
+                cache_position,
+                position_embeddings,
+            )
+        else:
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+            )
+
+        hidden_states = layer_outputs[0]
+        layerwise_reps += (hidden_states,)
+        
+        if use_cache:
+            next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+        if output_attentions:
+            all_self_attns += (layer_outputs[1],)
+
+    hidden_states = self.norm(hidden_states)
+    layerwise_reps = [self.norm(rep) for rep in layerwise_reps]
+
+    # add hidden states from the last decoder layer
+    if output_hidden_states:
+        all_hidden_states += (hidden_states,)
+
+    next_cache = next_decoder_cache if use_cache else None
+    if return_legacy_cache:
+        next_cache = next_cache.to_legacy_cache()
+
+    if not return_dict:
+        return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+    return (BaseModelOutputWithPast(
+        last_hidden_state=hidden_states,
+        past_key_values=next_cache,
+        hidden_states=all_hidden_states,
+        attentions=all_self_attns,
+    ), layerwise_reps)
+
+def get_forward(model):
+    if isinstance(model.config, LlamaConfig):
+        return llama3_forward
+    elif isinstance(model.config, Qwen2Config):
+        return qwen2_forward
+    elif isinstance(model.config, Phi3Config):
+        return phi3_forward
+    elif isinstance(model.config, MistralConfig):
+        return mistral_forward
+    else:
+        return None