Fix inference code

inference.py

@@ -5,7 +5,7 @@ model = AutoModelForCausalLM.from_pretrained(checkpoint, trust_remote_code=True)
 tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
 
 prompt = 'The mastermind behind the plan was, all along, '
-inputs = tokenizer(prompt, return_tensors='pt')
+inputs = tokenizer(prompt, return_tensors='pt', add_special_tokens=True)
 
 generate_ids = model.generate(inputs.input_ids,
                               max_new_tokens=40,

modeling_gpt1.py

@@ -149,13 +149,6 @@ class GPT1Model(GPT1PreTrainedModel):
             [GPT1DecoderLayer(config) for _ in range(config.num_hidden_layers)]
         )
 
-        causal_mask = torch.full((1, config.max_position_embeddings, config.max_position_embeddings),
-                                 fill_value=float('-inf'))
-        self.register_buffer('causal_mask',
-                             torch.triu(causal_mask, diagonal=1),
-                             persistent=False)
-        self.mask_cache_len = config.max_position_embeddings
-
         self.post_init()
 
     def get_input_embeddings(self):
@@ -164,7 +157,7 @@ class GPT1Model(GPT1PreTrainedModel):
     def set_input_embeddings(self, value):
         self.embs = value
 
-    def forward(self, input_ids, attention_mask=None, *args, **kwargs):
+    def forward(self, input_ids, *args, **kwargs):
         position_ids = torch.arange(input_ids.size(-1),
                                     dtype=torch.long,
                                     device=input_ids.device).unsqueeze_(0)
@@ -173,18 +166,12 @@ class GPT1Model(GPT1PreTrainedModel):
         position_embeds = self.pos_emb(position_ids)
         hidden_state = self.embs_dropout(input_embeds) + position_embeds
 
-        if attention_mask is not None and attention_mask.size(1) > self.mask_cache_len:
-            seq_len = attention_mask.size(1)
-            self.mask_cache_len = seq_len
-
-            causal_mask = torch.full((seq_len, seq_len),
-                                     fill_value=float('-inf'))
-            self.register_buffer('causal_mask',
-                                 torch.triu(causal_mask, diagonal=1),
-                                 persistent=False)
+        seq_len = input_ids.size(-1)
+        attn_mask = torch.full((seq_len, seq_len), fill_value=float('-inf'))
+        attn_mask = torch.triu(attn_mask, diagonal=1)
 
-        causal_mask = self.causal_mask.to(dtype=input_embeds.dtype,
-                                          device=input_embeds.device)
+        causal_mask = attn_mask.to(dtype=input_embeds.dtype,
+                                   device=input_embeds.device)
 
         for layer in self.layers:
             hidden_state = layer(hidden_state, attn_mask=causal_mask)
@@ -225,9 +212,8 @@ class GPT1ForCausalLM(GPT1PreTrainedModel):
     def set_decoder(self, decoder):
         self.model = decoder
 
-    def forward(self, input_ids, labels=None, attention_mask=None,
-                *args, **kwargs):
-        output = self.model(input_ids, attention_mask)
+    def forward(self, input_ids, labels=None, *args, **kwargs):
+        output = self.model(input_ids)
 
         hidden_state = output[0]
         logits = self.lm_head(hidden_state).float()
@@ -247,16 +233,5 @@ class GPT1ForCausalLM(GPT1PreTrainedModel):
             logits=logits
         )
 
-    def prepare_inputs_for_generation(self, input_ids, attention_mask,
-                                      *args, **kwargs):
-        assert attention_mask.size(1) == input_ids.size(1)
-
-        seq_len = attention_mask.size(1)
-
-        attn_mask = torch.full((seq_len, seq_len), fill_value=float('-inf'))
-        attn_mask = torch.triu(attn_mask, diagonal=1)
-
-        return {
-            'input_ids': input_ids,
-            'attention_mask': attn_mask
-        }
+    def prepare_inputs_for_generation(self, input_ids, *args, **kwargs):
+        return { 'input_ids': input_ids }