Upload 2 files

inductor.py

@@ -43,11 +43,11 @@ class BartInductor(object):
         self.orion_hypothesis_generator_path = 'facebook/bart-large' if not continue_pretrain_hypo_generator else ORION_HYPO_GENERATOR
 
         if group_beam:
-            self.orion_hypothesis_generator = BartForConditionalGeneration_GroupBeam.from_pretrained(self.orion_hypothesis_generator_path).to(device).eval().half()
+            self.orion_hypothesis_generator = BartForConditionalGeneration_GroupBeam.from_pretrained(self.orion_hypothesis_generator_path).to(device).eval()
         else:
-            self.orion_hypothesis_generator = BartForConditionalGeneration.from_pretrained(self.orion_hypothesis_generator_path).to(device).eval().half()
+            self.orion_hypothesis_generator = BartForConditionalGeneration.from_pretrained(self.orion_hypothesis_generator_path).to(device).eval()
         
-        self.orion_instance_generator = BartForConditionalGeneration.from_pretrained(self.orion_instance_generator_path).to(device).eval().half()
+        self.orion_instance_generator = BartForConditionalGeneration.from_pretrained(self.orion_instance_generator_path).to(device).eval()
     
         self.tokenizer = BartTokenizer.from_pretrained("facebook/bart-large")
         self.word_length = 2
@@ -128,15 +128,18 @@ class BartInductor(object):
 
     def extract_words_for_tA_bart(self, tA, k=6, print_it = False):
         spans = [t.lower().strip() for t in tA[:-1].split('<mask>')]
-        generated_ids = self.tokenizer([tA], padding='longest', return_tensors='pt')['input_ids'].to(device)
-        generated_ret = self.orion_instance_generator.generate(generated_ids.astype(np.float32), num_beams=max(120, k),
-                                                       min_length=min_length,
-                                                       max_length=max_length,
-                                                       early_stopping=early_stopping,
-                                                       **generator_kwargs)
-
+        generated_ids = self.tokenizer([tA], padding='longest', return_tensors='pt')['input_ids'].to(device).to(torch.int64)
+        generated_ret = self.orion_instance_generator.generate(generated_ids, num_beams=max(120, k),
+                                            #num_beam_groups=max(120, k),
+                                            max_length=generated_ids.size(1) + 15,
+                                            num_return_sequences=max(120, k), #min_length=generated_ids.size(1),
+                                            #diversity_penalty=2.0,
+                                            #length_penalty= 0.8,
+                                            #early_stopping=True, bad_words_ids=bad_words_ids, no_repeat_ngram_size=2,
+                                            output_scores=True,
+                                            return_dict_in_generate=True)
         summary_ids = generated_ret['sequences']
-        probs = F.softmax(generated_ret['sequences_scores'])
+        probs = F.softmax(generated_ret['sequences_scores'].to(torch.float32))
         txts = [self.tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True) for g in summary_ids]
         ret = []
 
@@ -221,7 +224,7 @@ class BartInductor(object):
                 index_words[len(index_words)] = '\t'.join(words)
             # index_words[len(templates)-1] = '\t'.join(words)
             if (len(templates) == batch_size) or enum==len(words_prob_sorted)-1 or (words_prob_sorted[enum+1][2]!=words_prob_sorted[enum][2]):
-                generated_ids = self.tokenizer(templates, padding="longest", return_tensors='pt')['input_ids'].to(device)
+                generated_ids = self.tokenizer(templates, padding="longest", return_tensors='pt')['input_ids'].to(device).to(torch.int64)
                 generated_ret = self.orion_hypothesis_generator.generate(generated_ids, num_beams=num_beams,
                                                     num_beam_groups=num_beams,
                                                     max_length=28, #template_length+5,
@@ -236,7 +239,7 @@ class BartInductor(object):
                                                     top_p=0.95,
                                                     )
                 summary_ids = generated_ret['sequences'].reshape((len(templates),num_beams,-1))
-                probs = F.softmax(generated_ret['sequences_scores'].reshape((len(templates),num_beams)),dim=1)
+                probs = F.softmax(generated_ret['sequences_scores'].reshape((len(templates),num_beams)),dim=1).to(torch.float32)
                 for ii in range(summary_ids.size(0)):
                     txts = [self.tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True) for g in
                             summary_ids[ii]]
@@ -303,7 +306,7 @@ class BartInductor(object):
 
 class CometInductor(object):
     def __init__(self):
-        self.model = AutoModelForSeq2SeqLM.from_pretrained("adamlin/comet-atomic_2020_BART").to(device).eval() # .half()
+        self.model = AutoModelForSeq2SeqLM.from_pretrained("adamlin/comet-atomic_2020_BART").to(device).eval().float() # .half()->float
         self.tokenizer = AutoTokenizer.from_pretrained("adamlin/comet-atomic_2020_BART")
         self.task = "summarization"
         self.use_task_specific_params()