Upload 2 files

expbert.py

@@ -42,15 +42,17 @@ GENERATED_EXP = {
     "disease": "data/exp/orion_disease_explanation.txt",
 }
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 def set_random_seed(seed):
     random.seed(seed)
     np.random.seed(seed)
     torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
 
 
 def print_config(config):
@@ -77,12 +79,12 @@ class ExpBERT(nn.Module):
     
     def forward(self, inputs):
         for k, v in inputs["encoding"].items():
-            inputs["encoding"][k] = v.cuda()
+            inputs["encoding"][k] = v.to(device)
         pooler_output = self.model(**inputs["encoding"]).last_hidden_state[:, 0, :].reshape(1, self.exp_num * self.config.hidden_size)
         pooler_output = self.dropout(pooler_output)
         logits = self.linear(pooler_output)
 
-        loss = self.criterion(logits, torch.LongTensor([inputs["label"]]).cuda())
+        loss = self.criterion(logits, torch.LongTensor([inputs["label"]]).to(device))
         prediction = torch.argmax(logits) 
 
         return {
@@ -187,7 +189,7 @@ class Trainer(object):
 
         self.train_dataset = REDataset(TASK2PATH['{}-train'.format(args.task)], exp, self.tokenizer)
         self.test_dataset = REDataset(TASK2PATH['{}-test'.format(args.task)], exp, self.tokenizer)
-        self.model = AutoModelForSequenceClassification.from_pretrained(args.model).cuda() if self.args.no_exp else ExpBERT(args, len(exp)).cuda()
+        self.model = AutoModelForSequenceClassification.from_pretrained(args.model).to(device) if self.args.no_exp else ExpBERT(args, len(exp)).to(device)
 
         self.train_loader = DataLoader(
             self.train_dataset,
@@ -220,7 +222,7 @@ class Trainer(object):
                     self.model.zero_grad()
                     if self.args.no_exp:
                         for k, v in examples.items():
-                            examples[k] = v.cuda()
+                            examples[k] = v.to(device)
                         outputs = self.model(**examples)
                         outputs.loss.backward()
 
@@ -245,7 +247,7 @@ class Trainer(object):
                 for step, examples in enumerate(self.test_loader):
                     if self.args.no_exp:
                         for k, v in examples.items():
-                            examples[k] = v.cuda()
+                            examples[k] = v.to(device)
                         outputs = self.model(**examples)
                         loss.append(outputs.loss.float())
                         labels.extend(examples["labels"].tolist())

inductor.py

@@ -11,6 +11,8 @@ from src.bart_with_group_beam import BartForConditionalGeneration_GroupBeam
 from src.utils import (construct_template, filter_words,
                        formalize_tA, post_process_template)
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 ORION_HYPO_GENERATOR = 'chenxran/orion-hypothesis-generator'
 ORION_INS_GENERATOR = 'chenxran/orion-instance-generator'
 
@@ -41,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).cuda().eval().half()
+            self.orion_hypothesis_generator = BartForConditionalGeneration_GroupBeam.from_pretrained(self.orion_hypothesis_generator_path).to(device).eval().half()
         else:
-            self.orion_hypothesis_generator = BartForConditionalGeneration.from_pretrained(self.orion_hypothesis_generator_path).cuda().eval().half()
+            self.orion_hypothesis_generator = BartForConditionalGeneration.from_pretrained(self.orion_hypothesis_generator_path).to(device).eval().half()
         
-        self.orion_instance_generator = BartForConditionalGeneration.from_pretrained(self.orion_instance_generator_path).cuda().eval().half()
+        self.orion_instance_generator = BartForConditionalGeneration.from_pretrained(self.orion_instance_generator_path).to(device).eval().half()
     
         self.tokenizer = BartTokenizer.from_pretrained("facebook/bart-large")
         self.word_length = 2
@@ -69,7 +71,7 @@ class BartInductor(object):
         self.bad_words_ids = [self.tokenizer.encode(bad_word)[1:-1] for bad_word in ['also', ' also']]
         stop_index = self.tokenizer(self.stop_sub_list, max_length=4, padding=True)
         stop_index = torch.tensor(stop_index['input_ids'])[:, 1]
-        stop_weight = torch.zeros(1, self.tokenizer.vocab_size).cuda()
+        stop_weight = torch.zeros(1, self.tokenizer.vocab_size).to(device)
         stop_weight[0, stop_index] -= 100
         self.stop_weight = stop_weight[0, :]
 
@@ -99,9 +101,9 @@ class BartInductor(object):
         if required_token <= self.word_length:
             k = min(k, 2)
         ret = []
-        generated_ids = self.tokenizer(tA, max_length=128, padding='longest', return_tensors='pt')  # ["input_ids"].cuda()
+        generated_ids = self.tokenizer(tA, max_length=128, padding='longest', return_tensors='pt')  # ["input_ids"].to(device)
         for key in generated_ids.keys():
-            generated_ids[key] = generated_ids[key].cuda()
+            generated_ids[key] = generated_ids[key].to(device)
         mask_index = torch.where(generated_ids["input_ids"][0] == self.tokenizer.mask_token_id)
         generated_ret = self.orion_instance_generator(**generated_ids)
         #logits = generated_ret.logits
@@ -126,7 +128,7 @@ 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'].cuda()
+        generated_ids = self.tokenizer([tA], padding='longest', return_tensors='pt')['input_ids'].to(device)
         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,
@@ -222,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'].cuda()
+                generated_ids = self.tokenizer(templates, padding="longest", return_tensors='pt')['input_ids'].to(device)
                 generated_ret = self.orion_hypothesis_generator.generate(generated_ids, num_beams=num_beams,
                                                     num_beam_groups=num_beams,
                                                     max_length=28, #template_length+5,
@@ -304,7 +306,7 @@ class BartInductor(object):
 
 class CometInductor(object):
     def __init__(self):
-        self.model = AutoModelForSeq2SeqLM.from_pretrained("adamlin/comet-atomic_2020_BART").cuda().eval() # .half()
+        self.model = AutoModelForSeq2SeqLM.from_pretrained("adamlin/comet-atomic_2020_BART").to(device).eval() # .half()
         self.tokenizer = AutoTokenizer.from_pretrained("adamlin/comet-atomic_2020_BART")
         self.task = "summarization"
         self.use_task_specific_params()
@@ -387,8 +389,8 @@ class CometInductor(object):
             input_ids, attention_mask = self.trim_batch(**batch, pad_token_id=self.tokenizer.pad_token_id)
 
             summaries = self.model.generate(
-                input_ids=input_ids.cuda(),
-                attention_mask=attention_mask.cuda(),
+                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,