Upload AMPLIFY

amplify.py

@@ -295,6 +295,14 @@ class AMPLIFY(AMPLIFYPreTrainedModel):
         # Initialize
         hidden_states, attentions = [], []
 
+        # We will output all the hidden_states that have an index higher than output_hidden_index
+        if type(output_hidden_states) == bool and not output_hidden_states:
+            output_hidden_index = self.config.num_hidden_layers + 1
+        elif type(output_hidden_states) == int:
+            output_hidden_index = output_hidden_states
+        else:
+            output_hidden_index = 0
+
         # Expand and repeat: (Batch, Length) -> (Batch, Heads, Length, Length)
         if pad_mask is not None:
             pad_mask = pad_mask.unsqueeze(1).unsqueeze(1).repeat(1, self.config.num_attention_heads, pad_mask.size(-1), 1)
@@ -325,9 +333,9 @@ class AMPLIFY(AMPLIFYPreTrainedModel):
             x = self.layer_norm_1(x)
 
         # Transformer encoder
-        for layer in self.transformer_encoder:
+        for idx, layer in enumerate(self.transformer_encoder):
             x, attn = layer(x, pad_mask, freqs_cis, output_attentions, max_seqlen, cu_seqlens)
-            if output_hidden_states:
+            if idx >= output_hidden_index:
                 hidden_states.append(x)
             if output_attentions:
                 attentions.append(attn)