wing-nus commited on
Commit
3c83061
1 Parent(s): 1a0701e

Update controlled_summarization.py

Browse files
Files changed (1) hide show
  1. controlled_summarization.py +4 -1
controlled_summarization.py CHANGED
@@ -4,11 +4,14 @@ from SciAssist import Summarization
4
  import os
5
  import requests
6
  from datasets import load_dataset
 
 
 
7
 
8
  acl_data = load_dataset("dyxohjl666/CocoScisum_ACL", revision="refs/convert/parquet")
9
  device = "gpu" if torch.cuda.is_available() else "cpu"
10
 
11
- ctrlsum_pipeline = Summarization(os_name="nt",device=device)
12
 
13
  acl_dict = {}
14
  recommended_kw = {}
 
4
  import os
5
  import requests
6
  from datasets import load_dataset
7
+ print(f"Is CUDA available: {torch.cuda.is_available()}")
8
+ # True
9
+ print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
10
 
11
  acl_data = load_dataset("dyxohjl666/CocoScisum_ACL", revision="refs/convert/parquet")
12
  device = "gpu" if torch.cuda.is_available() else "cpu"
13
 
14
+ ctrlsum_pipeline = Summarization(os_name="nt",checkpoint="dyxohjl666/flant5-xl-cocoscisum",device=device)
15
 
16
  acl_dict = {}
17
  recommended_kw = {}