numBery commited on
Commit
44f0002
1 Parent(s): e568201

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -1
app.py CHANGED
@@ -1,6 +1,7 @@
1
  import torch
2
 
3
  from keybert import KeyBERT
 
4
  from keyphrase_vectorizers import KeyphraseCountVectorizer
5
  from transformers import T5ForConditionalGeneration,T5Tokenizer
6
 
@@ -17,7 +18,8 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
17
  HfFolder.save_token(st.secrets["hf-auth-token"])
18
 
19
  # Load KeyBert Model
20
- kw_extractor = KeyBERT('valurank/MiniLM-L6-Keyword-Extraction', use_auth_token=True)
 
21
 
22
  # Load T5 for Paraphrasing
23
  t5_model = T5ForConditionalGeneration.from_pretrained('valurank/t5-paraphraser', use_auth_token=True)
 
1
  import torch
2
 
3
  from keybert import KeyBERT
4
+ from sentence_transformers import SentenceTransformer
5
  from keyphrase_vectorizers import KeyphraseCountVectorizer
6
  from transformers import T5ForConditionalGeneration,T5Tokenizer
7
 
 
18
  HfFolder.save_token(st.secrets["hf-auth-token"])
19
 
20
  # Load KeyBert Model
21
+ tmp_model = SentenceTransformer('valurank/MiniLM-L6-Keyword-Extraction', use_auth_token=True)
22
+ kw_extractor = KeyBERT(tmp_model)
23
 
24
  # Load T5 for Paraphrasing
25
  t5_model = T5ForConditionalGeneration.from_pretrained('valurank/t5-paraphraser', use_auth_token=True)