Fix couple of issues with inference, dataset folder

inference.py

@@ -8,7 +8,7 @@ prompt = 'The mastermind behind the plan was, all along, '
 inputs = tokenizer(prompt, return_tensors='pt')
 
 generate_ids = model.generate(inputs.input_ids,
-                              max_length=50,
+                              max_new_tokens=40,
                               num_beams=1,
                               do_sample=True,
                               top_p=0.9,

modeling_gpt1.py

@@ -163,7 +163,7 @@ class GPT1Model(GPT1PreTrainedModel):
     def set_input_embeddings(self, value):
         self.embs = value
 
-    def forward(self, input_ids, *args, **kwargs):
+    def forward(self, input_ids, attention_mask=None, *args, **kwargs):
         position_ids = torch.arange(input_ids.size(-1),
                                     dtype=torch.long,
                                     device=input_ids.device).unsqueeze_(0)
@@ -172,8 +172,12 @@ class GPT1Model(GPT1PreTrainedModel):
         position_embeds = self.pos_emb(position_ids)
         hidden_state = self.embs_dropout(input_embeds) + position_embeds
 
-        causal_mask = self.causal_mask.to(dtype=input_embeds.dtype,
-                                          device=input_embeds.device)
+        if attention_mask is not None:
+            causal_mask = attention_mask.to(dtype=input_embeds.dtype,
+                                            device=input_embeds.device)
+        else:
+            causal_mask = self.causal_mask.to(dtype=input_embeds.dtype,
+                                              device=input_embeds.device)
 
         for layer in self.layers:
             hidden_state = layer(hidden_state, attn_mask=causal_mask)
@@ -214,8 +218,9 @@ class GPT1ForCausalLM(GPT1PreTrainedModel):
     def set_decoder(self, decoder):
         self.model = decoder
 
-    def forward(self, input_ids, labels = None, *args, **kwargs):
-        output = self.model(input_ids)
+    def forward(self, input_ids, labels=None, attention_mask=None,
+                *args, **kwargs):
+        output = self.model(input_ids, attention_mask)
 
         hidden_state = output[0]
         logits = self.lm_head(hidden_state).float()
@@ -236,4 +241,12 @@ class GPT1ForCausalLM(GPT1PreTrainedModel):
         )
 
     def prepare_inputs_for_generation(self, input_ids, *args, **kwargs):
-        return { 'input_ids': input_ids }
+        seq_len = input_ids.size(1)
+
+        attn_mask = torch.full((1, 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
+        }

pre_training.py

@@ -11,17 +11,13 @@ from datasets import load_from_disk
 from configuration_gpt1 import GPT1Config
 from modeling_gpt1 import GPT1Model, GPT1ForCausalLM
 
-# a few more things to try to get the model to train (in this order)
-# actually manually check the input (the books), and the tokenizer output (i
-# don't know if it tokenizes correctly, if it adds eos_token etc.)
-
 
 GPT1Config.register_for_auto_class()
 GPT1Model.register_for_auto_class('AutoModel')
 GPT1ForCausalLM.register_for_auto_class('AutoModelForCausalLM')
 
-# load the already tokenized dataset (see training_preprocessing.py)
-tokenized_datasets = load_from_disk('tokenized_bookcorpusopen')
+# load the already tokenized dataset (see preprocessing.py)
+tokenized_datasets = load_from_disk('data')
 
 # shuffle for good measure
 tokenized_datasets = tokenized_datasets.shuffle(seed=42)

preprocessing.py

@@ -26,4 +26,4 @@ tokenized_datasets = raw_datasets.map(
     remove_columns=raw_datasets['train'].column_names,
 )
 
-tokenized_datasets.save_to_disk('tokenized_bookcorpusopen')
+tokenized_datasets.save_to_disk('data')