Update inductor.py

inductor.py

@@ -312,100 +312,3 @@ class BartInductor(object):
         return ret
 
 
-class CometInductor(object):
-    def __init__(self):
-        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()
-        self.decoder_start_token_id = None
-
-    def drop_repeat(self, old_list):
-        new_list = []
-        for item in old_list:
-            if item not in new_list:
-                new_list.append(item)
-        
-        return new_list
-
-    def chunks(self, lst, n):
-        """Yield successive n-sized chunks from lst."""
-        for i in range(0, len(lst), n):
-            yield lst[i : i + n]
-
-    def use_task_specific_params(self):
-        """Update config with summarization specific params."""
-        task_specific_params = self.model.config.task_specific_params
-
-        if task_specific_params is not None:
-            pars = task_specific_params.get(self.task, {})
-            self.model.config.update(pars)
-
-    def trim_batch(
-        self, input_ids, pad_token_id, attention_mask=None,
-    ):
-        """Remove columns that are populated exclusively by pad_token_id"""
-        keep_column_mask = input_ids.ne(pad_token_id).any(dim=0)
-        if attention_mask is None:
-            return input_ids[:, keep_column_mask]
-        else:
-            return (input_ids[:, keep_column_mask], attention_mask[:, keep_column_mask])
-
-    def generate(self, inputs, k, topk):
-        outputs = []
-        words = ['PersonX', 'PersonY']
-        for i, _ in enumerate(re.findall("<mask>", inputs)):
-            index = inputs.index('<mask>')
-            inputs = inputs[:index] + words[i] + inputs[index + len('<mask>'):]
-
-        for relation in RELATIONS:
-            inputs = "{} {} [GEN]".format(inputs[:-1], relation)
-            gen = self.generate_(inputs, num_generate=10)
-            switch = 0
-            for output in gen[0]:
-                output = output.strip()
-                if re.search("PersonX|X", output) and re.search("PersonY|Y", output):
-                    temp = re.sub("PersonX|X|PersonY|Y", "<mask>", output.strip())
-                    if temp.endswith("."):
-                        outputs.append(temp)
-                    else:
-                        outputs.append(temp + ".")
-                    switch = 1
-                    break
-
-            if switch == 0:
-                output = gen[0][0]
-                temp = re.sub("PersonX|X|PersonY|Y", "<mask>", output.strip())
-                if temp.endswith("."):
-                    outputs.append(temp)
-                else:
-                    outputs.append(temp + ".")
-        
-        outputs = [output.replace('PersonX', '<mask>').replace('PersonY', '<mask>') for output in outputs]
-        return outputs
-
-    def generate_(
-            self, 
-            queries,
-            decode_method="beam", 
-            num_generate=5, 
-        ):
-
-        with torch.no_grad():
-            decs = []
-            batch = self.tokenizer(queries, return_tensors="pt", padding="longest")
-            input_ids, attention_mask = self.trim_batch(**batch, pad_token_id=self.tokenizer.pad_token_id)
-
-            summaries = self.model.generate(
-                input_ids=input_ids.to(device),
-                attention_mask=attention_mask.to(device),
-                decoder_start_token_id=self.decoder_start_token_id,
-                num_beams=num_generate,
-                num_return_sequences=num_generate,
-                )
-
-            dec = self.tokenizer.batch_decode(summaries, skip_special_tokens=True, clean_up_tokenization_spaces=False)
-            decs.append(dec)
-
-            return decs
-