Upload model

BranchyModel.py

@@ -5,13 +5,15 @@ import torch.nn as nn
 import torch.nn.functional as F
 import copy
 
-from dataclasses import dataclass
+from dataclasses import dataclass 
 from torch import Tensor
 from .BranchyModelConfig import BranchyModelConfig
 from typing import List, Optional, Dict, Tuple
 from transformers import AutoModelForCausalLM, PreTrainedModel
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.utils import ModelOutput
+from transformers.cache_utils import Cache, DynamicCache
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -306,7 +308,7 @@ class BranchyCausalModel(PreTrainedModel):
                  config: BranchyModelConfig):
         super().__init__(config)
         self.model = BranchyModel(config)
-        self.head_thresholds = torch.tensor(config.head_thresholds).to(self.model.device)
+        self.head_thresholds = torch.tensor(config.head_thresholds).to(config.device)
         if config.confidence_metric == "breaking_ties":
             self.confidence_metric_fn = breaking_ties
         elif config.confidence_metric == "max":
@@ -314,6 +316,63 @@ class BranchyCausalModel(PreTrainedModel):
         else:
             raise ValueError("confidence_metric must be 'breaking_ties' or 'max'.")
         self.post_init()
+
+    
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values is not None:
+            if isinstance(past_key_values, Cache):
+                cache_length = past_key_values.get_seq_length()
+                past_length = past_key_values.seen_tokens
+                max_cache_length = past_key_values.get_max_length()
+            else:
+                cache_length = past_length = past_key_values[0][0].shape[2]
+                max_cache_length = None
+
+            # Keep only the unprocessed tokens:
+            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
+            # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
+            # input)
+            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
+                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
+            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
+            # input_ids based on the past_length.
+            elif past_length < input_ids.shape[1]:
+                input_ids = input_ids[:, past_length:]
+            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
+
+            # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
+            if (
+                max_cache_length is not None
+                and attention_mask is not None
+                and cache_length + input_ids.shape[1] > max_cache_length
+            ):
+                attention_mask = attention_mask[:, -max_cache_length:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -input_ids.shape[1] :]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
     
     def forward(self,
                 input_ids: torch.LongTensor = None,

README.md

@@ -1,8 +1,8 @@
 ---
-library_name: transformers
-license: mit
 language:
 - en
+license: mit
+library_name: transformers
 pipeline_tag: text-generation
 ---
 

config.json

@@ -15,6 +15,7 @@
   "branch_number": 4,
   "confidence_metric": "breaking_ties",
   "copy_lm_head": false,
+  "device": "cuda:1",
   "head_thresholds": [
     10.0,
     10.0,

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.40.2"
+}