Minor bug fixes and code cleanup

Author: tingofurro

model_generator.py

@@ -1,4 +1,4 @@
-from transformers.modeling_gpt2 import GPT2LMHeadModel, GPT2Config
+from transformers import GPT2LMHeadModel, GPT2Config
 
 import torch.utils.data.dataset
 import utils_tokenizer
@@ -16,8 +16,8 @@ def __init__(self, max_output_length=25, max_input_length=300, device='cpu', tok
         elif tokenizer_type == "bpecap":
             self.tokenizer = utils_tokenizer.BPETokenizer(bpe_model)
             config = GPT2Config.from_dict({"finetuning_task": None, "initializer_range": 0.02,
-                            "layer_norm_epsilon": 1e-05, "n_ctx": 1024, "n_embd": 768, "n_head": 12, "n_layer": 12, "n_positions": 1024, "num_labels": 1,
-                            "resid_pdrop": 0.1, "use_bfloat16": False, "vocab_size": self.tokenizer.vocab_size})
+                                           "layer_norm_epsilon": 1e-05, "n_ctx": 1024, "n_embd": 768, "n_head": 12, "n_layer": 12, "n_positions": 1024, "num_labels": 1,
+                                           "resid_pdrop": 0.1, "use_bfloat16": False, "vocab_size": self.tokenizer.vocab_size})
         else:
             print("Tokenizer unrecognized. Should be gpt2 or bpecap.")
             exit()
@@ -36,7 +36,7 @@ def __init__(self, max_output_length=25, max_input_length=300, device='cpu', tok
         self.mode = "train"
 
     def reload(self, from_file):
-        print(self.model.load_state_dict(torch.load(from_file)))
+        print(self.model.load_state_dict(torch.load(from_file), strict=False))
 
     def save(self, to_file):
         torch.save(self.model.state_dict(), to_file)
@@ -132,22 +132,22 @@ def decode_batch(self, bodies, special_append=None, max_output_length=100, sampl
         elif return_scores:
             return outputs, scores.tolist()
         else:
-            return outputs, end_indices
+            return outputs
 
     def decode_beam_batch(self, bodies, beam_size=3, max_output_length=100, sample=False):
         if self.mode != 'eval':
             print("BEWARE. Model is not in eval mode.")
-        self.eval() ## << Surely you are not training with beam decode?
+        self.eval() # << Surely you are not training with beam decode?
 
         batch_size = len(bodies)
         N = batch_size * beam_size
         inputs = self.preprocess_input(bodies)
         next_words = torch.LongTensor([self.tokenizer.start_id] * N).to(self.device).unsqueeze(1)
         build_up = None
         scores = torch.zeros((N)).to(self.device)
-        
+
         one_every_k = torch.FloatTensor([1] + [0] * (beam_size-1)).repeat(batch_size*beam_size).to(self.device)
-        
+
         # Sometimes, we process the same input, as we run it once as a sampled, and once as an argmax, in which case we should reuse the computation
         _, input_past = self.model(input_ids=inputs, past_key_values=None)
         input_past = [torch.repeat_interleave(p, repeats=beam_size, dim=1) for p in input_past]
@@ -157,23 +157,23 @@ def decode_beam_batch(self, bodies, beam_size=3, max_output_length=100, sample=F
             logits, past = self.model(input_ids=next_words, past_key_values=past)
             probs = torch.nn.functional.softmax(logits, dim=2).squeeze(1)
             logprobs = torch.nn.functional.log_softmax(logits, dim=2)
-            
+
             if sample:
                 all_selects = torch.multinomial(probs, beam_size).unsqueeze(1)
             else:
                 _, all_selects = torch.topk(logprobs, k=beam_size, dim=2)
-        
+
             if build_up is not None:
                 not_finished = (1-torch.any(build_up==self.tokenizer.end_id, dim=1).float()).to(self.device)
             else:
-                not_finished = torch.ones_like(scores, dtype=torch.float, device=self.device)        
-            
+                not_finished = torch.ones_like(scores, dtype=torch.float, device=self.device)
+
             expanded_not_finished = torch.repeat_interleave(not_finished, repeats=beam_size)
-            
+
             expanded_score = torch.repeat_interleave(scores, repeats=beam_size) # This should be batch_size * beam_size²
             added_score = logprobs[torch.repeat_interleave(torch.arange(N), repeats=beam_size), 0, all_selects.view(-1)]
             expanded_score += (expanded_not_finished*added_score)
-            
+
             # We don't want you to select from finished beams
             expanded_score -= (1-expanded_not_finished)*(1-one_every_k)*1000.0
 
@@ -182,11 +182,11 @@ def decode_beam_batch(self, bodies, beam_size=3, max_output_length=100, sample=F
             if build_up is None:
                 choices = torch.arange(beam_size, device=self.device).repeat(batch_size)
                 batched_choices = choices.view(batch_size, beam_size)
-                
+
             else:
                 _, batched_choices = torch.topk(batched_scores, k=beam_size, dim=1) # Going from k² choices per element to k choices.
-        
-            batched_tracks = batched_choices / beam_size
+
+            batched_tracks = (batched_choices / beam_size).long()
             tracks = beam_size*torch.repeat_interleave(torch.arange(batch_size), repeats=beam_size).to(self.device) + batched_tracks.view(-1)
 
             selected_scores = batched_scores[torch.repeat_interleave(torch.arange(batch_size), repeats=beam_size), batched_choices.view(-1)]
@@ -200,7 +200,7 @@ def decode_beam_batch(self, bodies, beam_size=3, max_output_length=100, sample=F
             if build_up is not None:
                 build_up = build_up[tracks, :]
             past = [p[:, tracks, :] for p in past]
-            
+
             # Update the latest scores, and the current_build
             if build_up is None:
                 build_up = next_words
@@ -228,7 +228,7 @@ def decode(self, bodies, max_output_length=100, max_batch_size=8, beam_size=1, r
         if progress:
             iterator = tqdm.tqdm(iterator)
         for i in iterator:
-            batch_bodies = bodies[i:min(N,i+max_batch_size)]
+            batch_bodies = bodies[i:min(N, i+max_batch_size)]
             with torch.no_grad():
                 if beam_size > 1:
                     batch_outputs = self.decode_beam_batch(batch_bodies, beam_size=beam_size, max_output_length=max_output_length, sample=sample)

train_summary_loop.py

@@ -1,11 +1,11 @@
 from torch.utils.data import DataLoader, RandomSampler
-import torch, os, sys, time, argparse, numpy as np
 from utils_dataset import SQLDataset, HDF5Dataset
+import torch, os, time, argparse, numpy as np
 from transformers.optimization import AdamW
 from model_generator import GeneTransformer
-from datetime import datetime, timedelta
-from utils_logplot import LogPlot
 import utils_misc, utils_tokenizer
+from utils_logplot import LogPlot
+from datetime import datetime
 
 from model_coverage import KeywordCoverage
 from model_guardrails import PatternPenalty, LengthPenalty, RepeatPenalty
@@ -17,7 +17,6 @@
 parser.add_argument("--experiment", type=str, required=True, help="Experiment name. Will be used to save a model file and a log file.")
 parser.add_argument("--dataset_file", type=str, required=True, help="Which dataset file to use. Can be full path or the root folder will be attached.")
 
-parser.add_argument("--root_folder", type=str, default="/home/"+user+"/")
 parser.add_argument("--train_batch_size", type=int, default=5, help="Training batch size.")
 parser.add_argument("--n_epochs", type=int, default=3, help="Number of epochs to run over the data.")
 parser.add_argument("--optim_every", type=int, default=4, help="Optimize every x backprops. A multiplier to the true batch size.")
@@ -34,8 +33,8 @@
     os.environ["CUDA_VISIBLE_DEVICES"] = ""+str(freer_gpu)
     args.experiment += "_"+freer_gpu
 
-models_folder = "/home/ubuntu/models/"
-log_folder = "/home/ubuntu/logs/"
+models_folder = "/home/phillab/models/"
+log_folder = "/home/phillab/logs/"
 
 summarizer_model_start = os.path.join(models_folder, "gpt2_copier23.bin")
 
@@ -65,6 +64,7 @@ def collate_func(inps):
     else:
         return [inp[0].decode() for inp in inps]
 
+
 param_optimizer = list(summarizer.model.named_parameters())
 no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
 optimizer_grouped_parameters = [
@@ -88,9 +88,9 @@ def collate_func(inps):
 
 print("Loading scorers")
 
-coverage_model_file = os.path.join(models_folder, "bert_coverage.bin")
+coverage_model_file = os.path.join(models_folder, "bert_coverage_google_cnndm_length15_1.bin")
 coverage_keyword_model_file = os.path.join(models_folder, "keyword_extractor.joblib")
-fluency_news_model_file = os.path.join(models_folder, "fluency_news_bs32.bin")
+fluency_news_model_file = os.path.join(models_folder, "news_gpt2_bs32.bin")
 
 scorers = [{"name": "coverage", "importance": 10.0, "sign": 1.0, "model": KeywordCoverage(args.device, keyword_model_file=coverage_keyword_model_file, model_file=coverage_model_file)},
            {"name": "fluency", "importance": 2.0, "sign": 1.0, "model": GeneTransformer(max_output_length=args.max_output_length, device=args.device, starter_model=fluency_news_model_file)},
@@ -102,6 +102,7 @@ def collate_func(inps):
 def background_tokenizer(bodies, out_queue):
     out_queue.put([bert_tokenizer.encode(body) for body in bodies])
 
+
 my_queue = queue.Queue()
 print("Started training")
 
@@ -116,7 +117,7 @@ def background_tokenizer(bodies, out_queue):
 dataloader = DataLoader(dataset=dataset, batch_size=args.train_batch_size, sampler=RandomSampler(dataset), drop_last=True, collate_fn=collate_func)
 
 for epi in range(n_epochs):
-    print("=================== EPOCH",epi, "===================")
+    print("=================== EPOCH", epi, "===================")
     for ib, documents in enumerate(dataloader):
         Timer = {}
 
@@ -126,7 +127,7 @@ def background_tokenizer(bodies, out_queue):
         bodies = [" ".join(doc.split(" ")[:300]) for doc in documents]
 
         # We run tokenization in the background, as it is BERT tokenization only used after the summarizer has run. Saves about 5% of time.
-        thread1 = threading.Thread(target = background_tokenizer, args = (bodies, my_queue))
+        thread1 = threading.Thread(target=background_tokenizer, args=(bodies, my_queue))
         # bodies_bert_tokenized = [bert_tokenizer.enncode(body) for body in bodies] # This is the not background version
         thread1.start()
 
@@ -159,11 +160,11 @@ def background_tokenizer(bodies, out_queue):
             sampled_scores = torch.FloatTensor(sampled_scores).to(args.device)
 
             argmax_scores, _ = scorer['model'].score(argmax_summaries, bodies, bodies_tokenized=bodies_bert_tokenized, extra=extra, lengths=argmax_end_idxs)
-            argmax_scores  = torch.FloatTensor(argmax_scores).to(args.device)
+            argmax_scores = torch.FloatTensor(argmax_scores).to(args.device)
 
             Timer["scores_"+scorer['name']] = time.time()-T
             total_sampled_scores += (scorer['sign'])*(scorer['importance'])*sampled_scores
-            total_argmax_scores  += (scorer['sign'])*(scorer['importance'])*argmax_scores
+            total_argmax_scores += (scorer['sign'])*(scorer['importance'])*argmax_scores
             log_obj[scorer['name']+"_score"] = sampled_scores.mean().item()
             scores_track[scorer['name']+"_scores"] = sampled_scores
 
@@ -180,7 +181,7 @@ def background_tokenizer(bodies, out_queue):
         T6 = time.time()
         Timer['backward'] = T6-T5
 
-        if ib%args.optim_every == 0:
+        if ib % args.optim_every == 0:
             optimizer.step()
             optimizer.zero_grad()
 
@@ -220,7 +221,7 @@ def background_tokenizer(bodies, out_queue):
 
         if ckpt_every > 0 and len(total_score_history) > ckpt_lookback:
             current_score = np.mean(total_score_history[-ckpt_lookback:])
-            
+
             if time.time()-time_ckpt > ckpt_every:
                 revert_ckpt = best_ckpt_score is not None and current_score < min(1.2*best_ckpt_score, 0.8*best_ckpt_score) # Could be negative or positive
                 print("================================== CKPT TIME, "+str(datetime.now())+" =================================")
@@ -232,7 +233,7 @@ def background_tokenizer(bodies, out_queue):
                     optimizer.load_state_dict(torch.load(ckpt_optimizer_file))
                 time_ckpt = time.time()
                 print("==============================================================================")
-    
+
             if best_ckpt_score is None or current_score > best_ckpt_score:
                 print("[CKPT] Saved new best at: %.3f %s" % (current_score, "["+str(datetime.now())+"]"))
                 best_ckpt_score = current_score

utils_tokenizer.py

@@ -1,7 +1,6 @@
-from transformers.tokenization_gpt2 import GPT2Tokenizer as GPT2Tok
-from transformers.tokenization_bert import BertTokenizer as BertTok
+from transformers import GPT2Tokenizer as GPT2Tok
+from transformers import BertTokenizer as BertTok
 import sentencepiece as spm
-import nltk
 
 class Capita:
     def forward(self, text):
@@ -26,18 +25,25 @@ def forward(self, text):
     def backward(self, text):
         words = text.split(" ")
         final_words = []
-        all_caps = False; capitalized = False
+        all_caps = False
+        capitalized = False
         for w in words:
-            if w == "⇧": all_caps = True
-            elif w == "↑": capitalized = True
+            if w == "⇧":
+                all_caps = True
+            elif w == "↑":
+                capitalized = True
             else:
                 final_word = w
-                if all_caps: final_word = final_word.upper()
+                if all_caps:
+                    final_word = final_word.upper()
                 elif capitalized:
-                    if len(final_word) <= 1: final_word = final_word.upper()
-                    else: final_word = final_word[0].upper()+final_word[1:]
+                    if len(final_word) <= 1:
+                        final_word = final_word.upper()
+                    else:
+                        final_word = final_word[0].upper()+final_word[1:]
                 final_words.append(final_word)
-                all_caps = False; capitalized = False
+                all_caps = False
+                capitalized = False
         return " ".join(final_words)
 
 class BPETokenizer:
@@ -53,7 +59,7 @@ def __init__(self, bpe_model, use_capita=True):
 
         if self.use_capita:
             self.cpt = Capita()
-        
+
     def tokenize(self, text):
         if len(text) == 0:
             return []
@@ -67,12 +73,12 @@ def tokenize(self, text):
         if tokens[0] == "▁":
             tokens = tokens[1:]
         return tokens
-        
+
     def encode(self, text):
         tokens = self.tokenize(text)
         token_ids = [self.sp.piece_to_id(w) for w in tokens]
         return token_ids
-        
+
     def decode(self, token_ids):
         text = self.sp.decode_ids(token_ids).replace("⇧", " ⇧").replace("↑", " ↑")
         if self.use_capita:
@@ -108,8 +114,8 @@ def __init__(self):
 
         self.pad_id = 0
         self.start_id = self.tokenizer.encode(self.start_tok)[0]
-        self.end_id =   self.tokenizer.encode(self.end_tok)[0]
-        self.vocab_size =  self.tokenizer.vocab_size
+        self.end_id = self.tokenizer.encode(self.end_tok)[0]
+        self.vocab_size = self.tokenizer.vocab_size
 
     def tokenize(self, text):
         return self.tokenizer.tokenize(text)