Update textrank.py

textrank.py

@@ -1,19 +1,23 @@
 import numpy as np
 import pandas as pd
 import nltk
-nltk.download('punkt') # one time execution
 import re
-import warnings
-warnings.filterwarnings('ignore')
-from sklearn.metrics.pairwise import cosine_similarity
+
+import torch
 import networkx as nx
 from tqdm import tqdm
-
 from sentence_transformers import SentenceTransformer
+from sklearn.metrics.pairwise import cosine_similarity
+
+import warnings
+warnings.filterwarnings('ignore')
+
+nltk.download('punkt')
 
 model = SentenceTransformer('all-mpnet-base-v2')
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
 
-model.to('cuda')
 def get_summary(text, num_words: int=1000):
     sentences = nltk.sent_tokenize(text)
     embeddings = model.encode(sentences, show_progress_bar=False)