update app.py

app.py

@@ -85,13 +85,9 @@ veracity_model = SequenceClassificationModule.load_from_checkpoint(veracity_chec
 # Justification
 justification_tokenizer = BartTokenizer.from_pretrained('facebook/bart-large', add_prefix_space=True)
 bart_model = BartForConditionalGeneration.from_pretrained("facebook/bart-large")
-best_checkpoint = os.getcwd()+ '/averitec/pretrained_models/bart_justifications_verdict-epoch=13-val_loss=2.03-val_meteor=0.28.ckpt'
+best_checkpoint = os.getcwd() + '/averitec/pretrained_models/bart_justifications_verdict-epoch=13-val_loss=2.03-val_meteor=0.28.ckpt'
 justification_model = JustificationGenerationModule.load_from_checkpoint(best_checkpoint, tokenizer=justification_tokenizer, model=bart_model).to('cuda')
 # justification_model = JustificationGenerationModule.load_from_checkpoint(best_checkpoint, tokenizer=justification_tokenizer, model=bart_model).to(device)
-
-
-print("veracity_model_device_0:{}".format(veracity_model.device))
-print("justification_model_device_0:{}".format(justification_model.device))
 # ---------------------------------------------------------------------------
 
 # ----------------------------------------------------------------------------
@@ -285,9 +281,8 @@ def veracity_prediction(claim, evidence):
         return pred_label
 
     tokenized_strings, attention_mask = dataLoader.tokenize_strings(evidence_strings)
-    example_support = torch.argmax(veracity_model(tokenized_strings, attention_mask=attention_mask).logits, axis=1)
+    example_support = torch.argmax(veracity_model(tokenized_strings.to('cuda'), attention_mask=attention_mask.to('cuda')).logits, axis=1)
     # example_support = torch.argmax(veracity_model(tokenized_strings.to(device), attention_mask=attention_mask.to(device)).logits, axis=1)
-    print("veracity_model_device_1:{}".format(veracity_model.device))
 
     has_unanswerable = False
     has_true = False
@@ -349,9 +344,8 @@ def justification_generation(claim, evidence, verdict_label):
     #
     claim_str = extract_claim_str(claim, evidence, verdict_label)
     claim_str.strip()
-    pred_justification = justification_model.generate(claim_str)
+    pred_justification = justification_model.generate(claim_str, device='cuda')
     # pred_justification = justification_model.generate(claim_str, device=device)
-    print("justification_model_device_1:{}".format(justification_model.device))
 
     return pred_justification.strip()